Function fused_mm_sample_cuda

src/fused_mm_sampling/cuda_impl.py:44–72

Fused matrix-multiply & sampling using the CUDA C++ kernel.

(
    weights: torch.Tensor,  # [V, D] bfloat16
    hidden_states: torch.Tensor,  # [H, D] bfloat16
    num_samples: int,
    temperature: torch.Tensor,  # scalar (0-d)
    seed: int = 0,
    **kwargs,
)

def fused_mm_sample_cuda(
    weights: torch.Tensor,  # [V, D] bfloat16
    hidden_states: torch.Tensor,  # [H, D] bfloat16
    num_samples: int,
    temperature: torch.Tensor,  # scalar (0-d)
    seed: int = 0,
    **kwargs,
) -> torch.Tensor:
    """Fused matrix-multiply & sampling using the CUDA C++ kernel."""
    V, D = weights.shape  # noqa: N806
    H = hidden_states.shape[0]  # noqa: N806
    assert hidden_states.shape[1] == D

    n_tiles_v = (V + TILE_V - 1) // TILE_V

    # Temperature must be float32 on GPU
    if temperature.dtype != torch.float32:
        temperature = temperature.float()

    maxs = torch.empty((n_tiles_v, H, num_samples), dtype=torch.float32, device=weights.device)
    maxs_idx = torch.empty((n_tiles_v, H, num_samples), dtype=torch.long, device=weights.device)

    mod = _get_module()
    mod.fmms_stage1(weights, hidden_states, maxs, maxs_idx, temperature, seed)

    # Stage 2: reduce across V-tiles (identical to Triton wrapper)
    idxs = maxs.max(dim=0).indices  # [H, num_samples]
    samples = maxs_idx.gather(dim=0, index=idxs.unsqueeze(0)).squeeze(0)
    return samples  # [H, num_samples]