Function grid

src/fused_mm_sampling/core.py:277–283

(meta)

    # HBM (155 MB per decode step at Qwen3-1.7B / H=256) for a buffer the
    # kernel never touches. Pass a 1-element dummy in that case so the kernel
    # still has a valid pointer to receive.
    logits_shape = (V, H) if return_logits else (1,)
    logits_out = torch.empty(logits_shape, dtype=torch.float32, device=weights.device)

    grid_size = {"v": None}

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