Function _launch_kernel

src/fused_mm_sampling/tl_fused_mm_topk.py:119–143

(weights, hidden_states, out_vals, out_ids, V, D, H, top_k, temperature)

@torch.compiler.disable
def _launch_kernel(weights, hidden_states, out_vals, out_ids, V, D, H, top_k, temperature):  # noqa: N803
    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,
        hidden_states_ptr=hidden_states,
        topk_vals_ptr=out_vals,
        topk_ids_ptr=out_ids,
        vocab_size=V,
        hidden_size=D,
        n_hidden_states=H,
        top_k=top_k,
        temperature_ptr=temperature,
    )

    assert grid_size["v"] is not None
    return grid_size["v"]


@triton.autotune(