Function fused_mm_sample_triton

src/fused_mm_sampling/core.py:220–332

(
    weights: torch.Tensor,  # [V_local, D] (may be a TP shard)
    hidden_states: torch.Tensor,  # [n_hidden_states, D]
    num_samples: int,
    temperature: torch.Tensor,  # scalar (0-d)
    seed: int,
    greedy_sampling: bool = False,
    tp: "TPInfo" = TP1,
    return_logits: bool = False,
)

# @torch.compile(fullgraph=True)
@nvtx.annotate()
def fused_mm_sample_triton(
    weights: torch.Tensor,  # [V_local, D] (may be a TP shard)
    hidden_states: torch.Tensor,  # [n_hidden_states, D]
    num_samples: int,
    temperature: torch.Tensor,  # scalar (0-d)
    seed: int,
    greedy_sampling: bool = False,
    tp: "TPInfo" = TP1,
    return_logits: bool = False,
):
    assert torch.cuda.is_available(), "fused_mm_sample_triton requires CUDA"
    V, D = weights.shape  # noqa: N806
    H, D2 = hidden_states.shape  # noqa: N806
    if D2 != D:
        raise ValueError(
            f"hidden_states second dimension ({D2}) must match weights second dimension ({D})"
        )

    # The kernel uses TMA descriptors which need a runtime allocator. Some
    # autotuner configs (notably the ones picked on B200/sm_100) request global
    # scratch from this allocator; without it Triton raises a RuntimeError at launch.
    set_torch_allocator_for_tma_descriptors_cached()

    NUM_SMS = num_sms_cached(weights.device.index)  # noqa: N806

    max_grid_size_v = triton.cdiv(V, MIN_BLOCK_SIZE_V)
    if tp.size > 1:
        maxs, maxs_idx, symm_mem_hdl, storage_offset_maxs_idx = allocate_symm_mem_outputs(
            num_samples=num_samples,
            max_grid_size_v=max_grid_size_v,
            H=H,
        )
        kernel_maxs = maxs[tp.rank]
        kernel_maxs_idx = maxs_idx[tp.rank]
        symm_mem_buffer_ptrs = symm_mem_hdl.buffer_ptrs_dev
    else:
        maxs = torch.empty(
            (num_samples, max_grid_size_v, H),
            dtype=torch.bfloat16,
            device=weights.device,
        )
        maxs_idx = torch.empty_like(maxs, dtype=torch.long)
        kernel_maxs = maxs
        kernel_maxs_idx = maxs_idx
        storage_offset_maxs_idx = 0
        symm_mem_buffer_ptrs = maxs

    # logits_out is only read when RETURN_LOGITS=True. For the common path
    # (return_logits=False), allocating a (V, H) fp32 buffer per call is wasted
    # 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):