Function bsz_h

src/fused_mm_sampling/core.py:856–861

(H: int)

    return logits, indices


@nvtx.annotate()
def flashinfer_sampling_from_logits(
    weights: torch.Tensor,  # [V, D]
    hidden_states: torch.Tensor,  # [n_hidden_states, D]
    num_samples: int,
    temperature: torch.Tensor,  # scalar (0-d)
    tp: "TPInfo" = TP1,
    **_kwargs,