Function sample

src/fused_mm_sampling/core.py:29–61

(
    weights: torch.Tensor,  # [V, D] (may be a TP shard over dim V)
    hidden_states: torch.Tensor,  # [n_hidden_states, D]
    num_samples: int,
    temperature: torch.Tensor,  # scalar (0-d)
    return_probs: bool = False,
    seed: int = None,
    tl_matmul: bool = False,
    top_k: int | None = None,
    top_p: float | None = None,
    use_qitra: bool = False,
    tp: "TPInfo" = TP1,
)

# gather_dim != 0), breaking torch.compile(fullgraph=True) on _sample_compiled
# / greedy_sample_compiled. Upstream fix is in flight (PRs #180389, #182435);
# patch the rule map until it lands.
manual_torch_name_rule_map.setdefault("torch._utils._maybe_view_chunk_cat", UserFunctionVariable)


def sample(
    weights: torch.Tensor,  # [V, D] (may be a TP shard over dim V)
    hidden_states: torch.Tensor,  # [n_hidden_states, D]
    num_samples: int,
    temperature: torch.Tensor,  # scalar (0-d)
    return_probs: bool = False,
    seed: int = None,
    tl_matmul: bool = False,
    top_k: int | None = None,
    top_p: float | None = None,
    use_qitra: bool = False,
    tp: "TPInfo" = TP1,
):
    if seed is not None:
        torch.manual_seed(seed)
    if tl_matmul:
        logits = matmul(hidden_states, weights)  # [n_hidden_states, V]
    else:
        logits = hidden_states @ weights.T  # [n_hidden_states, V]
    if tp.size > 1:
        logits = _allgather_logits(logits)  # shape [H, V_local] -> [H, V]
    # Upcast to float32 before temperature scaling: Qitra asserts float32, and
    # torch.multinomial produces imprecise distributions with bfloat16.
    # See findings/upcasting-before-softmax.md.
    logits = logits.float() / temperature
    if use_qitra:
        probs = apply_top_k_top_p_qitra(logits, top_k, top_p)
    else:
        probs = apply_top_k_top_p(logits, top_k, top_p)
    samples = torch.multinomial(probs, num_samples, replacement=True)
    if return_probs:
        return samples, probs