Function grid

src/fused_mm_sampling/tl_matmul.py:172–176

(meta)

    # This enables processing many N blocks for the same M block,
    # allowing A matrix (small M dimension) to be reused from L2 cache.
    def grid(meta):
        return (
            triton.cdiv(N, meta["BLOCK_SIZE_N"]),
            triton.cdiv(M, meta["BLOCK_SIZE_M"]),
        )

    matmul_kernel[grid](
        a,