diff --git a/megatron/core/inference/config.py b/megatron/core/inference/config.py index c4b092309c2..756f4fbe4be 100644 --- a/megatron/core/inference/config.py +++ b/megatron/core/inference/config.py @@ -2,7 +2,7 @@ from dataclasses import InitVar, dataclass from enum import Enum -from typing import List, Optional, Tuple +from typing import List, Literal, Optional, Tuple import torch @@ -297,6 +297,9 @@ class InferenceConfig: Defaults to 0, which means no logging. """ + sampling_backend: Literal['torch', 'flashinfer'] = 'torch' + """Which sampling kernels to use during inference.""" + request_metadata_types: Optional[List[Tuple[str, torch.dtype]]] = None """ A list of the per-request metadata types to track. Each entry is a tuple @@ -320,3 +323,12 @@ def __post_init__(self, verbose: bool): f"prefix_caching_routing_alpha must be in [0, 1], " f"got {self.prefix_caching_routing_alpha}" ) + + if self.sampling_backend == 'flashinfer': + try: + import flashinfer # noqa: F401 + except ImportError as e: + raise ImportError( + "sampling_backend='flashinfer' requires the flashinfer package; " + "install it or set sampling_backend='torch'." + ) from e diff --git a/megatron/core/inference/contexts/dynamic_context.py b/megatron/core/inference/contexts/dynamic_context.py index 159e1f90b34..42b5739934c 100644 --- a/megatron/core/inference/contexts/dynamic_context.py +++ b/megatron/core/inference/contexts/dynamic_context.py @@ -928,16 +928,20 @@ def initialize_all_tensors(self) -> None: # transferred to GPU each step via transfer_bookkeeping_to_gpu(). # Layout matches ContextGPUView._buf so a single cudaMemcpyAsync # suffices. Int64 token fields come first (8-byte aligned automatically), - # then int32 token fields, then int32 request-staging fields. - # token_to_input_ids (int64, max_tokens) - # token_to_pos_ids (int64, max_tokens) - # token_to_block_idx (int32, max_tokens) - # token_to_local_position_within_kv_block (int32, max_tokens) - # token_to_request_idx (int32, max_tokens) - # token_to_position_in_request (int32, max_tokens) - # request_in_prefill_status (staging) (int32, max_requests) - # request_query_lengths (staging) (int32, max_requests) - # request_kv_length_offsets (staging) (int32, max_requests) + # then int32 token fields, then int32/float32 request-staging fields. + # token_to_input_ids (int64, max_tokens) + # token_to_pos_ids (int64, max_tokens) + # token_to_block_idx (int32, max_tokens) + # token_to_local_position_within_kv_block (int32, max_tokens) + # token_to_request_idx (int32, max_tokens) + # token_to_position_in_request (int32, max_tokens) + # request_in_prefill_status (staging) (int32, max_requests) + # request_query_lengths (staging) (int32, max_requests) + # request_kv_length_offsets (staging) (int32, max_requests) + # temperature (staging) (float32, max_requests) + # top_k (staging) (int32, max_requests) + # top_p (staging) (float32, max_requests) + # active_request_last_token_idxs (alias) (int32, max_requests) # # Token fields are aliased with the source-of-truth attributes # (`self.token_to_input_ids`, etc.) because the forward pass reads @@ -948,7 +952,8 @@ def initialize_all_tensors(self) -> None: # slice from the persistent `request_*` tensors above. _tok_int64_bytes = self.max_tokens * 8 _tok_int32_bytes = self.max_tokens * 4 - _req_int32_bytes = self.max_requests * 4 + # Request-level fields are all 4 bytes wide (5 int32 + 2 float32 = 7 fields). + _req_4byte_bytes = self.max_requests * 4 # MHA section: 5 fields (int32) shared between GraphedMHAMetadata and # NonGraphedMHAMetadata. max_bs == max_requests. _mha_query_lengths_bytes = self.max_requests * 4 @@ -983,7 +988,7 @@ def initialize_all_tensors(self) -> None: _total_bytes = ( 2 * _tok_int64_bytes + 4 * _tok_int32_bytes - + 3 * _req_int32_bytes + + 7 * _req_4byte_bytes + _mha_query_lengths_bytes + _mha_cu_query_seq_lengths_bytes + _mha_kv_seq_lengths_bytes @@ -1043,17 +1048,40 @@ def initialize_all_tensors(self) -> None: # CPU (refreshed from persistent tensors in transfer_bookkeeping_to_gpu); # read-only on GPU via matching slots in ContextGPUView._buf. self._staging_request_in_prefill_status = self._cpu_bookkeeping_buf[ - _off : _off + _req_int32_bytes + _off : _off + _req_4byte_bytes ].view(torch.int32) - _off += _req_int32_bytes + _off += _req_4byte_bytes self._staging_request_query_lengths = self._cpu_bookkeeping_buf[ - _off : _off + _req_int32_bytes + _off : _off + _req_4byte_bytes ].view(torch.int32) - _off += _req_int32_bytes + _off += _req_4byte_bytes self._staging_request_kv_length_offsets = self._cpu_bookkeeping_buf[ - _off : _off + _req_int32_bytes + _off : _off + _req_4byte_bytes ].view(torch.int32) - _off += _req_int32_bytes + _off += _req_4byte_bytes + + # Sampling-parameter staging slots, refreshed from `active_request_metadata` + # in transfer_bookkeeping_to_gpu(). FlashInfer reads these via + # `gpu_view.{temperature, top_k, top_p}`. + self._staging_temperature = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view( + torch.float32 + ) + _off += _req_4byte_bytes + self._staging_top_k = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view( + torch.int32 + ) + _off += _req_4byte_bytes + self._staging_top_p = self._cpu_bookkeeping_buf[_off : _off + _req_4byte_bytes].view( + torch.float32 + ) + _off += _req_4byte_bytes + + # Per-request last-token row indices. Aliased with the matching gpu_view slot: + # build_active_slices/pad_active_slices populate this CPU view. + self.active_request_last_token_idxs = self._cpu_bookkeeping_buf[ + _off : _off + _req_4byte_bytes + ].view(torch.int32) + _off += _req_4byte_bytes # Static tensor addresses to make `last_token_logits` graphable with speculative decoding. max_logit_idxs = self.max_requests * (self.num_speculative_tokens + 1) @@ -1362,6 +1390,13 @@ def build_active_slices(self, batch_size: int): self.request_metadata[label][padded_slice], non_blocking=True ) + torch.cumsum( + self.request_query_lengths[padded_slice], + dim=0, + out=self.active_request_last_token_idxs[:batch_size], + ) + self.active_request_last_token_idxs[:batch_size].sub_(1) + def pad_active_slices(self): """Pad the active slices of specific tensors.""" active_request_count = self.total_request_count - self.paused_request_count @@ -1388,6 +1423,15 @@ def pad_active_slices(self): self.active_logit_idxs[active_decode_token_count + active_prefill_count :].zero_() + padding_request_slice = slice(active_request_count, self.padded_active_request_count) + + # Sampling metadata: pad with neutral defaults, so that the kernel early-exits. + self.active_request_metadata["temperature"][padding_request_slice].fill_(1.0) + self.active_request_metadata["top_k"][padding_request_slice].fill_(0) + self.active_request_metadata["top_p"][padding_request_slice].fill_(0.0) + # Padded gather indices fan in to row 0 harmlessly when used by FlashInfer. + self.active_request_last_token_idxs[padding_request_slice].fill_(0) + def append_key_value_cache(self, layer_number: int, key: Tensor, value: Tensor) -> None: """Append to KV cache. @@ -2226,7 +2270,7 @@ def transfer_bookkeeping_to_gpu(self) -> None: All copies use non_blocking=True with pinned CPU memory. CUDA stream ordering guarantees the forward pass sees completed transfers. - The 9 bookkeeping fields are backed by one contiguous pinned CPU buffer + The bookkeeping fields are backed by one contiguous pinned CPU buffer and one contiguous GPU buffer; a single cudaMemcpyAsync suffices. Request-level staging slots are refreshed from the persistent CPU tensors immediately before the H2D (GPU reads them at `[:n_active]` @@ -2237,9 +2281,9 @@ def transfer_bookkeeping_to_gpu(self) -> None: padded_active = max(n_active, self.padded_active_request_count) # Refresh request-level staging slots from the persistent CPU source. - # CPU-to-CPU slice assignment on pinned memory (~7.5 KB total for 3 - # int32 fields at max_requests=624). Negligible vs. the launch overhead - # we save by merging 9 H2D memcpys into 1. + # CPU-to-CPU slice assignment on pinned memory (~15 KB total for 6 + # 4-byte fields at max_requests=624). Negligible vs. the launch overhead + # we save by merging the H2D memcpys into 1. self._staging_request_in_prefill_status[:n_active] = self.request_in_prefill_status_tensor[ active_slice ] @@ -2247,6 +2291,15 @@ def transfer_bookkeeping_to_gpu(self) -> None: self._staging_request_kv_length_offsets[:n_active] = self.request_kv_length_offsets[ active_slice ] + # Sampling-parameter staging slots: read from `active_request_metadata`, + # which `build_active_slices` + `pad_active_slices` already populated for + # `[:padded_active]` (active values + neutral padding defaults). + self._staging_temperature[:padded_active] = self.active_request_metadata["temperature"][ + :padded_active + ] + self._staging_top_k[:padded_active] = self.active_request_metadata["top_k"][:padded_active] + self._staging_top_p[:padded_active] = self.active_request_metadata["top_p"][:padded_active] + # Full-iteration CUDA graphs may have captured GPU consumers with the # padded graph request count. Keep those padded staging rows bounded so # graph replay never builds indices from stale request lengths. diff --git a/megatron/core/inference/contexts/gpu_view.py b/megatron/core/inference/contexts/gpu_view.py index 17d6e19ea03..65c401163b0 100644 --- a/megatron/core/inference/contexts/gpu_view.py +++ b/megatron/core/inference/contexts/gpu_view.py @@ -17,11 +17,11 @@ class ContextGPUView: ``context.foo`` -> CPU (source of truth, used by bookkeeping) ``context.gpu_view.foo`` -> GPU (snapshot, used by forward pass) - Layout note: the 9 bookkeeping fields are backed by a single contiguous + Layout note: the bookkeeping fields are backed by a single contiguous ``uint8`` buffer (``self._buf``). Each field is a ``view(dtype)`` onto a slice of that buffer. This matches the pinned-CPU-buffer layout in :class:`DynamicInferenceContext` so that the per-step H2D transfer is a - single ``cudaMemcpyAsync`` instead of nine small ones. + single ``cudaMemcpyAsync`` instead of one per field. """ def __init__( @@ -39,7 +39,10 @@ def __init__( # max_mamba_chunks == 0). tok_int64_bytes = max_tokens * 8 # 2 fields of int64 = 8 bytes/elem tok_int32_bytes = max_tokens * 4 # 4 fields of int32 = 4 bytes/elem - req_int32_bytes = max_requests * 4 # 3 fields of int32 + # Request-level fields are all 4 bytes wide. 3 int32 (in_prefill_status, + # query_lengths, kv_length_offsets) + 1 int32 (top_k) + 2 float32 + # (temperature, top_p) + 1 int32 (active_request_last_token_idxs) = 7 fields. + req_4byte_bytes = max_requests * 4 # MHA section: 5 fields shared by both graphed and non-graphed MHAMetadata # (only one is active per step, so sharing storage is fine). @@ -90,7 +93,7 @@ def __init__( total_bytes = ( 2 * tok_int64_bytes + 4 * tok_int32_bytes - + 3 * req_int32_bytes + + 7 * req_4byte_bytes + mha_query_lengths_bytes + mha_cu_query_seq_lengths_bytes + mha_kv_seq_lengths_bytes @@ -128,12 +131,28 @@ def __init__( off += tok_int32_bytes # Request-level tensors (consumed by sampling, log-probs, speculative verification, MTP). - self.request_in_prefill_status = self._buf[off : off + req_int32_bytes].view(torch.int32) - off += req_int32_bytes - self.request_query_lengths = self._buf[off : off + req_int32_bytes].view(torch.int32) - off += req_int32_bytes - self.request_kv_length_offsets = self._buf[off : off + req_int32_bytes].view(torch.int32) - off += req_int32_bytes + self.request_in_prefill_status = self._buf[off : off + req_4byte_bytes].view(torch.int32) + off += req_4byte_bytes + self.request_query_lengths = self._buf[off : off + req_4byte_bytes].view(torch.int32) + off += req_4byte_bytes + self.request_kv_length_offsets = self._buf[off : off + req_4byte_bytes].view(torch.int32) + off += req_4byte_bytes + # Sampling parameters (consumed by FlashInfer sampling). + # Mirror the active slice of `active_request_metadata[{label}]`; + # padded slots get neutral defaults from `pad_active_slices` (T=1.0, top_k=0, top_p=0.0). + self.temperature = self._buf[off : off + req_4byte_bytes].view(torch.float32) + off += req_4byte_bytes + self.top_k = self._buf[off : off + req_4byte_bytes].view(torch.int32) + off += req_4byte_bytes + self.top_p = self._buf[off : off + req_4byte_bytes].view(torch.float32) + off += req_4byte_bytes + # Per-request last-token row indices (consumed by sampling kernels as `gather_indices`). + # The CPU side of this slot IS `context.active_request_last_token_idxs`, + # populated by `build_active_slices` and `pad_active_slices`. + self.active_request_last_token_idxs = self._buf[off : off + req_4byte_bytes].view( + torch.int32 + ) + off += req_4byte_bytes # MHA flash-attention metadata (shared between GraphedMHAMetadata and # NonGraphedMHAMetadata — only one is active per step). diff --git a/megatron/core/inference/engines/dynamic_engine.py b/megatron/core/inference/engines/dynamic_engine.py index ff5454bbaa3..fcee2c1daef 100644 --- a/megatron/core/inference/engines/dynamic_engine.py +++ b/megatron/core/inference/engines/dynamic_engine.py @@ -395,6 +395,12 @@ def create_cuda_graphs(self, reset_context: bool = True): with torch.inference_mode(): controller._dynamic_step_forward_logits(input_ids, position_ids) + if controller._sampling_backend == "flashinfer": + if controller.num_speculative_tokens > 0: + controller._dynamic_step_sample_logits_and_verify_tokens(input_ids) + else: + controller._dynamic_step_sample_logits() + # MTP CUDA graph warmup for this batch dimension. if mtp_warmup_enabled: n = cuda_graph_batch_dimension.req_count diff --git a/megatron/core/inference/sampling/__init__.py b/megatron/core/inference/sampling/__init__.py new file mode 100644 index 00000000000..b2941b33c9e --- /dev/null +++ b/megatron/core/inference/sampling/__init__.py @@ -0,0 +1,7 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from megatron.core.inference.sampling.base import Sampling +from megatron.core.inference.sampling.flashinfer_sampling import FlashInferSampling +from megatron.core.inference.sampling.torch_sampling import TorchSampling + +__all__ = ["Sampling", "TorchSampling", "FlashInferSampling"] diff --git a/megatron/core/inference/sampling/base.py b/megatron/core/inference/sampling/base.py new file mode 100644 index 00000000000..8aa4c416c27 --- /dev/null +++ b/megatron/core/inference/sampling/base.py @@ -0,0 +1,89 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from abc import ABC, abstractmethod +from typing import Any, Optional + +import torch +from torch import Tensor + + +class Sampling(ABC): + """Abstract base for inference sampling backends. + + Subclasses implement `sample_kernel`. CUDA graphs are added via `CudaGraphManager`. + """ + + @abstractmethod + def sample_kernel( + self, + logits: Tensor, + n: int, + context, + *, + gather_indices: Optional[Tensor] = None, + token_to_request_index: Optional[Tensor] = None, + eager: bool = False, + cache_key: Any = None, + ) -> Tensor: + """Sample `n` tokens from `logits` and return them. + + Args: + logits: Logits tensor of shape `[>=n, vocab_size]`. + n: Number of rows to sample. + context: The active DynamicInferenceContext. + 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. + eager, cache_key: Consumed by `CudaGraphManager` when it wraps this kernel. + + Returns: + Sampled token ids of shape `[n]`. Under CUDA graph replay, this is a static buffer. + """ + ... + + def sample_speculative( + self, + required_logits: Tensor, + num_decode: int, + num_prefill: int, + num_speculative_tokens: int, + context, + *, + gather_indices: Optional[Tensor] = None, + eager: bool = False, + cache_key: Any = None, + ) -> 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 `sample_kernel`. + The `sample_kernel` is forced eager so its own `CudaGraphManager` wrapper does not fire. + + 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). + """ + # CudaGraphManager consumes these args, if it exists. + del eager, cache_key + + n_spec = num_speculative_tokens + num_decode_tokens = num_decode * (1 + n_spec) + num_tokens = num_decode_tokens + num_prefill + device = required_logits.device + + token_to_request_index = torch.cat( + [ + torch.arange(num_decode, device=device).repeat_interleave( + 1 + n_spec, output_size=num_decode_tokens + ), + torch.arange(num_decode, num_decode + num_prefill, device=device), + ] + ) + return self.sample_kernel( + required_logits, + num_tokens, + context, + gather_indices=gather_indices, + token_to_request_index=token_to_request_index, + eager=True, + ) diff --git a/megatron/core/inference/sampling/flashinfer_sampling.py b/megatron/core/inference/sampling/flashinfer_sampling.py new file mode 100644 index 00000000000..c89093daeac --- /dev/null +++ b/megatron/core/inference/sampling/flashinfer_sampling.py @@ -0,0 +1,101 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from typing import Any, Optional + +import torch +from torch import Tensor + +try: + import flashinfer +except ImportError: + flashinfer = None + +from megatron.core.inference.sampling.base import Sampling +from megatron.core.transformer.cuda_graphs import CudaGraphManager + + +class FlashInferSampling(Sampling): + """Fused FlashInfer sampling, with optional CUDA graph capture/replay.""" + + def __init__( + self, vocab_size: int, rng: torch.Generator, config=None, enable_cuda_graph: bool = False + ) -> None: + self._vocab_size = vocab_size + self._rng = rng + if enable_cuda_graph and config is not None and config.cuda_graph_impl == "local": + CudaGraphManager( + config, + self, + function_name="sample_kernel", + need_backward=False, + inline_capture=True, + ) + CudaGraphManager( + config, + self, + function_name="sample_speculative", + need_backward=False, + inline_capture=True, + ) + + def sample_kernel( + self, + logits: Tensor, + n: int, + context, + *, + gather_indices: Optional[Tensor] = None, + token_to_request_index: Optional[Tensor] = None, + eager: bool = False, + cache_key: Any = None, + ) -> Tensor: + """FlashInfer fused top-k / top-p sampling kernel. + + Args: + logits: Logits tensor of shape `[>=n, vocab_size]`. + n: Number of rows to sample. + context: The active DynamicInferenceContext. + 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 (used by the speculative path). + eager, cache_key: Consumed by `CudaGraphManager` when it wraps this kernel. + + Returns: + Sampled token ids of shape `[n]`. Under CUDA graph replay, this is a static buffer. + """ + # CudaGraphManager consumes these args, if it exists. + del eager, cache_key + + # Read GPU sampling parameters from the per-step gpu_view mirror. The + # CPU source-of-truth (`active_request_metadata`) is pinned but resident + # on CPU, so reading it here would mix devices with `logits`. + gv = context.gpu_view + if token_to_request_index is None: + temperature = gv.temperature[:n] + top_k = gv.top_k[:n] + top_p = gv.top_p[:n] + else: + temperature = gv.temperature[token_to_request_index] + top_k = gv.top_k[token_to_request_index] + top_p = gv.top_p[token_to_request_index] + + # Clamp temperature to avoid division by 0. + temperature = temperature.clamp(min=1e-6) + if gather_indices is None: + scaled = logits[:n] / temperature.unsqueeze(1) + else: + scaled = logits[gather_indices[:n], :] / temperature.unsqueeze(1) + probs = torch.softmax(scaled, dim=-1) + + # Sentinel values disable filtering: + # top_k=vocab_size keeps all tokens, top_p=1.0 keeps the full probability mass. + # TODO: Consider changing the disable flags in the `InferenceRequest`. + top_k_safe = top_k.masked_fill(top_k == 0, self._vocab_size) + top_p_safe = top_p.masked_fill(top_p == 0.0, 1.0) + output = torch.empty(n, device=logits.device, dtype=torch.int64) + output.copy_( + flashinfer.sampling.top_k_top_p_sampling_from_probs( + probs, top_k_safe, top_p_safe, generator=self._rng + ) + ) + return output diff --git a/megatron/core/inference/sampling/torch_sampling.py b/megatron/core/inference/sampling/torch_sampling.py new file mode 100644 index 00000000000..79491add5ab --- /dev/null +++ b/megatron/core/inference/sampling/torch_sampling.py @@ -0,0 +1,167 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from collections import defaultdict +from typing import Any, List, Optional, Tuple + +import torch +from torch import Tensor + +from megatron.core.inference.sampling.base import Sampling + + +class TorchSampling(Sampling): + """Sampling via bucketed `torch.multinomial`. + + Groups requests into unique buckets by `(temperature, top_k, top_p)` for separate launches. + """ + + def __init__(self, rng: torch.Generator, vocab_size: int) -> None: + self._rng = rng + self._vocab_size = vocab_size + + @staticmethod + def sample_from_logits( + last_token_logits: Tensor, + temperature: float, + top_k: int, + top_p: float, + *, + generator: torch.Generator, + vocab_size: Optional[int] = None, + ) -> Tensor: + """Sample tokens from logits with temperature, top-k, and top-p filtering. + + Shared between dynamic batching and static batching. + + Args: + 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]`. + """ + assert isinstance(top_p, float) + assert isinstance(top_k, int) + assert not (top_k > 0 and top_p > 0.0), "Cannot have top-p and top-k both greater than zero" + assert top_p <= 1.0, "top-p should be in (0,1]" + + def modify_logits_for_top_k_filtering(logits, top_k): + """Set the logits for none top-k values to -inf.""" + filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None] + logits.masked_fill_(filter_, float("-Inf")) + + def modify_logits_for_top_p_filtering(logits, top_p): + """Set the logits for none top-p values to -inf.""" + sorted_logits, sorted_indices = torch.sort(logits, descending=True) + cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1) + + filter_ = cumulative_probs > top_p + # Clone needed: filter_[:, 1:] and filter_[:, :-1] are overlapping views; + # without clone, each write would corrupt the next read during the shift. + filter_[:, 1:] = filter_[:, :-1].clone() + filter_[..., 0] = 0 + + filter_ = filter_.scatter(1, sorted_indices, filter_) + logits.masked_fill_(filter_, float("-Inf")) + + if top_k == 1: + return torch.argmax(last_token_logits, dim=-1) + + # Clone needed: .div_() and masked_fill_() below modify in-place. + last_token_logits = last_token_logits.clone() + if temperature != 1.0: + last_token_logits.div_(temperature) + if top_k > 1: + assert top_k <= last_token_logits.size(1), "top-k is larger than logit size." + if vocab_size: + assert top_k < vocab_size, "top-k is larger than vocab size." + modify_logits_for_top_k_filtering(last_token_logits, top_k) + elif top_p > 0.0: + modify_logits_for_top_p_filtering(last_token_logits, top_p) + + probabilities = last_token_logits.softmax(dim=-1) + sampled = torch.multinomial(probabilities, num_samples=1, generator=generator).view(-1) + + if vocab_size: + sampled = torch.clamp(sampled, min=0, max=(vocab_size - 1)) + + return sampled + + def sample_kernel( + self, + logits: Tensor, + n: int, + context, + *, + gather_indices: Optional[Tensor] = None, + token_to_request_index: Optional[Tensor] = None, + eager: bool = False, + cache_key: Any = None, + ) -> Tensor: + """Bucket active requests by `(temperature, top_k, top_p)` and sample each bucket. + + Args: + logits: Logits tensor of shape `[>=n, vocab_size]`. + n: Number of rows to sample. + context: The active DynamicInferenceContext. + 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). + eager: Accepted for API symmetry; ignored (TorchSampling has no graph wrapper). + cache_key: Accepted for API symmetry; ignored. + + Returns: + Sampled token ids of shape `[n]`. + """ + # CudaGraphManager consumes these args, if it exists. + del eager, cache_key + + # Group active requests into sampling buckets by (temperature, top_k, top_p). + active_request_count = context.total_request_count - context.paused_request_count + md = context.active_request_metadata + device = torch.cuda.current_device() + + bucket_map: dict = defaultdict(list) + temp = md["temperature"][:active_request_count].tolist() + top_k = md["top_k"][:active_request_count].tolist() + top_p = md["top_p"][:active_request_count].tolist() + for request_index, (t, k, p) in enumerate(zip(temp, top_k, top_p)): + bucket_map[(t, k, p)].append(request_index) + + buckets: List[Tuple] = [(indices, *params) for params, indices in bucket_map.items()] + bucket_index_tensors: List[Tensor] = [ + torch.tensor(indices, device=device, dtype=torch.long) for indices, *_ in buckets + ] + + if gather_indices is not None: + logits = logits[gather_indices[:n], :] + + output = torch.empty(n, device=logits.device, dtype=torch.int64) + token_list = [] + indices_list = [] + for idx_tensor, (_, temp, top_k, top_p) in zip(bucket_index_tensors, buckets): + if token_to_request_index is None: + row_indices = idx_tensor + else: + row_indices = torch.where(torch.isin(token_to_request_index, idx_tensor))[0] + token_list.append( + TorchSampling.sample_from_logits( + logits[row_indices, :], + temp, + top_k, + top_p, + generator=self._rng, + vocab_size=self._vocab_size, + ) + ) + indices_list.append(row_indices) + + sampled_tokens = torch.cat(token_list, dim=0) + sampled_indices = torch.cat(indices_list, dim=0) + output[sampled_indices] = sampled_tokens + return output diff --git a/megatron/core/inference/text_generation_controllers/text_generation_controller.py b/megatron/core/inference/text_generation_controllers/text_generation_controller.py index faa9e5babd6..87edddea566 100644 --- a/megatron/core/inference/text_generation_controllers/text_generation_controller.py +++ b/megatron/core/inference/text_generation_controllers/text_generation_controller.py @@ -60,6 +60,7 @@ HAVE_TE = False from megatron.core.inference.batch_dimensions_utils import InferenceBatchDimensions +from megatron.core.inference.sampling import FlashInferSampling, Sampling, TorchSampling from megatron.core.inference.text_generation_controllers.mtp_utils_pytorch import rewind_kv_cache from megatron.core.inference.text_generation_controllers.mtp_utils_triton import ( mamba_state_selective_copy, @@ -154,7 +155,7 @@ def _init_dynamic_sampling_tensors(self): device = torch.cuda.current_device() logits_dtype = self.inference_wrapped_model.config.params_dtype - self._sampling_backend = "torch" + self._sampling_backend = context.config.sampling_backend self._enable_cuda_graph = self.model_config.cuda_graph_impl == "local" # Initialize bookkeeping tensors. @@ -164,11 +165,24 @@ def _init_dynamic_sampling_tensors(self): ) else: self._all_logits_cuda = None - self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device) - - # Used for inefficient torch sampling. - if self._sampling_backend == "torch": - self._torch_sampling_buckets: List[Tuple] = [] + # Speculative path: + # - `self._sampled_tokens_cuda` is pre-allocated by `_init_mtp_sampling_tensors`. + # - The tensor cannot be reused between the Triton kernel and the sampling graph. + # Non-speculative path: + # - `self._sampled_tokens_cuda` is rebound to the output of `sample_kernel`, + # which uses CudaGraphManager syntactic sugar to keep it as a static tensor. + self._sampled_tokens_cuda = None + + # Sampling backend: provides the sampling kernel. + if self._sampling_backend == "flashinfer": + self._sampling: Sampling = FlashInferSampling( + self.vocab_size, + self.sampling_rng, + config=self.model_config, + enable_cuda_graph=self._enable_cuda_graph, + ) + else: + self._sampling: Sampling = TorchSampling(self.sampling_rng, self.vocab_size) # Cache values that are constant across inference steps. self._unwrapped_model = unwrap_model(self.inference_wrapped_model.model) @@ -192,6 +206,7 @@ def _init_mtp_sampling_tensors(self): context = self.inference_wrapped_model.inference_context max_requests = context.max_requests device = torch.cuda.current_device() + self._sampled_tokens_cuda = torch.empty(max_requests, dtype=torch.int64, device=device) self._sampled_mtp_tokens_cuda = torch.empty( [self.num_speculative_tokens, max_requests], dtype=torch.int64, device=device ) @@ -321,95 +336,6 @@ def detokenize_generations( return text, prompts_plus_generations_segments - def _torch_sampling_func( - self, - last_token_logits: torch.Tensor, - temperature: float, - top_k: int, - top_p: float, - vocab_size: Optional[int] = None, - ): - """Samples the logits to generate outputs - - Given the logits of the last token, this function samples it - according to the parameters defined in sampling_params - and returns the samples. If sampling parameters top_n_logprobs > 0 - at each step it also updates the top_n_logprobs dict. - - Args: - last_token_logits (torch.Tensor): The last token logits. A tensor of - size [batch_size, vocab_size]. - temperature (float): The temperature to use for sampling. - top_k (int): The top-k value to use for sampling. - top_p (float): The top-p value to use for sampling. - vocab_size (int): Obtained from the tokenizer. Defaults to None. - - Returns: - sampled_logits (torch.Tensor): 1D tensor with [batch_size] elements - """ - assert isinstance(top_p, float) - assert isinstance(top_k, int) - assert not (top_k > 0 and top_p > 0.0), "Cannot have top-p and top-k both greater than zero" - assert top_p <= 1.0, "top-p should be in (0,1]" - - def modify_logits_for_top_k_filtering(logits, top_k): - """Set the logits for none top-k values to -inf.""" - filter_ = logits < torch.topk(logits, top_k)[0][..., -1, None] - logits.masked_fill_(filter_, float("-Inf")) - - def modify_logits_for_top_p_filtering(logits, top_p): - """Set the logits for none top-p values to -inf.""" - # First sort and calculate cumulative sum of probabilities. - sorted_logits, sorted_indices = torch.sort(logits, descending=True) - cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1) - - # Filteration based on the cumulative sum. - filter_ = cumulative_probs > top_p - # This shift by 1 is weird and I cannot justify it. This existed - # in the original implementation: - # https://github.com/ari-holtzman/degen/blob/master/gen.py - # and I guess it is needed so keeping it for now. - # Clone needed: filter_[:, 1:] and filter_[:, :-1] are overlapping views; - # without clone, each write would corrupt the next read during the shift. - filter_[:, 1:] = filter_[:, :-1].clone() - # Make sure we at least have one token to select from. - filter_[..., 0] = 0 - - # Fill in the filtered part - filter_ = filter_.scatter(1, sorted_indices, filter_) - logits.masked_fill_(filter_, float("-Inf")) - - # Greedy sampling - if top_k == 1: - sampled_logits = torch.argmax(last_token_logits, dim=-1) - else: - # Clone needed: .div_() and masked_fill_() below modify in-place, - # which would mutate the caller's tensor without this clone. - last_token_logits = last_token_logits.clone() - if temperature != 1.0: - last_token_logits.div_(temperature) - if top_k > 1: - assert top_k <= last_token_logits.size(1), "top-k is larger than logit size." - if vocab_size: - assert top_k < vocab_size, "top-k is larger than vocab size." - modify_logits_for_top_k_filtering(last_token_logits, top_k) - - elif top_p > 0.0: - modify_logits_for_top_p_filtering(last_token_logits, top_p) - - # After filtering, we need to recalculate the distribution. - probabilities = last_token_logits.softmax(dim=-1) - - sampled_logits = torch.multinomial( - probabilities, num_samples=1, generator=self.sampling_rng - ).view(-1) - - # If vocab size is provided, make sure the samples are in in the range [0, vocab-size). - if vocab_size: - sampled_logits = torch.clamp(sampled_logits, min=0, max=(vocab_size - 1)) - - return sampled_logits - def sample_from_logits( self, last_token_logits: torch.Tensor, @@ -496,7 +422,14 @@ def sample_from_logits( top_k = sampling_params.top_k temperature = sampling_params.temperature - return self._torch_sampling_func(last_token_logits, temperature, top_k, top_p, vocab_size) + return TorchSampling.sample_from_logits( + last_token_logits, + temperature, + top_k, + top_p, + generator=self.sampling_rng, + vocab_size=vocab_size, + ) def update_generation_status( self, @@ -731,36 +664,6 @@ def _dynamic_step_forward_logits(self, input_ids: Tensor, position_ids: Tensor): else: self._all_logits_cuda = logits - def _dynamic_step_sample_bookkeeping(self): - """Perform bookkeeping necessary to sample logits for dynamic batching.""" - context = self.inference_wrapped_model.inference_context - active_request_count = context.total_request_count - context.paused_request_count - - if self._sampling_backend == "torch": - # Bucketize the core sampling parameters. - # Doing so via list comprehension is orders of magnitude faster than via torch. - bucket_map = defaultdict(list) - - # Shorthands for the dictionary comprehension. - temp = context.active_request_metadata["temperature"][:active_request_count].tolist() - top_k = context.active_request_metadata["top_k"][:active_request_count].tolist() - top_p = context.active_request_metadata["top_p"][:active_request_count].tolist() - - for request_index, (t, k, p) in enumerate(zip(temp, top_k, top_p)): - sampling_params = (t, k, p) - bucket_map[sampling_params].append(request_index) - - # Just unpack the key directly! - device = torch.cuda.current_device() - self._torch_sampling_buckets = [ - (indices, *sampling_params) for sampling_params, indices in bucket_map.items() - ] - # Pre-compute index tensors on GPU to avoid per-step H2D copies. - self._torch_sampling_bucket_index_tensors = [ - torch.tensor(indices, device=device, dtype=torch.long) - for indices, *_ in self._torch_sampling_buckets - ] - def _rewind_kv_cache(self) -> tuple: """Update the KV cache bookkeeping for speculative decoding. @@ -836,18 +739,12 @@ def _sample_from_logits_2d(self, logits_2d: Tensor) -> Tensor: Returns: Tensor: Sampled tokens of shape [num_requests]. """ - spec_token_list = [] - for idx_tensor, (_, temp, top_k, top_p) in zip( - self._torch_sampling_bucket_index_tensors, self._torch_sampling_buckets - ): - spec_token_list.append( - self._torch_sampling_func(logits_2d[idx_tensor, :], temp, top_k, top_p) - ) - - spec_tokens = torch.empty(logits_2d.shape[0], device=logits_2d.device, dtype=torch.int64) - for tokens, indices in zip(spec_token_list, self._torch_sampling_bucket_index_tensors): - spec_tokens[indices] = tokens - return spec_tokens + return self._sampling.sample_kernel( + logits_2d, + logits_2d.shape[0], + self.inference_wrapped_model.inference_context, + eager=True, + ) def _compute_serial_mtp_and_sample(self): """Compute MTP logits serially after verification and sample speculative tokens. @@ -993,61 +890,6 @@ def _compute_serial_mtp_and_sample(self): if has_mtp: del unwrapped_model._decoder_hidden_states_cache - def _sample_speculative_logits( - self, required_logits: Tensor, request_in_prefill_status_tensor: Tensor - ) -> tuple: - """Sample tokens from logits using sampling buckets. - - For torch sampling buckets: [request_indices, temp, top_k, top_p] - - Example with 5 requests: - token_to_request_idx : [ 0 0 0 | 1 1 1 | 2 2 2 | 3 | 4 ] - required_logits : [ a5l a6l a7l | b3l b4l b5l | c6l c7l c8l | d2l | e4l ] # Shape [11, vocab_size] - - Sampling buckets: [[[0,2], temp1, top_k1, top_p1], [[1], temp3, top_k3, top_p3], [[3, 4], temp2, top_k2, top_p2]] - - Final output tokens : [a5s a6s a7s c6s c7s c8s b3s b4s b5s d2s e4s] # Shape [11] - (Rearranged from sampling bucket order back to input order using token_order) - - Returns: - tuple: (output_tokens, repeats) where output_tokens has shape [total_required_tokens] - """ - # request_in_prefill_status_tensor is already on GPU (from gpu_view). - repeats = torch.where( - request_in_prefill_status_tensor == 0, 1 + self.num_speculative_tokens, 1 - ) - token_to_request_index = torch.repeat_interleave( - torch.arange( - len(request_in_prefill_status_tensor), - device=request_in_prefill_status_tensor.device, - ), - repeats, - ) - - output_tokens_jumbled_list = [] - token_order_list = [] - - for idx_tensor, (_, temp, top_k, top_p) in zip( - self._torch_sampling_bucket_index_tensors, self._torch_sampling_buckets - ): - required_indices = torch.where(torch.isin(token_to_request_index, idx_tensor))[0] - output_tokens_jumbled_list.append( - self._torch_sampling_func(required_logits[required_indices, :], temp, top_k, top_p) - ) - token_order_list.append(required_indices) - - output_tokens_jumbled = torch.cat(output_tokens_jumbled_list, dim=0) - output_tokens = torch.empty( - len(output_tokens_jumbled), - device=output_tokens_jumbled.device, - dtype=output_tokens_jumbled.dtype, - ) - token_order = torch.cat(token_order_list, dim=0) - # Rearrange output tokens from sampling_bucket request order back to input ids order - output_tokens[token_order] = output_tokens_jumbled - - return output_tokens, repeats - def _verify_speculative_tokens( self, output_tokens: Tensor, @@ -1072,32 +914,59 @@ def _dynamic_step_sample_logits_and_verify_tokens(self, input_ids: Tensor): context = self.inference_wrapped_model.inference_context active_request_count = context.total_request_count - context.paused_request_count - # Use gpu_view for data consumed by GPU operations (sampling, verification). - request_in_prefill_status_tensor = context.gpu_view.request_in_prefill_status[ - :active_request_count - ] + # Sampling-side request counts: padded when running a captured graph. + # Verify uses the actual counts so the Triton kernels operate on the real workload. + use_graph_for_sampling = ( + self._sampling_backend == "flashinfer" + and self._enable_cuda_graph + and context.using_cuda_graph_this_step() + ) + if use_graph_for_sampling: + sample_num_decode = context.padded_batch_dimensions.decode_req_count + sample_num_prefill = context.padded_batch_dimensions.prefill_req_count + else: + sample_num_decode = context.num_decode_requests + sample_num_prefill = context.num_prefill_requests - # Get the logit indices for tokens that need sampling. - # These indices are always needed for input_ids slicing and tracking - # accepted sequence positions, even when logits are pre-sliced. + # Logit indices for tokens that need sampling. + # Padded under graph capture so the captured `gather_indices` input has a stable shape. + # Padded slots resolve to row 0; verify and prepare-next read only the actual prefix, + # so the padded-row samples produced by the captured kernel are discarded. nvtx_range_push("mtp-spec-decoding/verify/logit-indices") # Use pre-allocated buffer for CUDA graph compatibility. logits = self._all_logits_cuda + # `speculative_required_logit_indices()` already returns padded indices when + # running a captured graph (`num_last_token_logits` uses the padded counts and + # `pad_active_slices` zero-pads the trailing slots), so the call site does not + # need to re-pad here. required_logit_indices = context.speculative_required_logit_indices() if context.config.materialize_only_last_token_logits: # last_token_logits already selected exactly the required positions. - required_logits = logits.squeeze(0) + sample_logits = logits.squeeze(0) + sample_gather_indices = None else: - required_logits = logits.squeeze(0)[ - required_logit_indices, : - ] # Shape [num_required, vocab_size] + # Push the gather inside the captured kernel: + # pass the full per-token logits buffer (constant shape) plus the padded indices. + sample_logits = logits.squeeze(0) + sample_gather_indices = required_logit_indices nvtx_range_pop("mtp-spec-decoding/verify/logit-indices") # Sample tokens from logits nvtx_range_push("mtp-spec-decoding/verify/sample") - output_tokens, repeats = self._sample_speculative_logits( - required_logits, request_in_prefill_status_tensor + output_tokens = self._sampling.sample_speculative( + sample_logits, + sample_num_decode, + sample_num_prefill, + self.num_speculative_tokens, + context, + gather_indices=sample_gather_indices, + eager=not use_graph_for_sampling, + cache_key=( + ("sample_speculative", sample_num_decode, sample_num_prefill) + if use_graph_for_sampling + else None + ), ) nvtx_range_pop("mtp-spec-decoding/verify/sample") @@ -1168,40 +1037,31 @@ def _dynamic_step_sample_logits(self): # TODO(ksanthanam): Evaluate whether it makes more sense to sample on 1 rank # and then broadcast the sampled tokens rather than broadcasting the raw logits. - # Last token logits. context = self.inference_wrapped_model.inference_context active_request_count = context.total_request_count - context.paused_request_count - - if context.config.materialize_only_last_token_logits: - # When materialize_only_last_token_logits is true, last_token_logits is - # already called in the forward pass of GPT. - required_token_logits = self._all_logits_cuda.squeeze(0)[:active_request_count, :] - else: - required_token_logits = context.last_token_logits( - self._all_logits_cuda[:, : context.padded_active_token_count, :] - ) - - if self._sampling_backend == "torch": - # Concatenate the outputs once to prevent repeated small writes. - token_list = [] - indices_list = [] - - # e.g torch sample buckets will be - # i.e (for all unique comibnation of t, topk, topk what are the associated - # requests indices (based on the active slices) - # [ [req at index 0, req at index 2], t1, topk1, topp1 ]] - # [ [req at index 1, req at index 3, req at index 4] , t2, topk2, topp2] - for indices, temp, top_k, top_p in self._torch_sampling_buckets: - token_list.append( - self._torch_sampling_func(required_token_logits[indices, :], temp, top_k, top_p) - ) - indices_list.append(torch.tensor(indices)) - - # Single write to the output tensor. - sampled_tokens = torch.cat(token_list, dim=0) - sampled_indices = torch.cat(indices_list, dim=0) - - self._sampled_tokens_cuda[sampled_indices] = sampled_tokens + use_graph = ( + self._sampling_backend == "flashinfer" + and self._enable_cuda_graph + and context.using_cuda_graph_this_step() + ) + # Padded count when running a captured graph (cache key buckets); actual otherwise. + n = context.padded_active_request_count if use_graph else active_request_count + # When `materialize_only_last_token_logits` is true the forward pass already + # selected the right rows. Otherwise we point the kernel at the per-request + # last-token positions via `gather_indices`; padded slots safely fan in to row 0. + gather_indices = ( + None + if context.config.materialize_only_last_token_logits + else context.gpu_view.active_request_last_token_idxs + ) + self._sampled_tokens_cuda = self._sampling.sample_kernel( + self._all_logits_cuda.squeeze(0), + n, + context, + gather_indices=gather_indices, + eager=not use_graph, + cache_key=("sample", n) if use_graph else None, + ) def _dynamic_step_log_probs_bookkeeping(self) -> Tuple[bool, bool]: """Perform bookkeeping necessary to compute log probs for dynamic batching. @@ -1901,8 +1761,6 @@ async def async_generate_output_tokens_dynamic_batch( range_push("sampling") return_log_probs, return_top_n_logprobs = self._dynamic_step_log_probs_bookkeeping() - self._dynamic_step_sample_bookkeeping() - if self.num_speculative_tokens > 0: # Phase 1: Verify speculative tokens using base logits only. nvtx_range_push("mtp-spec-decoding/verify") diff --git a/megatron/inference/utils.py b/megatron/inference/utils.py index bb65c754ab1..559ac57f496 100644 --- a/megatron/inference/utils.py +++ b/megatron/inference/utils.py @@ -374,6 +374,7 @@ def get_inference_config_from_model_and_args(model: MegatronModule, args): logging_step_interval=args.inference_logging_step_interval, num_speculative_tokens=args.num_speculative_tokens, use_synchronous_zmq_collectives=args.inference_use_synchronous_zmq_collectives, + sampling_backend=args.inference_dynamic_batching_sampling_backend, ) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 17b27187b2e..e79334b7c03 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1007,6 +1007,16 @@ def validate_args(args, defaults={}): ): raise ValueError("MXFP8 with inference optimized layers requires FlashInfer >= 0.6.4") + if args.inference_dynamic_batching_sampling_backend == 'flashinfer': + try: + import flashinfer # noqa: F401 + except ImportError as e: + raise ImportError( + "--inference-dynamic-batching-sampling-backend=flashinfer requires " + "the flashinfer package; install it or pass " + "--inference-dynamic-batching-sampling-backend=torch." + ) from e + if args.use_megatron_fsdp: # NOTE: The flag `use_custom_fsdp` is deprecated and will be removed in future versions. # Please use `use_megatron_fsdp` instead, as all functionality will be migrated there. @@ -1994,6 +2004,12 @@ def _add_inference_args(parser): group.add_argument('--inference-dynamic-batching-cuda-graph-mixed-prefill-count', type=int, default=16, help='Number of mixed prefill requests to capture in a cuda graph.') + group.add_argument('--inference-dynamic-batching-sampling-backend', + type=str, default='torch', + choices=['torch', 'flashinfer'], + help='Which sampling kernels to use during inference. ' + 'Falls back to "torch" with a warning if "flashinfer" ' + 'is requested but the package is not installed.') group.add_argument('--inference-logging-step-interval', type=int, default=0, help='Step interval for logging inference metrics. ' 'Default to 0 to disable inference logging.') diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/golden_values_dev_dgx_h100.json new file mode 100644 index 00000000000..956754f44c1 --- /dev/null +++ b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/golden_values_dev_dgx_h100.json @@ -0,0 +1,286 @@ +{ + "0": { + "input_prompt": "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies.", + "generated_text": " You are not alone. You are not alone. You are not alone. You are not alone. You are not alone. You are not alone.", + "generated_tokens": [ + 3213, + 1584, + 1605, + 9412, + 1046, + 3213, + 1584, + 1605, + 9412, + 1046, + 3213, + 1584, + 1605, + 9412, + 1046, + 3213, + 1584, + 1605, + 9412, + 1046, + 3213, + 1584, + 1605, + 9412, + 1046, + 3213, + 1584, + 1605, + 9412, + 1046 + ], + "latency": 1.878054141998291, + "ttft": 0.07786321640014648, + "cuda_graph_request_count_map": null, + "step_count": 30, + "top_n_logprobs": null, + "prompt_top_n_logprobs": null, + "prompt_logprobs": [ + -9.498085021972656, + -3.787536859512329, + -3.0404648780822754, + -1.7445809841156006, + -0.29672086238861084, + -1.3661342859268188, + -2.3458175659179688, + -1.83931303024292, + -1.4894113540649414, + -6.440437316894531, + -0.8176816701889038, + -1.790361762046814, + -3.662419557571411, + -3.7036173343658447, + -1.6009434461593628, + -1.5501081943511963, + -2.846059799194336, + -6.732302665710449, + -0.06605878472328186, + -1.334327220916748, + -6.083745002746582, + -9.440131187438965, + -10.473882675170898, + -1.5964821577072144, + -4.702763557434082, + -0.7514524459838867, + -2.1461901664733887, + -0.012340382672846317, + -0.03605639934539795, + -3.0907557010650635, + -8.744739532470703, + -1.5410845279693604, + -5.84979772567749, + -3.0918972492218018, + -3.9814329147338867, + -3.78017520904541, + -2.5227086544036865, + -2.258594036102295, + -0.4719255566596985, + -1.0329649448394775, + -5.3284382820129395, + -8.25335693359375, + -0.015789249911904335, + -2.854100227355957, + -1.2236379384994507, + -3.905193328857422, + -0.9268187284469604, + -0.0030202509369701147, + -3.224249839782715, + -11.11172103881836, + -3.8121743202209473, + -2.3400487899780273, + -4.672845363616943, + -0.09729652851819992, + -0.06232408434152603, + -1.336004614830017, + -2.054157257080078, + -4.390933036804199, + -0.44248226284980774, + -3.9417736530303955, + -0.5888474583625793, + -0.26697415113449097, + -2.9271092414855957, + -13.515066146850586, + -0.10294333100318909, + -3.5007452964782715, + -0.8535972237586975, + -5.173652648925781, + -0.330394983291626, + -2.304553508758545, + -0.5418462753295898, + -1.300589919090271, + -4.9136152267456055, + -15.558022499084473, + -4.918652534484863, + -0.22206512093544006, + -6.589188575744629, + -0.9015690684318542, + -2.2228457927703857, + -1.8689247369766235, + -0.2006368339061737, + -5.918689727783203, + -0.006355076562613249, + -7.532094955444336, + -3.2708187103271484, + -3.743263006210327, + -2.011824131011963 + ], + "generated_logprobs": [ + -3.0331737995147705, + -1.9080564975738525, + -2.52506947517395, + -2.325258493423462, + -1.180279016494751, + -1.1824196577072144, + -0.39788734912872314, + -1.110222578048706, + -1.5034958124160767, + -0.9765141606330872, + -0.9300433397293091, + -0.15196305513381958, + -0.2200203537940979, + -0.06051275506615639, + -0.6840062737464905, + -1.0964292287826538, + -0.17654964327812195, + -0.18547140061855316, + -0.06710249185562134, + -0.4758152365684509, + -0.6657928228378296, + -0.10342729091644287, + -0.10059614479541779, + -0.046978313475847244, + -0.410809725522995, + -0.428723007440567, + -0.06053968518972397, + -0.06518109142780304, + -0.030038274824619293, + -0.3271780014038086 + ], + "logprobs": [ + -9.498085021972656, + -3.787536859512329, + -3.0404648780822754, + -1.7445809841156006, + -0.29672086238861084, + -1.3661342859268188, + -2.3458175659179688, + -1.83931303024292, + -1.4894113540649414, + -6.440437316894531, + -0.8176816701889038, + -1.790361762046814, + -3.662419557571411, + -3.7036173343658447, + -1.6009434461593628, + -1.5501081943511963, + -2.846059799194336, + -6.732302665710449, + -0.06605878472328186, + -1.334327220916748, + -6.083745002746582, + -9.440131187438965, + -10.473882675170898, + -1.5964821577072144, + -4.702763557434082, + -0.7514524459838867, + -2.1461901664733887, + -0.012340382672846317, + -0.03605639934539795, + -3.0907557010650635, + -8.744739532470703, + -1.5410845279693604, + -5.84979772567749, + -3.0918972492218018, + -3.9814329147338867, + -3.78017520904541, + -2.5227086544036865, + -2.258594036102295, + -0.4719255566596985, + -1.0329649448394775, + -5.3284382820129395, + -8.25335693359375, + -0.015789249911904335, + -2.854100227355957, + -1.2236379384994507, + -3.905193328857422, + -0.9268187284469604, + -0.0030202509369701147, + -3.224249839782715, + -11.11172103881836, + -3.8121743202209473, + -2.3400487899780273, + -4.672845363616943, + -0.09729652851819992, + -0.06232408434152603, + -1.336004614830017, + -2.054157257080078, + -4.390933036804199, + -0.44248226284980774, + -3.9417736530303955, + -0.5888474583625793, + -0.26697415113449097, + -2.9271092414855957, + -13.515066146850586, + -0.10294333100318909, + -3.5007452964782715, + -0.8535972237586975, + -5.173652648925781, + -0.330394983291626, + -2.304553508758545, + -0.5418462753295898, + -1.300589919090271, + -4.9136152267456055, + -15.558022499084473, + -4.918652534484863, + -0.22206512093544006, + -6.589188575744629, + -0.9015690684318542, + -2.2228457927703857, + -1.8689247369766235, + -0.2006368339061737, + -5.918689727783203, + -0.006355076562613249, + -7.532094955444336, + -3.2708187103271484, + -3.743263006210327, + -2.011824131011963, + -3.0331737995147705, + -1.9080564975738525, + -2.52506947517395, + -2.325258493423462, + -1.180279016494751, + -1.1824196577072144, + -0.39788734912872314, + -1.110222578048706, + -1.5034958124160767, + -0.9765141606330872, + -0.9300433397293091, + -0.15196305513381958, + -0.2200203537940979, + -0.06051275506615639, + -0.6840062737464905, + -1.0964292287826538, + -0.17654964327812195, + -0.18547140061855316, + -0.06710249185562134, + -0.4758152365684509, + -0.6657928228378296, + -0.10342729091644287, + -0.10059614479541779, + -0.046978313475847244, + -0.410809725522995, + -0.428723007440567, + -0.06053968518972397, + -0.06518109142780304, + -0.030038274824619293, + -0.3271780014038086 + ] + }, + "mem-max-allocated-bytes": 53350692864, + "lifetime_prefill_token_count": 88 +} \ No newline at end of file diff --git a/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/model_config.yaml b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/model_config.yaml new file mode 100644 index 00000000000..e989be22f7e --- /dev/null +++ b/tests/functional_tests/test_cases/hybrid/hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer/model_config.yaml @@ -0,0 +1,76 @@ +ENV_VARS: + CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 + NCCL_ALGO: Ring + CUBLAS_WORKSPACE_CONFIG: :4096:8 + TRITON_CACHE_AUTOTUNING: 0 + MAMBA_DETERMINISTIC: 1 +TEST_TYPE: frozen-start +MODE: inference +MODEL_ARGS: + --log-num-zeros-in-grad: true + --log-validation-ppl-to-tensorboard: true + --log-timers-to-tensorboard: true + --log-memory-to-tensorboard: true + --timing-log-level: 0 + --load: ${CHECKPOINT_LOAD_PATH}/model/mamba_hybrid_2b/dcp/mcore-v1_bf16/checkpoint + --tokenizer-model: ${CHECKPOINT_LOAD_PATH}/model/mamba_hybrid_2b/dcp/mcore-v1_bf16/multiMixV8.gpt4o_nc_sd.500000.128k.vocab.json + --tokenizer-type: TikTokenizer + --tiktoken-pattern: v2 + --distributed-backend: nccl + --log-interval: 1 + --transformer-impl: transformer_engine + --tensor-model-parallel-size: 1 + --pipeline-model-parallel-size: 1 + --expert-model-parallel-size: 1 + --use-mcore-models: true + --model-provider: hybrid + --init-method-std: 0.0198 + --untie-embeddings-and-output-weights: true + --disable-bias-linear: true + --init-method-std: 0.014 + --position-embedding-type: none + --hidden-size: 2048 + --ffn-hidden-size: 11264 + --num-attention-heads: 16 + --kv-channels: 128 + --hybrid-layer-pattern: M-M-M-M*-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M- + --spec: megatron.core.models.hybrid.hybrid_layer_specs hybrid_stack_spec + --normalization: RMSNorm + --swiglu: true + --attention-dropout: 0.0 + --hidden-dropout: 0.0 + --seq-length: 4096 + --max-position-embeddings: 4096 + --micro-batch-size: 1 + --ckpt-format: torch_dist + --ckpt-fully-parallel-save: true + --ckpt-fully-parallel-load: true + --ckpt-assume-constant-structure: true + --dist-ckpt-strictness: log_unexpected + --bf16: true + --attention-backend: flash + --no-create-attention-mask-in-dataloader: true + --num-workers: 8 + --use-checkpoint-args: true + --no-use-tokenizer-model-from-checkpoint-args: true + --no-load-optim: true + --deterministic-mode: true + --save-interval: 2000 + --temperature: 1.0 + --top_k: 1 + --inference-dynamic-batching-sampling-backend: flashinfer + --return-log-probs: true + --num-tokens-to-generate: 30 + --max-tokens-to-oom: 3600000 + --inference-max-seq-length: 4096 + --output-path: ${INFERENCE_OUTPUT_PATH} + --prompts: "Time travel to 2008, and go to a bar or a club or one of the myriad disco-basements on the Lower East Side that does not quite know which of those it is. Dance awkwardly in a room full of other glittered-up nerds, and wait for something to happen, buoyed on the feeling that this is the big swollen heart of life, that this is New York like the movies." + --incoming-requests-per-step: 32 + --inference-repeat-n: 3 + --no-record-throughput: true + --mamba-inference-conv-states-dtype: fp32 + --mamba-inference-ssm-states-dtype: fp32 +METRICS: + - "generated_tokens" + - "logprobs" diff --git a/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml b/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml index 495a4d130b0..aa78cdb8316 100644 --- a/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml +++ b/tests/test_utils/recipes/h100/mamba-dynamic-inference.yaml @@ -65,3 +65,8 @@ products: - environment: [dev] scope: [mr, mr-github] platforms: [dgx_h100] + - test_case: [hybrid_dynamic_inference_tp1_pp1_dp8_583m_flashinfer] + products: + - environment: [dev] + scope: [mr, mr-github] + platforms: [dgx_h100] diff --git a/tests/unit_tests/inference/engines/test_dynamic_engine.py b/tests/unit_tests/inference/engines/test_dynamic_engine.py index 24efaea9e1d..7bcf21882c1 100644 --- a/tests/unit_tests/inference/engines/test_dynamic_engine.py +++ b/tests/unit_tests/inference/engines/test_dynamic_engine.py @@ -146,6 +146,7 @@ class DynamicEngineTestConfig: track_generated_token_events: bool = False num_speculative_tokens: int = 0 position_embedding_type: str = "learned_absolute" + sampling_backend: str = 'torch' def __post_init__(self): @@ -274,6 +275,7 @@ def _build_inference_context( unified_memory_level=0, # unit tests currently broken with UVM track_generated_token_events=test_config.track_generated_token_events, num_speculative_tokens=test_config.num_speculative_tokens, + sampling_backend=test_config.sampling_backend, ), ) @@ -2261,9 +2263,15 @@ def set_epoch(epoch): not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) @torch.inference_mode() - def test_speculative_decoding_with_early_termination(self): + @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"]) + @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False]) + def test_speculative_decoding_with_early_termination( + self, materialize_only_last_token_logits, sampling_backend + ): """Test that speculative decoding handles premature request termination safely (e.g. hitting max_sequence_length mid-speculative-batch).""" + if sampling_backend == "flashinfer": + pytest.importorskip("flashinfer") # Set max_sequence_length tight so it terminates during a speculative step test_config = DynamicEngineTestConfig( @@ -2274,7 +2282,8 @@ def test_speculative_decoding_with_early_termination(self): max_sequence_length=7, # Will force termination after 3 tokens model_provider="gpt", num_speculative_tokens=3, - materialize_only_last_token_logits=False, + materialize_only_last_token_logits=materialize_only_last_token_logits, + sampling_backend=sampling_backend, ) env = self._build_test_env(test_config) @@ -2299,6 +2308,8 @@ def mock_mtp_forward(*args, **kwargs): unwrapped_model._decoder_hidden_states_cache = torch.zeros( tokens.size(1), 1, hidden_size, device=tokens.device, dtype=torch.bfloat16 ) + if test_config.materialize_only_last_token_logits: + base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0) return base_logits def mock_compute_mtp_single_step( @@ -2335,12 +2346,18 @@ def mock_compute_mtp_single_step( @pytest.mark.internal @torch.inference_mode() - def test_speculative_block_boundary_crossing(self): + @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"]) + @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False]) + def test_speculative_block_boundary_crossing( + self, materialize_only_last_token_logits, sampling_backend + ): """Test to verify KV cache block boundary crossing logic. When a request fills exactly one block and speculative decoding generates multiple tokens, the first new token shouldn't incorrectly overwrite the old block. """ + if sampling_backend == "flashinfer": + pytest.importorskip("flashinfer") test_config = DynamicEngineTestConfig( num_requests=1, min_prompt_length=256, @@ -2350,8 +2367,9 @@ def test_speculative_block_boundary_crossing(self): context_block_size_tokens=256, # Exactly matches prompt length context_max_requests=16, model_provider="gpt", - materialize_only_last_token_logits=False, + materialize_only_last_token_logits=materialize_only_last_token_logits, use_fixed_output_lengths=True, + sampling_backend=sampling_backend, ) env = self._build_test_env(test_config) @@ -2391,9 +2409,13 @@ def test_speculative_block_boundary_crossing(self): not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) @torch.inference_mode() - def test_speculative_stop_word_hit(self): + @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"]) + @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False]) + def test_speculative_stop_word_hit(self, materialize_only_last_token_logits, sampling_backend): """Test that if an accepted speculative token completes a stop word, the request correctly triggers the stop logic without crashing.""" + if sampling_backend == "flashinfer": + pytest.importorskip("flashinfer") test_config = DynamicEngineTestConfig( num_requests=0, # We will manually add our request cleanly @@ -2401,8 +2423,9 @@ def test_speculative_stop_word_hit(self): max_prompt_length=4, num_tokens_to_generate=10, num_speculative_tokens=2, - materialize_only_last_token_logits=False, + materialize_only_last_token_logits=materialize_only_last_token_logits, model_provider="gpt", + sampling_backend=sampling_backend, ) env = self._build_test_env(test_config) @@ -2424,6 +2447,8 @@ def mock_deterministic_forward(*args, **kwargs): unwrapped_model._decoder_hidden_states_cache = torch.zeros( s, 1, hidden_size, device=tokens.device, dtype=torch.bfloat16 ) + if test_config.materialize_only_last_token_logits: + base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0) return base_logits def mock_compute_mtp_single_step( @@ -2477,9 +2502,15 @@ def mock_compute_mtp_single_step( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) @torch.inference_mode() - def test_speculative_long_stop_word_hit(self): + @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"]) + @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False]) + def test_speculative_long_stop_word_hit( + self, materialize_only_last_token_logits, sampling_backend + ): """Test that if an accepted speculative token completes a long stop word (length > num_speculative_tokens), it is correctly detected.""" + if sampling_backend == "flashinfer": + pytest.importorskip("flashinfer") test_config = DynamicEngineTestConfig( num_requests=0, @@ -2487,8 +2518,9 @@ def test_speculative_long_stop_word_hit(self): max_prompt_length=4, num_tokens_to_generate=10, num_speculative_tokens=2, - materialize_only_last_token_logits=False, + materialize_only_last_token_logits=materialize_only_last_token_logits, model_provider="gpt", + sampling_backend=sampling_backend, ) env = self._build_test_env(test_config) @@ -2510,6 +2542,8 @@ def mock_deterministic_forward(*args, **kwargs): unwrapped_model._decoder_hidden_states_cache = torch.zeros( s, 1, hidden_size, device=tokens.device, dtype=torch.bfloat16 ) + if test_config.materialize_only_last_token_logits: + base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0) return base_logits def mock_compute_mtp_single_step( @@ -2559,7 +2593,11 @@ def mock_compute_mtp_single_step( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) @torch.inference_mode() - def test_speculative_stop_word_truncates_trailing_tokens(self): + @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"]) + @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False]) + def test_speculative_stop_word_truncates_trailing_tokens( + self, materialize_only_last_token_logits, sampling_backend + ): """Test that when a stop word lands in the middle of speculative tokens, the extra tokens generated after the stop word are removed. @@ -2567,6 +2605,8 @@ def test_speculative_stop_word_truncates_trailing_tokens(self): (1 base + 2 speculative). If the stop word is [6] and the engine generates [5, 6, 7] in one step, token 7 must be truncated so the output ends with the stop word [6].""" + if sampling_backend == "flashinfer": + pytest.importorskip("flashinfer") test_config = DynamicEngineTestConfig( num_requests=0, @@ -2574,8 +2614,9 @@ def test_speculative_stop_word_truncates_trailing_tokens(self): max_prompt_length=4, num_tokens_to_generate=10, num_speculative_tokens=2, - materialize_only_last_token_logits=False, + materialize_only_last_token_logits=materialize_only_last_token_logits, model_provider="gpt", + sampling_backend=sampling_backend, ) env = self._build_test_env(test_config) @@ -2597,6 +2638,8 @@ def mock_deterministic_forward(*args, **kwargs): unwrapped_model._decoder_hidden_states_cache = torch.zeros( s, 1, hidden_size, device=tokens.device, dtype=torch.bfloat16 ) + if test_config.materialize_only_last_token_logits: + base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0) return base_logits def mock_compute_mtp_single_step( @@ -2677,9 +2720,16 @@ def mock_compute_mtp_single_step( "non_divisible_boundary", ], ) + @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"]) + @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False]) @torch.inference_mode() def test_speculative_tokens_exceed_max_sequence_length( - self, prompt_length, num_tokens_to_generate, num_speculative_tokens + self, + prompt_length, + num_tokens_to_generate, + num_speculative_tokens, + materialize_only_last_token_logits, + sampling_backend, ): """Test that speculative decoding correctly trims output when speculative tokens would push the sequence beyond max_sequence_length. @@ -2689,6 +2739,8 @@ def test_speculative_tokens_exceed_max_sequence_length( speculative tokens are accepted and the boundary trimming logic is actually exercised. """ + if sampling_backend == "flashinfer": + pytest.importorskip("flashinfer") max_sequence_length = prompt_length + num_tokens_to_generate test_config = DynamicEngineTestConfig( @@ -2698,11 +2750,12 @@ def test_speculative_tokens_exceed_max_sequence_length( num_tokens_to_generate=num_tokens_to_generate, max_sequence_length=max_sequence_length, num_speculative_tokens=num_speculative_tokens, - materialize_only_last_token_logits=False, + materialize_only_last_token_logits=materialize_only_last_token_logits, model_provider="gpt", # Disable positional embeddings so speculative position IDs # beyond max_sequence_length don't cause out-of-bounds lookups. position_embedding_type="none", + sampling_backend=sampling_backend, ) env = self._build_test_env(test_config) @@ -2805,8 +2858,12 @@ def test_detokenize_stop_sequence_flag(self, detokenize_stop_sequence): @pytest.mark.parametrize( "acceptance_mode", ["all_rejected", "all_accepted"], ids=["all_rejected", "all_accepted"] ) + @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"]) + @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False]) @torch.inference_mode() - def test_speculative_sequence_length_double_counting(self, acceptance_mode): + def test_speculative_sequence_length_double_counting( + self, acceptance_mode, materialize_only_last_token_logits, sampling_backend + ): """Test to verify active_sequence_lengths is not double-counted. If active sequence length is double-counted during speculative decoding, @@ -2818,6 +2875,8 @@ def test_speculative_sequence_length_double_counting(self, acceptance_mode): a faulty formula that adds accepted_tokens on top of the KV length will over-count by 2 per step, finishing the request after only 4 of 6 tokens. """ + if sampling_backend == "flashinfer": + pytest.importorskip("flashinfer") test_config = DynamicEngineTestConfig( num_requests=0, min_prompt_length=4, @@ -2827,10 +2886,11 @@ def test_speculative_sequence_length_double_counting(self, acceptance_mode): context_max_requests=16, num_speculative_tokens=2, model_provider="gpt", - materialize_only_last_token_logits=False, + materialize_only_last_token_logits=materialize_only_last_token_logits, use_fixed_output_lengths=False, context_max_tokens=512, position_embedding_type="none", + sampling_backend=sampling_backend, ) env = self._build_test_env(test_config) @@ -2854,6 +2914,8 @@ def mock_mtp_forward(*args, **kwargs): model._decoder_hidden_states_cache = torch.zeros( s, 1, hidden_size, device=tokens.device, dtype=torch.bfloat16 ) + if test_config.materialize_only_last_token_logits: + base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0) return base_logits def mock_compute_mtp(*args_mtp, **kwargs_mtp): @@ -2925,12 +2987,18 @@ def deterministic_mtp( not is_fa_min_version("2.7.3"), reason="need latest flash attn for dynamic batching" ) @torch.inference_mode() - def test_speculative_decoding_with_eviction_and_swapping(self): + @pytest.mark.parametrize("sampling_backend", ["torch", "flashinfer"]) + @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False]) + def test_speculative_decoding_with_eviction_and_swapping( + self, materialize_only_last_token_logits, sampling_backend + ): """Test that speculative decoding works correctly when requests are paused and evicted. This exercises the `_swap_book_keeping_tensors` logic with the 2D `new_speculative_tokens` tensor, ensuring no dimensional mismatch or index errors occur during tensor swapping. """ + if sampling_backend == "flashinfer": + pytest.importorskip("flashinfer") # Very constrained memory environment to force pausing and eviction test_config = DynamicEngineTestConfig( num_requests=3, @@ -2942,8 +3010,9 @@ def test_speculative_decoding_with_eviction_and_swapping(self): context_buffer_size_gb=0.00064, # 640 KB context_paused_buffer_size_gb=0.0, # 0 paused buffer forces immediate eviction model_provider="gpt", - materialize_only_last_token_logits=False, + materialize_only_last_token_logits=materialize_only_last_token_logits, use_fixed_output_lengths=True, + sampling_backend=sampling_backend, ) env = self._build_test_env(test_config) @@ -2966,6 +3035,8 @@ def mock_safe_forward(*args, **kwargs): unwrapped_model._decoder_hidden_states_cache = torch.zeros( s, 1, hidden_size, device=tokens.device, dtype=torch.bfloat16 ) + if test_config.materialize_only_last_token_logits: + base_logits = env.engine.context.last_token_logits(base_logits).unsqueeze(0) return base_logits def mock_compute_mtp_single_step( @@ -4601,6 +4672,7 @@ def _build_engine(self, model, enable_chunked_prefill, num_cuda_graphs, context_ enable_chunked_prefill=enable_chunked_prefill, max_tokens=context_max_tokens, max_requests=128, + sampling_backend='torch', ) if mamba_config is not None: inference_config_kwargs.update(mamba_inference_state_config=mamba_config) diff --git a/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py b/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py index d3f1930561b..92890b22b53 100644 --- a/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py +++ b/tests/unit_tests/inference/engines/test_hybrid_prefix_caching_e2e.py @@ -200,6 +200,7 @@ def _build_engine( enable_prefix_caching=enable_prefix_caching, unified_memory_level=0, num_cuda_graphs=num_cuda_graphs, + sampling_backend='torch', ) if enable_prefix_caching: inference_config_kwargs.update( diff --git a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py index f4e76f6ab7c..8fd1f4a1154 100644 --- a/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py +++ b/tests/unit_tests/inference/test_mtp_cuda_graph_inference.py @@ -153,6 +153,7 @@ def _build_engine( block_size_tokens=256, max_requests=max_requests, num_cuda_graphs=-1, + sampling_backend='torch', ), ) wrapped = GPTInferenceWrapper(model, context) @@ -420,10 +421,10 @@ def test_cuda_graph_sp_padding_end_to_end(self, mtp_use_repeated_layer): ctrl._mtp_resolved_padded_count = padded_count context._using_cuda_graph_this_step = True - ctrl._torch_sampling_buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)] - ctrl._torch_sampling_bucket_index_tensors = [ - torch.arange(active_request_count, device='cuda', dtype=torch.long) - ] + # Greedy sampling for all active requests. + context.active_request_metadata["temperature"][:active_request_count] = 1.0 + context.active_request_metadata["top_k"][:active_request_count] = 1 + context.active_request_metadata["top_p"][:active_request_count] = 0.0 ctrl._compute_serial_mtp_and_sample() @@ -521,10 +522,10 @@ def _run_mtp(use_cuda_graph): unwrapped._decoder_hidden_states_cache = local_hidden ctrl._last_accepted_seq_indices = torch.arange(active_request_count, device='cuda') - ctrl._torch_sampling_buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)] - ctrl._torch_sampling_bucket_index_tensors = [ - torch.arange(active_request_count, device='cuda', dtype=torch.long) - ] + # Greedy sampling for all active requests. + context.active_request_metadata["temperature"][:active_request_count] = 1.0 + context.active_request_metadata["top_k"][:active_request_count] = 1 + context.active_request_metadata["top_p"][:active_request_count] = 0.0 ctrl._compute_serial_mtp_and_sample() # CUDA graph replay is asynchronous, and this test reuses the controller's @@ -900,6 +901,7 @@ def _build_context( num_cuda_graphs=num_cuda_graphs, use_cuda_graphs_for_non_decode_steps=use_cuda_graphs_for_non_decode_steps, max_requests=max_requests, + sampling_backend='torch', ), ) diff --git a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py index 2c1a9902a1f..5c96836e9ba 100644 --- a/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py +++ b/tests/unit_tests/inference/text_generation_controllers/test_text_generation_controller.py @@ -67,6 +67,7 @@ def setup_model( expert_model_parallel_size: int = 1, num_moe_experts: int = None, hybrid_layer_pattern: str = None, + sampling_backend: str = 'torch', cuda_graph_impl: str = 'none', ): if use_training_random_init: @@ -164,6 +165,7 @@ def setup_model( enable_prefix_caching=enable_prefix_caching, max_requests=max_requests, mamba_inference_state_config=mamba_inference_state_config, + sampling_backend=sampling_backend, ), ) @@ -280,28 +282,34 @@ def detokenize(self, inp, skip_special_tokens=False): sampled_logits >= expected_min_value ), f"The sampled logits should all be greater than {expected_min_value} but its {sampled_logits}" - @pytest.mark.parametrize("backend", ["torch"]) + @pytest.mark.parametrize("backend", ["torch", "flashinfer"]) @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False]) def test_sample_from_dynamic_logits( self, backend: str, materialize_only_last_token_logits: bool ): - batch_size = 12 + if backend == "flashinfer": + pytest.importorskip("flashinfer") + batch_size = 15 self.setup_model( torch.float32, batch_size=batch_size, static=False, materialize_only_last_token_logits=materialize_only_last_token_logits, + sampling_backend=backend, ) self.mock_tokenizer.eod = self.vocab_size context = self.text_generation_controller.inference_wrapped_model.inference_context # Prepare sampling params in human-readable format, to aid with test maintenance. + # The temperature=0 / top_k=1 bucket exercises the greedy path: torch short-circuits + # to argmax, flashinfer relies on its temperature clamp to avoid divide-by-zero. sampling_test_cases: List[Tuple[SamplingParams, List[int]]] = [ (SamplingParams(temperature=0.1, top_p=0.01), [9, 6, 10]), (SamplingParams(temperature=5.0, top_k=15), [0, 3, 2]), (SamplingParams(top_p=0.8), [4, 1, 7]), (SamplingParams(temperature=10.0, top_k=5), [11, 5, 8]), + (SamplingParams(temperature=0.0, top_k=1), [12, 13, 14]), ] # For non-torch backends, test simultaneous top_k and top_p sampling. if backend != "torch": @@ -320,7 +328,6 @@ def test_sample_from_dynamic_logits( context.active_request_metadata["temperature"][:batch_size].copy_(temp_values) context.active_request_metadata["top_k"][:batch_size].copy_(top_k_values) context.active_request_metadata["top_p"][:batch_size].copy_(top_p_values) - self.text_generation_controller._sampling_backend = backend context.padded_active_token_count = batch_size context.request_query_lengths = torch.ones(batch_size, dtype=torch.int32, device='cuda') @@ -329,9 +336,6 @@ def test_sample_from_dynamic_logits( context.num_prefill_requests = 0 context.pad_active_slices() - # Bookkeeping. - self.text_generation_controller._dynamic_step_sample_bookkeeping() - # Sampling. logits = torch.arange(0, self.vocab_size).repeat(batch_size, 1).unsqueeze(0).float().cuda() self.text_generation_controller._all_logits_cuda = logits @@ -355,13 +359,13 @@ def test_sample_from_dynamic_logits( sampled_l.masked_fill_(top_k_mask, 0.0) top_p_mask = sampled_l.cumsum(dim=-1) > top_p_values.unsqueeze(1) + # When `top_p` is enabled, but the cumulative probs don't actually filter anything, + # our constraint reduces to top_k alone. + start_idx = torch.clamp(self.vocab_size - top_k_values, min=0).long() first_excluded = torch.where( - top_p_mask.any(dim=-1), - top_p_mask.float().argmax(dim=-1), - torch.full((batch_size,), self.vocab_size, device=top_p_mask.device), + top_p_mask.any(dim=-1), top_p_mask.float().argmax(dim=-1), start_idx + 1 ) last_included = torch.clamp(first_excluded - 1, min=0) - start_idx = torch.clamp(self.vocab_size - top_k_values, min=0).long() last_included = torch.max(last_included, start_idx) expected_min_values = l.gather(1, last_included.unsqueeze(1)).squeeze(1) assert torch.all( @@ -958,10 +962,22 @@ def test_dynamic_top_n_logprobs_calculation( ), f"Request {req_idx}, token {token_idx}: expected {top_n} indices" @pytest.mark.internal - def test_speculative_verify_tokens(self): + @pytest.mark.parametrize("backend", ["torch", "flashinfer"]) + @pytest.mark.parametrize("materialize_only_last_token_logits", [True, False]) + def test_speculative_verify_tokens( + self, backend: str, materialize_only_last_token_logits: bool + ): """Test consecutive token acceptance logic for speculative decoding.""" + if backend == "flashinfer": + pytest.importorskip("flashinfer") self.setup_model( - torch.float32, static=False, num_speculative_tokens=2, max_requests=2, mtp_num_layers=2 + torch.float32, + static=False, + num_speculative_tokens=2, + max_requests=2, + mtp_num_layers=2, + sampling_backend=backend, + materialize_only_last_token_logits=materialize_only_last_token_logits, ) # Enable speculative decoding @@ -998,12 +1014,8 @@ def mock_sampling_func(logits, *args, **kwargs): # The verification logic only uses base tokens, so we can return zeros here. return torch.zeros((12,), dtype=torch.long, device='cuda') - # Override sampling to return our predictable mock outputs - self.text_generation_controller._torch_sampling_buckets = [([0, 1], 1.0, 1, 0.0)] - self.text_generation_controller._torch_sampling_bucket_index_tensors = [ - torch.tensor([0, 1], device='cuda', dtype=torch.long) - ] - self.text_generation_controller._torch_sampling_func = mock.MagicMock( + # Override sampling to return our predictable mock outputs. + self.text_generation_controller._sampling.sample_kernel = mock.MagicMock( side_effect=mock_sampling_func ) @@ -1238,15 +1250,11 @@ def test_speculative_multinomial_sampling(self): # Base logits shape: [1, 8, vocab_size] logits = torch.randn(1, 8, self.vocab_size, device='cuda') - # Set up a bucket that forces multinomial sampling (top_p = 0.9, top_k = 0) - # _torch_sampling_buckets format: (indices, temp, top_k, top_p) - self.text_generation_controller._torch_sampling_buckets = [([0, 1], 1.0, 0, 0.9)] - self.text_generation_controller._torch_sampling_bucket_index_tensors = [ - torch.tensor([0, 1], device='cuda', dtype=torch.long) - ] - - # Since we are actually testing the internal math of `_torch_sampling_func` handling the shapes, - # we DO NOT mock `_torch_sampling_func` here. We want it to run natively to prove it doesn't crash. + # Drive sampling onto the multinomial path (top_p > 0, top_k == 0) via metadata. + # We do NOT mock the sampling kernel: we want it to run natively to prove it doesn't crash. + ctx.active_request_metadata["temperature"][:2] = 1.0 + ctx.active_request_metadata["top_k"][:2] = 0 + ctx.active_request_metadata["top_p"][:2] = 0.9 self.text_generation_controller._all_logits_cuda = logits try: @@ -1491,10 +1499,9 @@ def test_mtp_sp_padding_real_ranks(self, active_request_count): ctrl._last_accepted_seq_indices = torch.arange(active_request_count, device='cuda') # Greedy sampling: top_k=1 selects the argmax token deterministically. - ctrl._torch_sampling_buckets = [(list(range(active_request_count)), 1.0, 1, 0.0)] - ctrl._torch_sampling_bucket_index_tensors = [ - torch.arange(active_request_count, device='cuda', dtype=torch.long) - ] + ctx.active_request_metadata["temperature"][:active_request_count] = 1.0 + ctx.active_request_metadata["top_k"][:active_request_count] = 1 + ctx.active_request_metadata["top_p"][:active_request_count] = 0.0 # Run the MTP forward pass ctrl._compute_serial_mtp_and_sample()