Is this page helpful?

core.inference.sampling.torch_sampling#

Module Contents#

Classes#

TorchSampling

Sampling via bucketed torch.multinomial.

API#

class core.inference.sampling.torch_sampling.TorchSampling(rng: torch.Generator, vocab_size: int)#

Bases: megatron.core.inference.sampling.base.Sampling

Sampling via bucketed torch.multinomial.

Groups requests into unique buckets by (temperature, top_k, top_p) for separate launches.

Initialization

static _modify_logits_for_top_k_filtering(
logits: torch.Tensor,
top_k: int,
) → None#

In-place: set logits outside the top-k set to -inf.

static _modify_logits_for_top_p_filtering(
logits: torch.Tensor,
top_p: float,
) → None#

In-place: set logits outside the top-p (nucleus) set to -inf.

static filter_logits(
last_token_logits: torch.Tensor,
temperature: float,
top_k: int,
top_p: float,
*,
vocab_size: Optional[int] = None,
) → torch.Tensor#

Temperature-scale then top-k/top-p filter logits; filtered entries become -inf.

Returns a new tensor (input unmodified). Shared by sample_from_logits and log_probs_kernel so sampling and processed log-probs apply the same filter.

static sample_from_logits(
last_token_logits: torch.Tensor,
temperature: float,
top_k: int,
top_p: float,
*,
generator: torch.Generator,
vocab_size: Optional[int] = None,
) → torch.Tensor#

Sample tokens from logits with temperature, top-k, and top-p filtering.

Shared between dynamic batching and static batching.

Parameters:

  • last_token_logits – Logits of shape [batch_size, vocab_size].

  • temperature – Temperature scaling factor.

  • top_k – Top-k filtering value (0 = disabled).

  • top_p – Top-p (nucleus) filtering value (0.0 = disabled).

  • generator – RNG used by torch.multinomial.

  • vocab_size – When provided, asserts top_k < vocab_size and clamps the sampled ids to [0, vocab_size - 1].

Returns:

Sampled token ids of shape [batch_size].

log_probs_kernel(
logits: torch.Tensor,
context,
*,
token_to_request_index: Optional[torch.Tensor] = None,
) → torch.Tensor#

Per-row log-probs of the temperature, top-k/top-p sampling distribution.

Buckets rows by identical (temperature, top_k, top_p) and reuses filter_logits (the same filter as sample_from_logits) so log-probs match how this backend samples.

Parameters:

  • logits (Tensor) – Raw logits with shape [num_rows, vocab_size].

  • context – Active dynamic inference context providing CPU sampling metadata.

  • token_to_request_index (Optional[Tensor]) – Optional CPU mapping from each logits row to its request index.

Returns:

Per-row log probabilities for the processed distribution.

Return type:

Tensor

sample_kernel(
logits: torch.Tensor,
n: int,
context,
*,
no_top_k: bool,
no_top_p: bool,
gather_indices: Optional[torch.Tensor] = None,
token_to_request_index: Optional[torch.Tensor] = None,
output: Optional[torch.Tensor] = None,
eager: bool = False,
cache_key: Any = None,
) → torch.Tensor#

Bucket active requests by sampling parameters and sample each bucket.

Parameters:

  • logits – Logits tensor of shape [>=n, vocab_size].

  • n – Number of rows to sample.

  • context – The active DynamicInferenceContext.

  • no_top_k – Batch-level dispatch flags (part of the shared kernel contract); ignored here since the exact per-bucket sort already handles any top-k / top-p combination.

  • no_top_p – Batch-level dispatch flags (part of the shared kernel contract); ignored here since the exact per-bucket sort already handles any top-k / top-p combination.

  • gather_indices – When set, sample from logits[gather_indices[:n], :].

  • token_to_request_index – When set, the loop dispatches per-token rather than per-request (used by the speculative path).

  • output – Optional caller-owned destination tensor of shape [n].

  • eager – Accepted for API symmetry; ignored (TorchSampling has no graph wrapper).

  • cache_key – Accepted for API symmetry; ignored.

Returns:

Sampled token IDs in output, or a newly allocated tensor when it is not provided.