Function grid

src/fused_mm_sampling/tl_fused_mm_topk.py:122–128

(meta)

    grid_size = {"v": None}

    def grid(meta):
        grid_size_v = triton.cdiv(V, meta["BLOCK_SIZE_V"])
        grid_size["v"] = grid_size_v
        return (
            grid_size_v,
            triton.cdiv(H, meta["BLOCK_SIZE_H"]),
        )

    fused_mm_topk_kernel[grid](
        weights_ptr=weights,