diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index 20d6a2506851..02068ee7d0a2 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -322,6 +322,56 @@ llm.generate(["Hello, my name is", "Hello, my name is"], sampling_params) ``` +### Compact variable-width beam execution (prototype) + +`compact_beam_widths` is an opt-in PyTorch engine setting for variable-width beam +search. It lists the **input beam widths** (at least 2) for which the model forward should +execute only active rows. The default, `[]`, preserves capacity-width execution. Width-one steps also +retain capacity-width execution, preserving the beam attention kernel family. +The KV cache allocation and sampler layout still use `max_beam_width`; this +option does not reduce the reserved beam capacity. + +```python +llm = LLM( + model="Qwen/Qwen3-0.6B", + max_batch_size=2, + max_seq_len=1024, + max_num_tokens=1024, + max_beam_width=500, + compact_beam_widths=[20, 300], +) +params = SamplingParams( + use_beam_search=True, + best_of=500, + n=5, + max_tokens=6, + beam_width_array=[20, 20, 300, 300, 500, 500], +) +llm.generate(["The capital of France is"], params) +``` + +The prototype supports single-GPU dense Llama, Qwen2 and Qwen3 decoder models +with TRTLLM attention. It requires one token per beam and the same input width +for all real requests in a decode batch. Context/mixed batches and widths absent +from the list use capacity-width execution. Existing beam-search admission and +mixed-width restrictions still apply. Speculative decoding, encoder-decoder and +multimodal models, sparse/hybrid attention, torch compile and in-graph sampling +use capacity-width execution. Configuring this option with LoRA is rejected. + +With CUDA graphs enabled, each configured compact width adds a graph per +captured batch size. Input gathering and logits padding are captured; graph +metadata is refreshed in place before replay. This adds graph memory and warmup +time. Eager execution is also supported. The sampler receives capacity-stride +logits, with inactive rows filled with negative infinity. + +Changing the forward shape can select different attention/GEMM kernels and +floating-point reduction orders. **Bitwise equality, identical beam rankings, +and identical top-1 outputs are not guaranteed**, including with deterministic +sampling. Evaluate output quality for the intended model and workload before +enabling the option. Small beam capacities can be slower because metadata +packing and full-vocabulary logits padding add work; benchmark the complete +workload rather than assuming a speedup from the row-count reduction. + ### Over the OpenAI-compatible API `length_penalty` and `early_stopping` now default to `null` in the HTTP schema, deferring to diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 46df832ac489..1bf19212d005 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/tensorrt_llm/_torch/models/modeling_llama.py @@ -14,7 +14,7 @@ import copy import os -from typing import Any, Dict, List, Literal, Optional, Tuple, Union +from typing import Any, ClassVar, Dict, List, Literal, Optional, Tuple, Union import torch from PIL.Image import Image @@ -1148,6 +1148,8 @@ def forward( @register_auto_model("LlamaForCausalLM") class LlamaForCausalLM(SpecDecOneEngineForCausalLM[LlamaModel, LlamaConfig]): + supports_compact_beam: ClassVar[bool] = True + @classmethod def get_preferred_kv_cache_manager_version( cls, diff --git a/tensorrt_llm/_torch/models/modeling_qwen.py b/tensorrt_llm/_torch/models/modeling_qwen.py index 68b5e610ea9a..d0d4ab4bb7b2 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen.py +++ b/tensorrt_llm/_torch/models/modeling_qwen.py @@ -1,4 +1,7 @@ -from typing import Optional, Tuple +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +from typing import ClassVar, Optional, Tuple import torch from torch import nn @@ -248,6 +251,8 @@ def forward( @register_auto_model("Qwen2ForCausalLM") class Qwen2ForCausalLM(DecoderModelForCausalLM[QwenModel, Qwen2Config]): + supports_compact_beam: ClassVar[bool] = True + def __init__( self, model_config: ModelConfig[Qwen2Config], diff --git a/tensorrt_llm/_torch/models/modeling_qwen3.py b/tensorrt_llm/_torch/models/modeling_qwen3.py index 50474525a6b8..bc2428d0c182 100644 --- a/tensorrt_llm/_torch/models/modeling_qwen3.py +++ b/tensorrt_llm/_torch/models/modeling_qwen3.py @@ -1,5 +1,8 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + import copy -from typing import Optional +from typing import ClassVar, Optional import torch from torch import nn @@ -274,6 +277,8 @@ def forward( @register_auto_model("Qwen3ForCausalLM") class Qwen3ForCausalLM(SpecDecOneEngineForCausalLM[Qwen3Model, Qwen3Config]): + supports_compact_beam: ClassVar[bool] = True + def __init__( self, model_config: ModelConfig[Qwen3Config], diff --git a/tensorrt_llm/_torch/pyexecutor/beam_compaction.py b/tensorrt_llm/_torch/pyexecutor/beam_compaction.py new file mode 100644 index 000000000000..d76b1cc67563 --- /dev/null +++ b/tensorrt_llm/_torch/pyexecutor/beam_compaction.py @@ -0,0 +1,205 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Execute active beam rows while retaining the executor's capacity-stride layout.""" + +import copy +from typing import Any + +import torch +import torch.nn.functional as F + +from tensorrt_llm._utils import prefer_pinned + +from ..attention.backends.trtllm import TrtllmAttentionMetadata +from .scheduler import ScheduledRequests + + +class CompactBeamView: + """Stable decode metadata buffers for one (batch size, execution width) shape. + + Rows are request-major. Selecting the first ``width`` rows of each request + preserves physical KV page IDs; neither the KV allocator nor the sampler + changes its capacity-stride layout. + """ + + def __init__(self, source: TrtllmAttentionMetadata, width: int, capacity: int): + if not 1 < width < capacity: + raise ValueError("Compact execution width must be at least 2 and smaller than capacity") + self.source = source + self.capacity = capacity + self.width = width + self.batch_size = source.num_seqs // capacity + self.row_indices_cpu = ( + torch.arange(self.batch_size, device="cpu")[:, None] * capacity + + torch.arange(width, device="cpu")[None, :] + ).flatten() + self.row_indices_list = self.row_indices_cpu.tolist() + self.row_indices = self.row_indices_cpu.to(source.kv_cache_block_offsets.device) + self.metadata = metadata = copy.copy(source) + count = self.batch_size * width + metadata.beam_width = width + metadata.max_num_requests = self.batch_size + metadata.max_num_sequences = count + metadata._seq_lens = torch.ones(count, dtype=torch.int32, device="cpu") + metadata._seq_lens_cuda = metadata._seq_lens.to(self.row_indices.device) + metadata._seq_lens_kv = None + metadata._seq_lens_kv_cuda = None + metadata._num_generations = count + metadata._num_tokens = count + metadata._num_ctx_tokens = 0 + metadata.padded_num_tokens = None + metadata.all_rank_num_tokens = [count] + metadata._saved_tensors = {} + metadata.fmha_plan_caches = {} + metadata.draft_replay_swapped_attrs = {} + metadata.kv_cache_params = copy.copy(source.kv_cache_params) + for name in self._row_fields: + original = getattr(source, name) + setattr( + metadata, + name, + torch.empty( + count, + dtype=original.dtype, + device=original.device, + pin_memory=not original.is_cuda and prefer_pinned(), + ), + ) + metadata.host_total_kv_lens = torch.empty_like(source.host_total_kv_lens) + offsets = source.kv_cache_block_offsets + metadata.kv_cache_block_offsets = offsets.new_empty( + offsets.shape[0], count, *offsets.shape[2:] + ) + metadata.cache_indirection = source.cache_indirection.new_empty( + self.batch_size, width, source.cache_indirection.shape[-1] + ) + metadata._bind_runtime_views( + kv_lens_cuda=metadata.kv_lens_cuda, + kv_lens=metadata.kv_lens, + prompt_lens_cuda=metadata.prompt_lens_cuda, + prompt_lens_cpu=metadata.prompt_lens_cpu, + host_request_types=metadata.host_request_types, + ) + + _row_fields = ( + "kv_lens", + "kv_lens_cuda", + "prompt_lens_cpu", + "prompt_lens_cuda", + "host_request_types", + ) + + def update(self) -> None: + """Refresh graph-visible buffers in place, outside graph capture/replay.""" + source, metadata = self.source, self.metadata + for name in self._row_fields: + original = getattr(source, name) + indices = self.row_indices if original.is_cuda else self.row_indices_cpu + torch.index_select(original, 0, indices, out=getattr(metadata, name)) + torch.index_select( + source.kv_cache_block_offsets, 1, self.row_indices, out=metadata.kv_cache_block_offsets + ) + metadata.cache_indirection.copy_(source.cache_indirection[: self.batch_size, : self.width]) + metadata.host_total_kv_lens[0] = 0 + metadata.host_total_kv_lens[1] = metadata.kv_lens.sum() + metadata.request_ids = source.request_ids + metadata.prompt_lens = [source.prompt_lens[i] for i in self.row_indices_list] + metadata.kv_cache_params.num_cached_tokens_per_seq = [ + source.kv_cache_params.num_cached_tokens_per_seq[i] for i in self.row_indices_list + ] + metadata.on_update_kv_lens() + + def pack_inputs(self, inputs: dict[str, Any]) -> dict[str, Any]: + """Gather inside the graph so replay can consume capacity-stride inputs.""" + return { + **inputs, + "attn_metadata": self.metadata, + "input_ids": inputs["input_ids"].index_select(0, self.row_indices), + "position_ids": inputs["position_ids"].index_select(-1, self.row_indices), + } + + def restore_logits(self, logits: torch.Tensor) -> torch.Tensor: + """Inactive rows remain impossible candidates for the unchanged sampler.""" + return F.pad( + logits.reshape(self.batch_size, self.width, -1), + (0, 0, 0, self.capacity - self.width), + value=-float("inf"), + ).reshape(self.batch_size * self.capacity, -1) + + +class BeamCompaction: + """Select configured widths and own metadata views for eager/CUDA graph use.""" + + def __init__(self, capacity: int, widths: tuple[int, ...]): + self.capacity = capacity + self.widths = widths + self._views: dict[tuple[int, int, int], CompactBeamView] = {} + + def clear(self) -> None: + self._views.clear() + + def execution_width(self, batch: ScheduledRequests, has_previous_tokens: bool) -> int: + if not self.widths or batch.context_requests: + return 0 + widths = { + request.get_beam_width_by_iter(for_next_iteration=has_previous_tokens) + for request in batch.generation_requests + if not request.is_dummy + } + if len(widths) != 1: + return 0 + width = widths.pop() + return width if width in self.widths else 0 + + def supports_inputs(self, inputs: dict[str, Any]) -> bool: + if ( + not self.widths + or inputs.get("multimodal_params") + or any( + inputs.get(key) is not None + for key in ( + "inputs_embeds", + "spec_metadata", + "lora_params", + "encoder_hidden_states", + "context_requests", + "mrope_delta_read_seq_slots", + "mrope_delta_write_seq_slots", + ) + ) + ): + return False + metadata = inputs["attn_metadata"] + return ( + type(metadata) is TrtllmAttentionMetadata + and not metadata.num_contexts + and not metadata.has_cross_sub_metadata + and not metadata.mamba_metadata + and metadata.beam_width == self.capacity + and metadata.num_seqs > 0 + and metadata.num_seqs % self.capacity == 0 + and metadata.num_tokens == metadata.num_seqs + and bool(torch.all(metadata.seq_lens == 1)) + and not metadata.kv_cache_params.num_extra_kv_tokens + and metadata.kv_cache_block_offsets is not None + and metadata.cache_indirection is not None + ) + + def prepare_inputs(self, inputs: dict[str, Any], width: int) -> None: + inputs.pop("beam_compaction", None) + if width not in self.widths or not self.supports_inputs(inputs): + return + source = inputs["attn_metadata"] + key = (id(source), source.num_seqs, width) + if key not in self._views: + if not source.is_cuda_graph: + # Eager views only need the latest batch shape for each width. + # CUDA graph views retain stable buffers for replay. + for cached_key in tuple(self._views): + if cached_key[0] == id(source) and cached_key[2] == width: + del self._views[cached_key] + self._views[key] = CompactBeamView(source, width, self.capacity) + view = self._views[key] + view.update() + inputs["beam_compaction"] = view diff --git a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py index 893ff3953224..acb523634ff0 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -30,6 +30,7 @@ from ..speculative.spec_sampler_base import SampleStateTensorsSpec from ..speculative.utils import get_draft_kv_cache_manager from ..utils import make_weak_ref, piecewise_cuda_graph +from .beam_compaction import BeamCompaction from .llm_request import LlmRequest, get_draft_token_length from .resource_manager import (BaseResourceManager, ResourceManager, ResourceManagerType) @@ -78,6 +79,8 @@ class KeyType(NamedTuple): num_encoder_tokens: int = 0 peft_cache_data_type: Optional[torch.dtype] = None use_lora_graph: bool = False + # Zero selects the original capacity-width forward. + beam_execution_width: int = 0 def _save_spec_decode_capture_state( @@ -152,6 +155,7 @@ class CUDAGraphRunnerConfig: sparse_attention_config: Optional[BaseSparseAttentionConfig] = None enable_encoder_decoder_mixed_cuda_graph: bool = False enable_in_graph_sampling: bool = False + compact_beam_widths: Tuple[int, ...] = () static_input_ids: Optional[torch.Tensor] = None """The engine's own int32 input-id buffer. When given (with ``static_position_ids``), the graphs' static inputs are views of the @@ -173,6 +177,8 @@ class CUDAGraphRunner: def __init__(self, config: CUDAGraphRunnerConfig): self.config = config + self.beam_compaction = BeamCompaction(config.max_beam_width, + config.compact_beam_widths) # High-level configuration from the config object self.enabled = config.use_cuda_graph @@ -442,16 +448,20 @@ def get_graph_key( num_encoder_tokens = sum( int(request.encoder_output_len) for request in context_requests if not request.py_skip_cross_kv_projection) - return KeyType(batch_size=batch_size, - draft_len=draft_len, - short_seq_len_mode=short_seq_len_mode, - is_all_greedy_sample=is_all_greedy_sample, - sample_type=sample_type, - num_contexts=num_contexts, - context_query_len=context_query_len, - num_encoder_tokens=num_encoder_tokens, - peft_cache_data_type=peft_cache_data_type, - use_lora_graph=use_lora_graph) + return KeyType( + batch_size=batch_size, + draft_len=draft_len, + short_seq_len_mode=short_seq_len_mode, + is_all_greedy_sample=is_all_greedy_sample, + sample_type=sample_type, + num_contexts=num_contexts, + context_query_len=context_query_len, + num_encoder_tokens=num_encoder_tokens, + peft_cache_data_type=peft_cache_data_type, + use_lora_graph=use_lora_graph, + beam_execution_width=(self.beam_compaction.execution_width( + batch, new_tensors_device is not None) if not use_lora_graph + and not draft_len else 0)) def set_capture_sample_type(self, sample_type: Optional[SampleType]) -> None: @@ -712,6 +722,35 @@ def capture(self, initial_inputs: Dict[str, Any], enable_spec_decode: bool = False, postprocess_fn: Optional[Callable] = None) -> Any: + """Warm up and/or capture each supported execution width. + + DecoderCudaGraphs.warmup_and_capture sizes shared workspaces + for all graph shapes in a separate warmup-only pass before capture. + """ + key = KeyType(*key) + widths = (key.beam_execution_width, ) + if (not key.beam_execution_width and not key.num_contexts + and not key.draft_len and not enable_spec_decode + and not key.use_lora_graph + and self.beam_compaction.supports_inputs(initial_inputs)): + widths = (*self.config.compact_beam_widths, 0) + result = None + for width in widths: + inputs = initial_inputs.copy() + self.beam_compaction.prepare_inputs(inputs, width) + output = self._capture(key._replace(beam_execution_width=width), + forward_fn, inputs, enable_spec_decode, + postprocess_fn) + if width == key.beam_execution_width: + result = output + return result + + def _capture(self, + key: KeyType, + forward_fn: Callable, + initial_inputs: Dict[str, Any], + enable_spec_decode: bool = False, + postprocess_fn: Optional[Callable] = None) -> Any: """Warm up and/or capture the forward pass for a graph key.""" batch_size = key.batch_size # [CUDA graph spec decode padding] @@ -765,8 +804,12 @@ def capture(self, attn_metadata, enable_spec_decode) self.graph_metadata[key] = { - "attn_metadata": attn_metadata, - "spec_metadata": initial_inputs.get("spec_metadata", None), + "attn_metadata": + attn_metadata, + "spec_metadata": + initial_inputs.get("spec_metadata", None), + "compact_metadata": (initial_inputs["beam_compaction"].metadata if + "beam_compaction" in initial_inputs else None), } output = None @@ -825,6 +868,13 @@ def replay( position corrections applied inside the captured speculative preprocessing. """ stored_meta = self.graph_metadata[key] + if key.beam_execution_width: + compact = current_inputs.get("beam_compaction") + if compact is None or compact.metadata is not stored_meta[ + "compact_metadata"]: + raise RuntimeError( + "Compact CUDA graph replay requires refreshed compact metadata" + ) assert current_inputs["attn_metadata"] is stored_meta["attn_metadata"] if stored_meta["spec_metadata"] is not None: assert current_inputs.get( @@ -1003,6 +1053,10 @@ def _record_strict_buffer_snapshot(self, key: KeyType) -> None: stored_meta = self.graph_metadata[key] stored_meta["attn_metadata_ptrs"] = self._snapshot_graph_tensor_ptrs( stored_meta["attn_metadata"]) + if stored_meta.get("compact_metadata") is not None: + stored_meta[ + "compact_metadata_ptrs"] = self._snapshot_graph_tensor_ptrs( + stored_meta["compact_metadata"]) if stored_meta["spec_metadata"] is not None: stored_meta[ "spec_metadata_ptrs"] = self._snapshot_graph_tensor_ptrs( @@ -1016,6 +1070,10 @@ def _check_strict_buffer_stability(self, stored_meta: Dict[str, if "attn_metadata_ptrs" in stored_meta: self._assert_graph_tensor_ptrs_stable( stored_meta["attn_metadata"], stored_meta["attn_metadata_ptrs"]) + if "compact_metadata_ptrs" in stored_meta: + self._assert_graph_tensor_ptrs_stable( + stored_meta["compact_metadata"], + stored_meta["compact_metadata_ptrs"]) if "spec_metadata_ptrs" in stored_meta: self._assert_graph_tensor_ptrs_stable( stored_meta["spec_metadata"], stored_meta["spec_metadata_ptrs"]) @@ -1367,6 +1425,7 @@ def clear(self, *, release_nccl_window_owners: bool = True): self.graphs.clear() self.graph_outputs.clear() self.graph_metadata.clear() + self.beam_compaction.clear() self._pre_replay_position_ids = None self.padding_dummy_requests = {} del self.memory_pool diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/config.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/config.py index c8e7d4fe4008..29de1a5349b8 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/config.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/config.py @@ -3,7 +3,7 @@ """Resolved settings for scheduled decoder execution.""" -from dataclasses import dataclass +from dataclasses import dataclass, field import torch @@ -49,3 +49,4 @@ class DecoderRunnerConfig(RunnerConfig): mm_encoder_cache_enabled: bool enable_autotuner: bool cuda_graph_specialize_lora: bool + compact_beam_widths: tuple[int, ...] = field(default=(), kw_only=True) diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py index 274732b2d5da..4dcbbc908f5f 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/decoder/runner.py @@ -45,6 +45,7 @@ from tensorrt_llm._utils import nvtx_range from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor, InputProcessor from tensorrt_llm.llmapi.llm_args import BaseSparseAttentionConfig, PrefillCudaGraphBackend +from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping from ...lora import LoraParamBuilder, make_cuda_graph_lora_manager @@ -410,6 +411,33 @@ def _initialize_cuda_graph_runner(self) -> CUDAGraphRunner | None: register_hooks(cuda_graph_runner) return cuda_graph_runner + def _supported_compact_beam_widths(self) -> tuple[int, ...]: + widths = self._config.compact_beam_widths + if not widths: + return () + # Require an explicit declaration on each model class: subclasses may + # introduce additional token-indexed state or outputs. + supported = ( + type(self.model).__dict__.get("supports_compact_beam", False) + and self._config.attention_backend.Metadata is TrtllmAttentionMetadata + and self.mapping.world_size == 1 + and not self._config.is_encode_only + and not self._config.is_spec_decode + and not self._config.use_mrope + and not self._config.is_multimodal + and not self._config.without_logits + and not self._config.torch_compile_enabled + and not self._config.enable_in_graph_sampling + and self.sparse_attention_config is None + ) + if not supported: + logger.warning( + "compact_beam_widths is unsupported for this model configuration; " + "using capacity-width execution." + ) + return () + return widths + def _cuda_graph_runner_config(self) -> CUDAGraphRunnerConfig: return CUDAGraphRunnerConfig( use_cuda_graph=( @@ -419,6 +447,7 @@ def _cuda_graph_runner_config(self) -> CUDAGraphRunnerConfig: cuda_graph_batch_sizes=self._config.cuda_graph_batch_sizes, max_cuda_graph_batch_size=self._config.max_cuda_graph_batch_size, max_beam_width=self._config.max_beam_width, + compact_beam_widths=self._supported_compact_beam_widths(), spec_config=self._config.spec_config, cuda_graph_mem_pool=self._cuda_graph_mem_pool, dynamic_draft_len_mapping=self._dynamic_draft_len_mapping, @@ -650,6 +679,12 @@ def _forward_decoder( gather_context_logits=gather_context_logits, ) inputs = prepared.kwargs + self.cuda_graph_runner.beam_compaction.prepare_inputs( + inputs, + self.cuda_graph_runner.beam_compaction.execution_width( + execution_requests, new_tensors_device is not None + ), + ) runtime_draft_len = prepared.runtime_draft_len if execution_promoted_context_ids: self._inputs.record_promoted_iter_states(scheduled_requests) @@ -730,12 +765,28 @@ def model_forward(self, *, is_dummy: bool, **kwargs: Any) -> Any: and isinstance(metadata, TrtllmAttentionMetadata) else contextlib.nullcontext() ) + compact = kwargs.pop("beam_compaction", None) moe_reclaimer = self._moe_workspace_reclaimer moe_scope = contextlib.nullcontext() if moe_reclaimer is not None: moe_scope = moe_reclaimer.forward(warmup=is_dummy, num_tokens=metadata.num_tokens) - with reclaim_scope, moe_scope: - return self._model_caller(**kwargs) + try: + with reclaim_scope, moe_scope: + if compact is None: + return self._model_caller(**kwargs) + # The reclaimer owns the original metadata. Bind scratch only + # after it has begun tracking this forward's workspace demand. + compact.metadata.workspace = metadata.workspace + compact.metadata.cuda_graph_workspace = metadata.cuda_graph_workspace + kwargs = compact.pack_inputs(kwargs) + outputs = self._model_caller(**kwargs) + metadata.workspace_reclaimable &= compact.metadata.workspace_reclaimable + return compact.restore_logits(outputs) + finally: + if compact is not None and not metadata.is_cuda_graph: + # Reclamation may replace the original workspace tensor; a + # cached view must not keep the previous allocation alive. + compact.metadata.workspace = None @nvtx_range("_forward_step") def _forward_step( diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py index c52cc1e20e8d..917b721bde46 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -744,6 +744,7 @@ def _decoder_runner_config( cuda_graph_specialize_lora=( self.llm_args.lora_config is not None and self.llm_args.lora_config.cuda_graph_specialize_lora), + compact_beam_widths=tuple(self.llm_args.compact_beam_widths), **extra_fields, ) diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 3ef5f8a05178..5f527181508b 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -6353,6 +6353,30 @@ class TorchLlmArgs(BaseLlmArgs): "Lower values trigger more frequent garbage collection.", status="beta") + compact_beam_widths: List[Annotated[int, Field(ge=2)]] = Field( + default_factory=list, + description= + ("Opt-in execution widths (at least 2) for compact variable-width beam search. " + "Empty disables compaction. Only single-GPU dense Llama, Qwen2 and " + "Qwen3 decode with TRTLLM attention is supported. CUDA graphs are " + "captured separately for each width. Changing forward shapes can " + "change floating-point rounding and beam rankings."), + status="prototype") + + @model_validator(mode='after') + def validate_compact_beam_widths(self) -> 'TorchLlmArgs': + if self.compact_beam_widths and (self.enable_lora + or self.lora_config is not None): + raise ValueError("compact_beam_widths does not support LoRA") + capacity = self.max_beam_width or 1 + if any(width > capacity for width in self.compact_beam_widths): + raise ValueError( + "compact_beam_widths must not exceed max_beam_width") + self.compact_beam_widths = sorted( + {width + for width in self.compact_beam_widths if width < capacity}) + return self + cuda_graph_config: Optional[CudaGraphConfigType] = Field( default_factory=CudaGraphConfig, description="CUDA graph config. If true, use CUDA graphs for decoding. \ diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index 341d16689a59..8f136c93abb4 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -4132,7 +4132,7 @@ }, "role": { "default": "both", - "description": "What this server does with the pool. 'both' reads and writes, 'producer' only writes, 'consumer' only reads, and 'capacity' does neither: its ranks lend memory without moving any KV, so they pin no staging buffers and pay no per-request cost.", + "description": "What this server does with the pool. 'both' reads and writes, 'producer' only writes, 'consumer' only reads, and 'capacity' does neither: its ranks lend memory without registering any KV cache, so it needs no GPUDirect RDMA.", "enum": [ "both", "producer", @@ -4172,8 +4172,8 @@ "title": "Segment Size" }, "stage_through_host": { - "default": true, - "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. The only transfer path supported today, so 'false' is ignored with a warning. Costs a copy each way and needs no GPUDirect RDMA.", + "default": false, + "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. For a host whose HCA cannot pin GPU pages; costs a copy each way. Ignored when role is 'capacity'.", "telemetry": { "exclude": true }, @@ -6440,6 +6440,16 @@ "description": "Live objects require the Python LLM API.", "type": "null" }, + "compact_beam_widths": { + "description": "Opt-in execution widths (at least 2) for compact variable-width beam search. Empty disables compaction. Only single-GPU dense Llama, Qwen2 and Qwen3 decode with TRTLLM attention is supported. CUDA graphs are captured separately for each width. Changing forward shapes can change floating-point rounding and beam rankings.", + "items": { + "minimum": 2, + "type": "integer" + }, + "status": "prototype", + "title": "Compact Beam Widths", + "type": "array" + }, "context_parallel_size": { "default": 1, "description": "The context parallel size.", @@ -6948,6 +6958,17 @@ "$ref": "#/$defs/LayerwiseBenchmarksConfig", "status": "prototype" }, + "lm_head_output_dtype": { + "default": "auto", + "description": "Output dtype written by the LM head GEMM. 'auto' uses the model dtype before logits are upcast to float32. 'float32' writes the float32 accumulator directly, avoiding intermediate bf16/fp16 rounding of logits before logprob computation. Weights and inputs remain in the model dtype. Only supported for unquantized LM heads of decoder-only models.", + "enum": [ + "auto", + "float32" + ], + "status": "prototype", + "title": "Lm Head Output Dtype", + "type": "string" + }, "load_format": { "anyOf": [ { diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index b57ddfb733f8..5ff054f2a7ba 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -1774,6 +1774,16 @@ "description": "Live objects require the Python LLM API.", "type": "null" }, + "compact_beam_widths": { + "description": "Opt-in execution widths (at least 2) for compact variable-width beam search. Empty disables compaction. Only single-GPU dense Llama, Qwen2 and Qwen3 decode with TRTLLM attention is supported. CUDA graphs are captured separately for each width. Changing forward shapes can change floating-point rounding and beam rankings.", + "items": { + "minimum": 2, + "type": "integer" + }, + "status": "prototype", + "title": "Compact Beam Widths", + "type": "array" + }, "context_parallel_size": { "default": 1, "description": "The context parallel size.", @@ -2314,6 +2324,17 @@ "$ref": "#/$defs/LayerwiseBenchmarksConfig", "status": "prototype" }, + "lm_head_output_dtype": { + "default": "auto", + "description": "Output dtype written by the LM head GEMM. 'auto' uses the model dtype before logits are upcast to float32. 'float32' writes the float32 accumulator directly, avoiding intermediate bf16/fp16 rounding of logits before logprob computation. Weights and inputs remain in the model dtype. Only supported for unquantized LM heads of decoder-only models.", + "enum": [ + "auto", + "float32" + ], + "status": "prototype", + "title": "Lm Head Output Dtype", + "type": "string" + }, "load_format": { "anyOf": [ { @@ -5682,7 +5703,7 @@ }, "role": { "default": "both", - "description": "What this server does with the pool. 'both' reads and writes, 'producer' only writes, 'consumer' only reads, and 'capacity' does neither: its ranks lend memory without moving any KV, so they pin no staging buffers and pay no per-request cost.", + "description": "What this server does with the pool. 'both' reads and writes, 'producer' only writes, 'consumer' only reads, and 'capacity' does neither: its ranks lend memory without registering any KV cache, so it needs no GPUDirect RDMA.", "enum": [ "both", "producer", @@ -5722,8 +5743,8 @@ "title": "Segment Size" }, "stage_through_host": { - "default": true, - "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. The only transfer path supported today, so 'false' is ignored with a warning. Costs a copy each way and needs no GPUDirect RDMA.", + "default": false, + "description": "Copy pages through a pinned host buffer instead of registering the KV pools with Mooncake. For a host whose HCA cannot pin GPU pages; costs a copy each way. Ignored when role is 'capacity'.", "telemetry": { "exclude": true }, @@ -7997,6 +8018,16 @@ "description": "Live objects require the Python LLM API.", "type": "null" }, + "compact_beam_widths": { + "description": "Opt-in execution widths (at least 2) for compact variable-width beam search. Empty disables compaction. Only single-GPU dense Llama, Qwen2 and Qwen3 decode with TRTLLM attention is supported. CUDA graphs are captured separately for each width. Changing forward shapes can change floating-point rounding and beam rankings.", + "items": { + "minimum": 2, + "type": "integer" + }, + "status": "prototype", + "title": "Compact Beam Widths", + "type": "array" + }, "conditional_disagg_config": { "anyOf": [ { @@ -8615,6 +8646,17 @@ "$ref": "#/$defs/LayerwiseBenchmarksConfig", "status": "prototype" }, + "lm_head_output_dtype": { + "default": "auto", + "description": "Output dtype written by the LM head GEMM. 'auto' uses the model dtype before logits are upcast to float32. 'float32' writes the float32 accumulator directly, avoiding intermediate bf16/fp16 rounding of logits before logprob computation. Weights and inputs remain in the model dtype. Only supported for unquantized LM heads of decoder-only models.", + "enum": [ + "auto", + "float32" + ], + "status": "prototype", + "title": "Lm Head Output Dtype", + "type": "string" + }, "load_format": { "anyOf": [ { diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index 077aa8be8481..092c4939853c 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -181,6 +181,11 @@ "kind": "categorical", "path": "checkpoint_io_policy" }, + { + "capture_policy": "list[int]", + "kind": "value", + "path": "compact_beam_widths" + }, { "capture_policy": "int", "kind": "value", diff --git a/tests/unittest/_torch/compilation/test_phase_selective_forward.py b/tests/unittest/_torch/compilation/test_phase_selective_forward.py index 613975c89b9b..7f223006f665 100644 --- a/tests/unittest/_torch/compilation/test_phase_selective_forward.py +++ b/tests/unittest/_torch/compilation/test_phase_selective_forward.py @@ -43,6 +43,7 @@ def _create_compile_restricted_runner( @pytest.mark.parametrize("enabled", [False, True]) def test_engine_passes_compile_policy_to_decoder_runner(enabled: bool) -> None: engine = Mock() + engine.llm_args.compact_beam_widths = [] engine._compile_only_piecewise_graphs = enabled config = PyTorchModelEngine._decoder_runner_config(engine) runner_type = Mock() diff --git a/tests/unittest/_torch/executor/test_beam_compaction.py b/tests/unittest/_torch/executor/test_beam_compaction.py new file mode 100644 index 000000000000..7cc99c06155d --- /dev/null +++ b/tests/unittest/_torch/executor/test_beam_compaction.py @@ -0,0 +1,622 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from tensorrt_llm._torch.attention.backends.interface import KVCacheParams +from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata +from tensorrt_llm._torch.models.modeling_llama import LlamaForCausalLM +from tensorrt_llm._torch.models.modeling_qwen import Qwen2ForCausalLM +from tensorrt_llm._torch.models.modeling_qwen3 import Qwen3ForCausalLM +from tensorrt_llm._torch.pyexecutor.beam_compaction import BeamCompaction +from tensorrt_llm._torch.pyexecutor.engine.runners.decoder.runner import DecoderRunner +from tensorrt_llm._torch.pyexecutor.workspace import EagerWorkspaceReclaimer + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required") + + +def make_metadata(batch, capacity): + count = batch * capacity + metadata = TrtllmAttentionMetadata( + max_num_requests=batch, + max_num_tokens=count, + max_num_sequences=count, + seq_lens=torch.ones(count, dtype=torch.int32), + num_contexts=0, + beam_width=capacity, + max_seq_len=64, + ) + metadata.kv_cache_params = KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[4] * count) + metadata.request_ids = list(range(batch)) + metadata.prompt_lens = [3] * count + metadata.kv_lens.fill_(5) + metadata.kv_lens_cuda.fill_(5) + metadata.prompt_lens_cpu.fill_(3) + metadata.prompt_lens_cuda.fill_(3) + metadata.host_request_types.fill_(1) + metadata.kv_cache_block_offsets = torch.arange( + count * 8, dtype=torch.int32, device="cuda" + ).reshape(1, count, 2, 4) + metadata.cache_indirection = torch.zeros(batch, capacity, 64, dtype=torch.int32, device="cuda") + return metadata + + +def make_inputs(batch, capacity): + ids = torch.arange(batch * capacity, device="cuda", dtype=torch.int32) + return dict( + input_ids=ids, + position_ids=ids[None, :] + 10, + attn_metadata=make_metadata(batch, capacity), + multimodal_params=[], + mrope_delta_read_seq_slots=None, + mrope_delta_write_seq_slots=None, + ) + + +def fake_forward(**inputs): + values = inputs["input_ids"].float() + inputs["position_ids"][0] + values += inputs["attn_metadata"].kv_lens_cuda[: values.numel()] + return torch.stack((values, values + 1), dim=-1) + + +def forward(inputs): + runner = SimpleNamespace( + _model_caller=fake_forward, + _eager_workspace_reclaimer=None, + _moe_workspace_reclaimer=None, + ) + return DecoderRunner.model_forward(runner, is_dummy=False, **inputs) + + +def request(width, next_width=None, dummy=False): + return SimpleNamespace( + is_dummy=dummy, + get_beam_width_by_iter=lambda for_next_iteration=False: ( + next_width if for_next_iteration and next_width is not None else width + ), + ) + + +@pytest.mark.parametrize("previous,expected", [(False, 2), (True, 4)]) +def test_width_tracks_input_iteration_and_ignores_padding(previous, expected): + batch = SimpleNamespace( + context_requests=[], + generation_requests=[request(2, 4), request(2, 4), request(5, dummy=True)], + ) + assert BeamCompaction(5, (2, 4)).execution_width(batch, previous) == expected + + +@pytest.mark.parametrize( + "requests,contexts", + [ + ([request(3)], []), + ([request(2), request(4)], []), + ([request(5, dummy=True)], []), + ([request(2)], [object()]), + ], +) +def test_unsupported_batches_use_capacity(requests, contexts): + batch = SimpleNamespace(context_requests=contexts, generation_requests=requests) + assert BeamCompaction(5, (2, 4)).execution_width(batch, False) == 0 + + +@pytest.mark.parametrize( + "batch,capacity,width", [(1, 5, 2), (2, 5, 4), (2, 500, 20), (2, 500, 300)] +) +@pytest.mark.parametrize("graph_enabled", [False, True]) +def test_active_rows_and_repeated_metadata_updates(batch, capacity, width, graph_enabled): + inputs = make_inputs(batch, capacity) + source = inputs["attn_metadata"] + compaction = BeamCompaction(capacity, (width,)) + compaction.prepare_inputs(inputs, width) + view = inputs["beam_compaction"] + pointer = view.metadata.kv_lens_cuda.data_ptr() + if graph_enabled: + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + for _ in range(3): + forward(inputs) + torch.cuda.current_stream().wait_stream(stream) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + result = forward(inputs) + for step in range(3): + inputs["input_ids"].add_(1) + inputs["position_ids"].add_(2) + source.kv_lens.add_(1) + source.kv_lens_cuda.add_(1) + source.kv_cache_block_offsets.add_(8) + # Parent IDs remain request-local; their numeric values are not remapped. + source.cache_indirection.fill_(step % width) + compaction.prepare_inputs(inputs, width) + assert inputs["beam_compaction"] is view + assert view.metadata.kv_lens_cuda.data_ptr() == pointer + if graph_enabled: + graph.replay() + else: + result = forward(inputs) + expected = fake_forward(**inputs).reshape(batch, capacity, 2) + actual = result.reshape(batch, capacity, 2) + torch.testing.assert_close(actual[:, :width], expected[:, :width], rtol=0, atol=0) + assert torch.isneginf(actual[:, width:]).all() + torch.testing.assert_close( + view.metadata.kv_cache_block_offsets, + source.kv_cache_block_offsets.index_select(1, view.row_indices), + ) + torch.testing.assert_close( + view.metadata.cache_indirection, source.cache_indirection[:, :width] + ) + + +def test_metadata_views_do_not_share_plans_or_cross_batch_rows(): + inputs = make_inputs(2, 5) + compaction = BeamCompaction(5, (2, 4)) + compaction.prepare_inputs(inputs, 2) + narrow = inputs["beam_compaction"] + compaction.prepare_inputs(inputs, 4) + wide = inputs["beam_compaction"] + assert narrow.row_indices_list == [0, 1, 5, 6] + assert wide.row_indices_list == [0, 1, 2, 3, 5, 6, 7, 8] + assert narrow.metadata.fmha_plan_caches is not wide.metadata.fmha_plan_caches + compaction.prepare_inputs(inputs, 0) + assert "beam_compaction" not in inputs + compaction.clear() + assert not compaction._views + + +def test_multitoken_and_lora_inputs_use_capacity(): + inputs = make_inputs(1, 5) + compaction = BeamCompaction(5, (2,)) + inputs["lora_params"] = object() + compaction.prepare_inputs(inputs, 2) + assert "beam_compaction" not in inputs + inputs.pop("lora_params") + inputs["attn_metadata"].seq_lens = torch.full((5,), 2, dtype=torch.int32) + compaction.prepare_inputs(inputs, 2) + assert "beam_compaction" not in inputs + + +def test_workspace_reclaimer_keeps_source_identity(): + inputs = make_inputs(1, 5) + source = inputs["attn_metadata"] + source.workspace = torch.empty(64, dtype=torch.uint8, device="cuda") + reclaimer = EagerWorkspaceReclaimer(source) + old_workspace = weakref.ref(source.workspace) + source.workspace.resize_(1024) + compaction = BeamCompaction(5, (2,)) + compaction.prepare_inputs(inputs, 2) + view = inputs["beam_compaction"] + + def model(**kwargs): + assert kwargs["attn_metadata"].workspace is source.workspace + source.workspace.resize_(max(source.workspace.numel(), 64)) + return fake_forward(**kwargs) + + runner = SimpleNamespace( + _model_caller=model, + _eager_workspace_reclaimer=reclaimer, + _moe_workspace_reclaimer=None, + ) + for _ in range(5): + DecoderRunner.model_forward(runner, is_dummy=False, **inputs) + assert view.metadata.workspace is None + assert source.workspace.untyped_storage().nbytes() == 64 + assert old_workspace() is None + + +@pytest.mark.parametrize("shared_inputs", [False, True]) +def test_graph_runner_captures_widths_and_replays_updated_rows(monkeypatch, shared_inputs): + from tensorrt_llm._torch.pyexecutor import cuda_graph_runner as graph_module + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import ( + CUDAGraphRunner, + CUDAGraphRunnerConfig, + KeyType, + ) + from tensorrt_llm.mapping import Mapping + + monkeypatch.setattr(graph_module, "_STRICT_BUFFER_CHECK", True) + inputs = make_inputs(2, 5) + config = CUDAGraphRunnerConfig( + use_cuda_graph=True, + cuda_graph_padding_enabled=True, + cuda_graph_batch_sizes=[2], + max_cuda_graph_batch_size=2, + max_beam_width=5, + max_num_tokens=10, + spec_config=None, + cuda_graph_mem_pool=None, + use_mrope=False, + original_max_total_draft_tokens=0, + enable_attention_dp=False, + is_encoder_decoder=False, + batch_size=2, + mapping=Mapping(), + dist=None, + kv_cache_manager_key=None, + compact_beam_widths=(2, 4), + static_input_ids=inputs["input_ids"] if shared_inputs else None, + static_position_ids=inputs["position_ids"][0] if shared_inputs else None, + ) + runner = CUDAGraphRunner(config) + for name in ("input_ids", "position_ids"): + assert (runner.shared_static_tensors[name].data_ptr() == inputs[name].data_ptr()) == ( + shared_inputs + ) + inputs["attn_metadata"].is_cuda_graph = True + key = KeyType(batch_size=2, draft_len=0) + runner.is_warmup_only = True + runner.capture(key, forward, inputs) + runner.is_warmup_only = False + runner.capture(key, forward, inputs) + assert {k.beam_execution_width for k in runner.graphs} == {0, 2, 4} + eager_inputs = make_inputs(2, 5) + eager_ids, eager_positions = eager_inputs["input_ids"], eager_inputs["position_ids"] + for step, width in enumerate((2, 4, 0, 2)): + eager_count = (1 + step % 2) * 5 + eager_inputs["attn_metadata"].seq_lens = torch.ones(eager_count, dtype=torch.int32) + eager_inputs["input_ids"] = eager_ids[:eager_count] + eager_inputs["position_ids"] = eager_positions[:, :eager_count] + runner.beam_compaction.prepare_inputs(eager_inputs, 2) + inputs["input_ids"].add_(7) + inputs["position_ids"].add_(1) + inputs["attn_metadata"].kv_lens_cuda.add_(3) + runner.beam_compaction.prepare_inputs(inputs, width) + result = runner.replay(key._replace(beam_execution_width=width), inputs) + expected = fake_forward(**inputs).reshape(2, 5, 2) + actual = result.reshape(2, 5, 2) + active = width or 5 + torch.testing.assert_close(actual[:, :active], expected[:, :active], rtol=0, atol=0) + assert torch.isneginf(actual[:, active:]).all() + inputs.pop("beam_compaction") + with pytest.raises(RuntimeError, match="refreshed compact metadata"): + runner.replay(key._replace(beam_execution_width=2), inputs) + + +@pytest.mark.parametrize("widths,expected", [([], []), ([4, 2, 2, 5], [2, 4])]) +def test_config_normalizes_execution_widths(widths, expected): + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + + args = TorchLlmArgs.model_construct(max_beam_width=5, compact_beam_widths=widths) + assert args.validate_compact_beam_widths().compact_beam_widths == expected + + +def test_config_rejects_width_above_capacity(): + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + + args = TorchLlmArgs.model_construct(max_beam_width=5, compact_beam_widths=[6]) + with pytest.raises(ValueError, match="max_beam_width"): + args.validate_compact_beam_widths() + + +@pytest.mark.parametrize("width", [-1, 0, 1]) +def test_config_width_schema_requires_at_least_two(width): + from pydantic import TypeAdapter, ValidationError + + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + + schema = TypeAdapter(TorchLlmArgs.model_fields["compact_beam_widths"].annotation) + assert schema.json_schema()["items"]["minimum"] == 2 + with pytest.raises(ValidationError): + schema.validate_python([width]) + assert schema.validate_python([2, 4]) == [2, 4] + + +def test_width_one_uses_capacity_path(): + batch = SimpleNamespace(context_requests=[], generation_requests=[request(1)]) + assert BeamCompaction(5, (2, 4)).execution_width(batch, False) == 0 + + +def test_metadata_host_indices_ignore_default_device(): + inputs = make_inputs(2, 5) + compaction = BeamCompaction(5, (2,)) + with torch.device("cuda"): + compaction.prepare_inputs(inputs, 2) + view = inputs["beam_compaction"] + assert view.row_indices_cpu.device.type == "cpu" + assert view.metadata.seq_lens.device.type == "cpu" + assert view.row_indices.is_cuda + + +@pytest.mark.parametrize("capacity,width", [(5, 2), (5, 4), (500, 20), (500, 300)]) +def test_eager_cache_releases_previous_batch_views(capacity: int, width: int) -> None: + """Changing batch size releases obsolete views and keeps active rows correct.""" + inputs = make_inputs(3, capacity) + source = inputs["attn_metadata"] + input_ids, position_ids = inputs["input_ids"], inputs["position_ids"] + compaction = BeamCompaction(capacity, (width,)) + previous_view = previous_buffer = None + for batch in (1, 2, 3, 1, 3, 2): + count = batch * capacity + source.seq_lens = torch.ones(count, dtype=torch.int32) + source.request_ids = list(range(batch)) + source.kv_lens.add_(1) + source.kv_lens_cuda.add_(1) + inputs["input_ids"] = input_ids[:count] + inputs["position_ids"] = position_ids[:, :count] + compaction.prepare_inputs(inputs, width) + view = inputs["beam_compaction"] + if previous_view is not None: + assert previous_view() is None + assert previous_buffer() is None + assert len(compaction._views) == 1 + assert view.batch_size == batch + pointer = view.metadata.kv_lens_cuda.data_ptr() + compaction.prepare_inputs(inputs, width) + assert inputs["beam_compaction"] is view + assert view.metadata.kv_lens_cuda.data_ptr() == pointer + actual = forward(inputs).reshape(batch, capacity, 2) + expected = fake_forward(**inputs).reshape(batch, capacity, 2) + torch.testing.assert_close(actual[:, :width], expected[:, :width], rtol=0, atol=0) + assert torch.isneginf(actual[:, width:]).all() + previous_view = weakref.ref(view) + previous_buffer = weakref.ref(view.metadata.kv_lens_cuda) + del view + + +@pytest.mark.parametrize("graph_enabled", [False, True]) +def test_eager_cache_keeps_other_sources_and_widths(graph_enabled: bool) -> None: + """Eviction is isolated from other widths and eager or CUDA graph sources.""" + inputs = make_inputs(3, 5) + other_inputs = make_inputs(2, 5) + other_inputs["attn_metadata"].is_cuda_graph = graph_enabled + compaction = BeamCompaction(5, (2, 4)) + compaction.prepare_inputs(inputs, 4) + wide = inputs["beam_compaction"] + compaction.prepare_inputs(other_inputs, 2) + other = other_inputs["beam_compaction"] + other_pointer = other.metadata.kv_lens_cuda.data_ptr() + compaction.prepare_inputs(inputs, 2) + inputs["attn_metadata"].seq_lens = torch.ones(5, dtype=torch.int32) + inputs["input_ids"] = inputs["input_ids"][:5] + inputs["position_ids"] = inputs["position_ids"][:, :5] + compaction.prepare_inputs(inputs, 2) + assert len(compaction._views) == 3 + assert wide in compaction._views.values() + assert other in compaction._views.values() + compaction.prepare_inputs(other_inputs, 2) + assert other_inputs["beam_compaction"] is other + assert other.metadata.kv_lens_cuda.data_ptr() == other_pointer + + +@pytest.mark.parametrize( + "key", + ["multimodal_params", "mrope_delta_read_seq_slots", "mrope_delta_write_seq_slots"], +) +def test_multimodal_and_mrope_inputs_use_capacity(key: str) -> None: + inputs = make_inputs(2, 5) + compaction = BeamCompaction(5, (2,)) + compaction.prepare_inputs(inputs, 2) + assert "beam_compaction" in inputs + + inputs[key] = [object()] if key == "multimodal_params" else torch.tensor([0, 1]) + compaction.prepare_inputs(inputs, 2) + assert "beam_compaction" not in inputs + torch.testing.assert_close(forward(inputs), fake_forward(**inputs), rtol=0, atol=0) + + +def make_compact_runner(model: object) -> SimpleNamespace: + return SimpleNamespace( + model=model, + _config=SimpleNamespace( + compact_beam_widths=(2, 4), + attention_backend=SimpleNamespace(Metadata=TrtllmAttentionMetadata), + is_encode_only=False, + is_spec_decode=False, + use_mrope=False, + is_multimodal=False, + without_logits=False, + torch_compile_enabled=False, + enable_in_graph_sampling=False, + ), + mapping=SimpleNamespace(world_size=1), + sparse_attention_config=None, + ) + + +@pytest.mark.parametrize("model_class", [LlamaForCausalLM, Qwen2ForCausalLM, Qwen3ForCausalLM]) +def test_supported_models_declare_compact_beam(model_class: type) -> None: + runner = make_compact_runner(object.__new__(model_class)) + assert DecoderRunner._supported_compact_beam_widths(runner) == (2, 4) + + +def test_compact_beam_support_is_independent_of_class_name() -> None: + class CompactCapableModel: + supports_compact_beam = True + + runner = make_compact_runner(CompactCapableModel()) + assert DecoderRunner._supported_compact_beam_widths(runner) == (2, 4) + runner.model = object() + assert DecoderRunner._supported_compact_beam_widths(runner) == () + + +@pytest.mark.parametrize("model_class", [LlamaForCausalLM, Qwen2ForCausalLM, Qwen3ForCausalLM]) +@pytest.mark.parametrize("declaration", [None, False, True]) +def test_model_subclasses_require_explicit_compact_beam_support( + model_class: type, declaration: bool | None +) -> None: + attributes = {} if declaration is None else {"supports_compact_beam": declaration} + derived = type("DerivedModel", (model_class,), attributes) + runner = make_compact_runner(object.__new__(derived)) + expected = (2, 4) if declaration else () + assert DecoderRunner._supported_compact_beam_widths(runner) == expected + + +@pytest.mark.parametrize( + "field,value", + [ + ("compact_beam_widths", ()), + ("attention_backend", SimpleNamespace(Metadata=object)), + ("is_encode_only", True), + ("is_spec_decode", True), + ("use_mrope", True), + ("is_multimodal", True), + ("without_logits", True), + ("torch_compile_enabled", True), + ("enable_in_graph_sampling", True), + ], +) +def test_model_capability_preserves_config_guards(field: str, value: object) -> None: + runner = make_compact_runner(object.__new__(Qwen3ForCausalLM)) + setattr(runner._config, field, value) + assert DecoderRunner._supported_compact_beam_widths(runner) == () + + +@pytest.mark.parametrize("distributed", [False, True]) +def test_model_capability_preserves_parallel_and_sparse_guards(distributed: bool) -> None: + runner = make_compact_runner(object.__new__(Qwen3ForCausalLM)) + if distributed: + runner.mapping.world_size = 2 + else: + runner.sparse_attention_config = object() + assert DecoderRunner._supported_compact_beam_widths(runner) == () + + +@pytest.mark.parametrize("field,value", [("enable_lora", True), ("lora_config", object())]) +def test_config_rejects_lora_with_compact_beam(field: str, value: object) -> None: + from tensorrt_llm.llmapi.llm_args import TorchLlmArgs + + args = TorchLlmArgs.model_construct(max_beam_width=5, compact_beam_widths=[2], **{field: value}) + with pytest.raises(ValueError, match="does not support LoRA"): + args.validate_compact_beam_widths() + + +@pytest.mark.parametrize("graph_enabled", [False, True]) +@pytest.mark.parametrize("previous_tokens", [False, True]) +def test_decoder_prepares_compaction_for_actual_execution_batch( + monkeypatch, graph_enabled: bool, previous_tokens: bool +) -> None: + from contextlib import nullcontext + + from tensorrt_llm._torch.pyexecutor import ced_replay + from tensorrt_llm._torch.pyexecutor.cuda_graph_runner import KeyType + + monkeypatch.setattr(ced_replay, "uses_private_decoder_cache", lambda *_: False) + real = SimpleNamespace( + context_requests=[], generation_requests=[request(2, 4)], num_context_requests=0 + ) + padded = SimpleNamespace( + context_requests=[], + generation_requests=[request(2, 4), request(5, dummy=True)], + num_context_requests=0, + ) + metadata = make_metadata(2 if graph_enabled else 1, 5) + inputs = make_inputs(2 if graph_enabled else 1, 5) + inputs["attn_metadata"] = metadata + compaction = BeamCompaction(5, (2, 4)) + events = [] + + def prepare(batch, *args, **kwargs): + assert batch is (padded if graph_enabled else real) + events.append("prepare") + return SimpleNamespace(kwargs=inputs, runtime_draft_len=0) + + class InputsPrepared(Exception): + pass + + def record_copy(): + events.append("copy") + assert inputs["beam_compaction"].width == (4 if previous_tokens else 2) + assert inputs["beam_compaction"].batch_size == (2 if graph_enabled else 1) + raise InputsPrepared + + runner = SimpleNamespace( + kv_cache_manager_key="kv", + model=SimpleNamespace(model=None), + _config=SimpleNamespace(spec_config=None, use_mrope=False, is_spec_decode=False), + _graphs=SimpleNamespace( + attention_dp_execution_batch=lambda batch: batch, + promote_final_contexts=lambda *args, **kwargs: (real, frozenset()), + pad_batch=lambda *args: nullcontext(padded), + select=lambda *args, **kwargs: SimpleNamespace( + key=KeyType(batch_size=2, draft_len=0) if graph_enabled else None, + use_lora_graph=False, + attn_metadata=metadata, + spec_metadata=None, + ), + ), + _metadata=SimpleNamespace(setup_attention=lambda *args: metadata, attn_metadata=metadata), + _inputs=SimpleNamespace(prepare_inputs=prepare, record_input_copy=record_copy), + _trtllm_gen_jit_warmup=False, + moe_load_balancer=None, + _stage_in_graph_sampling=None, + _cuda_graph_metadata_prepare_scope=lambda *args, **kwargs: nullcontext(), + cuda_graph_runner=SimpleNamespace(beam_compaction=compaction), + ) + forward_inputs = SimpleNamespace( + batch=real, + new_tensors_device=object() if previous_tokens else None, + cache_indirection_buffer=None, + gather_context_logits=False, + enable_spec_decode=False, + runtime_draft_len=0, + ) + with pytest.raises(InputsPrepared): + DecoderRunner._forward_decoder( + runner, + forward_inputs, + SimpleNamespace(get_resource_manager=lambda _: object()), + is_dummy=False, + ) + assert events == ["prepare", "copy"] + + +@pytest.mark.parametrize("compact_enabled", [False, True]) +@pytest.mark.parametrize("fail", [False, True]) +@pytest.mark.parametrize("graph_enabled", [False, True]) +def test_compact_forward_preserves_both_workspace_scopes( + compact_enabled: bool, fail: bool, graph_enabled: bool +) -> None: + from contextlib import contextmanager + + inputs = make_inputs(1, 5) + metadata = inputs["attn_metadata"] + metadata.is_cuda_graph = graph_enabled + if compact_enabled: + BeamCompaction(5, (2,)).prepare_inputs(inputs, 2) + view = inputs.get("beam_compaction") + events = [] + + @contextmanager + def scope(name): + events.append(name + ":enter") + try: + yield + finally: + events.append(name + ":exit") + + def model(**kwargs): + assert "beam_compaction" not in kwargs + assert kwargs["attn_metadata"] is (view.metadata if view else metadata) + events.append("model") + if fail: + raise ValueError("forward failed") + return fake_forward(**kwargs) + + def moe_scope(*, warmup, num_tokens): + assert warmup is False + assert num_tokens == 5 + return scope("moe") + + runner = SimpleNamespace( + _model_caller=model, + _eager_workspace_reclaimer=SimpleNamespace(forward=lambda _: scope("attention")), + _moe_workspace_reclaimer=SimpleNamespace(forward=moe_scope), + ) + if fail: + with pytest.raises(ValueError, match="forward failed"): + DecoderRunner.model_forward(runner, is_dummy=False, **inputs) + else: + result = DecoderRunner.model_forward(runner, is_dummy=False, **inputs) + expected = fake_forward(**inputs) + width = 2 if compact_enabled else 5 + torch.testing.assert_close(result[:width], expected[:width], rtol=0, atol=0) + if compact_enabled: + assert torch.isneginf(result[width:]).all() + assert events == ["attention:enter", "moe:enter", "model", "moe:exit", "attention:exit"] + if view is not None and not graph_enabled: + assert view.metadata.workspace is None diff --git a/tests/unittest/_torch/executor/test_pytorch_model_engine.py b/tests/unittest/_torch/executor/test_pytorch_model_engine.py index fb7dcf8470a6..fb4cb9703bac 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -976,6 +976,7 @@ def test_sparse_sequence_mode_uses_promoted_context_cursor(self) -> None: def test_graph_key_forwards_promoted_context_ids(self) -> None: runner = Mock() runner.config = SimpleNamespace() + runner.beam_compaction.execution_width.return_value = 0 runner._resolve_sample_type.return_value = SampleType.FULL runner._get_seq_len_mode.return_value = True request = _make_request_stub(7) diff --git a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py index 4db84b3894be..031201324015 100644 --- a/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py +++ b/tests/unittest/llmapi/test_trtllm_serve_config_schemas.py @@ -54,6 +54,12 @@ def schemas(tmp_path_factory: pytest.TempPathFactory) -> dict[str, dict]: [ (generator.SERVE_SCHEMA, {}, True), (generator.SERVE_SCHEMA, {"max_batch_size": 8}, True), + (generator.SERVE_SCHEMA, {"compact_beam_widths": [], "max_beam_width": 5}, True), + (generator.SERVE_SCHEMA, {"compact_beam_widths": [2, 4], "max_beam_width": 5}, True), + (generator.SERVE_SCHEMA, {"compact_beam_widths": [1], "max_beam_width": 5}, False), + (generator.SERVE_SCHEMA, {"compact_beam_widths": [0], "max_beam_width": 5}, False), + (generator.SERVE_SCHEMA, {"compact_beam_widths": [-1], "max_beam_width": 5}, False), + (generator.SERVE_SCHEMA, {"compact_beam_widths": [2.5], "max_beam_width": 5}, False), (generator.SERVE_SCHEMA, {"max_batch_szie": 8}, False), (generator.SERVE_SCHEMA, {"max_batch_size": "eight"}, False), (generator.SERVE_SCHEMA, {"moe_config": {"backend": "CUTLASS"}}, True),