Is this page helpful?

core.inference.sampling.flashinfer_sampling#

Module Contents#

Classes#

FlashInferSampling

FlashInfer sampling with per-step top-p-only / top-k-only / joint dispatch.

API#

class core.inference.sampling.flashinfer_sampling.FlashInferSampling(
vocab_size: int,
rng: torch.Generator,
config=None,
enable_cuda_graph: bool = False,
)#

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

FlashInfer sampling with per-step top-p-only / top-k-only / joint dispatch.

Each step selects a kernel from the batch’s active filters: the dedicated exact top-p or top-k kernel when only one filter is in use, and the joint kernel only for genuinely mixed batches. The dispatch flags are read from the pinned CPU sampling metadata, so evaluating them costs no GPU sync.

The sampler runs eagerly. Its kernel choice is data-dependent (it varies with which filters the batch uses), so it cannot be captured in a CUDA graph; running eagerly also lets the controller’s seeded RNG generator advance its philox offset normally between steps – fresh randomness per step, reproducible from the seed. (FlashInfer bakes the philox state into a graph as a by-value constant at capture, so a captured sampler replays identical random numbers; see https://www.linkedin.com/pulse/pinned-rng-drifting-crash-from-cuda-graph-chenyang-zhao-csuac/)

Initialization

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#

Sample tokens, dispatching top-p-only / top-k-only / joint by filter flags.

Parameters:

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

  • n – Number of rows to sample.

  • context – The active DynamicInferenceContext.

  • no_top_k – Required batch-level dispatch flags (whether NO active request uses top-k / top-p). The caller computes them once from the pinned CPU sampling metadata (the controller’s _active_requests_sampling_filter_flags).

  • no_top_p – Required batch-level dispatch flags (whether NO active request uses top-k / top-p). The caller computes them once from the pinned CPU sampling metadata (the controller’s _active_requests_sampling_filter_flags).

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

  • token_to_request_index – When set, sampling parameters are gathered per-token rather than per-request (speculative decoding path).

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

  • eager – Accepted for API symmetry; ignored (no CUDA graph).

  • cache_key – Accepted for API symmetry; ignored (no CUDA graph).

Returns:

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

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

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

Parameters:

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

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

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

Returns:

Per-row log probabilities for the processed distribution.

Return type:

Tensor