Function grid

src/fused_mm_sampling/tl_gemv.py:19–20

(meta)

    out = torch.empty(B, device=x.device, dtype=x.dtype)

    def grid(meta):
        return (triton.cdiv(B, meta["BLOCK_SIZE_B"]),)

    tl_gemv_kernel[grid](x, Y, out, M, B)
    return out