Is this page helpful?

core.inference.sampling.base#

Module Contents#

Classes#

Sampling

Abstract base for inference sampling backends.

API#

class core.inference.sampling.base.Sampling#

Bases: abc.ABC

Abstract base for inference sampling backends.

Subclasses implement sample_kernel and log_probs_kernel. CUDA graphs are added via CudaGraphManager.

abstractmethod 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 n tokens from logits and return them.

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 (see the controller’s _active_requests_sampling_filter_flags), so the kernel never has to.

  • 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 (see the controller’s _active_requests_sampling_filter_flags), so the kernel never has to.

  • gather_indices – If provided, only sample from logits[gather_indices[:n], :].

  • token_to_request_index – Per-token request mapping; when set, sampling parameters are gathered per-token instead of per-request.

  • 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 of shape [n].

sample_speculative(
required_logits: torch.Tensor,
num_decode: int,
num_prefill: int,
num_speculative_tokens: int,
context,
*,
gather_indices: Optional[torch.Tensor] = None,
eager: bool = False,
cache_key: Any = None,
) → torch.Tensor#

Sample tokens for the speculative-verify path.

Decode requests contribute 1 + num_speculative_tokens rows; prefill requests contribute 1. Builds the per-token request mapping and dispatches to the return-valued sample_kernel.

When gather_indices is supplied, the kernel selects via logits[gather_indices[:n], :]. When gather_indices is None, required_logits is expected to be already pre-gathered to the layout described above (e.g. when materialize_only_last_token_logits=True upstream).

Parameters:

  • required_logits – Logits containing base and speculative rows.

  • num_decode – Number of decode requests.

  • num_prefill – Number of prefill requests.

  • num_speculative_tokens – Number of draft tokens per decode request.

  • context – The active DynamicInferenceContext.

  • gather_indices – Optional rows to gather from required_logits.

  • eager – Whether to bypass a wrapped CUDA graph.

  • cache_key – CUDA graph lookup key.

Returns:

Sampled token IDs for all required base and speculative rows.

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

Per-row log-probs of the distribution this backend samples from.

Parameters:

  • logits – [num_rows, vocab_size] raw logits.

  • context – The active DynamicInferenceContext.

  • token_to_request_index – Optional per-row request mapping. When omitted, each logits row maps to the request at the same index.

Returns:

[num_rows, vocab_size] log-probs; filtered-out tokens are -inf.