From 3de73011b6ecd677e268b82d8fdb2a40638aa8df Mon Sep 17 00:00:00 2001 From: yifanQi98 <73160040+yifanQi98@users.noreply.github.com> Date: Thu, 22 Oct 2026 15:28:28 +0800 Subject: [PATCH] [None][perf] Compact active beam rows for model forward Add opt-in compact execution for supported variable-width beam search while preserving capacity-stride KV allocation and sampler layouts. Use per-width CUDA graph metadata and restore inactive logits safely. Prepare views from the actual execution batch, preserve both workspace scopes, and cover input, graph, capability and cleanup boundaries. Refresh generated serving schemas and configuration coverage. Signed-off-by: yifanQi98 <73160040+yifanQi98@users.noreply.github.com> --- docs/source/features/sampling.md | 50 ++ tensorrt_llm/_torch/models/modeling_llama.py | 4 +- tensorrt_llm/_torch/models/modeling_qwen.py | 7 +- tensorrt_llm/_torch/models/modeling_qwen3.py | 7 +- .../_torch/pyexecutor/beam_compaction.py | 205 ++++++ .../_torch/pyexecutor/cuda_graph_runner.py | 83 ++- .../engine/runners/decoder/config.py | 3 +- .../engine/runners/decoder/runner.py | 55 +- .../_torch/pyexecutor/model_engine.py | 1 + tensorrt_llm/llmapi/llm_args.py | 24 + .../schemas/trtllm-serve-config.schema.json | 20 +- .../trtllm-serve-disagg-config.schema.json | 30 +- ...trtllm-serve-visual-gen-config.schema.json | 4 +- .../usage/llm_args_golden_manifest.json | 5 + .../test_phase_selective_forward.py | 1 + .../_torch/executor/test_beam_compaction.py | 615 ++++++++++++++++++ .../executor/test_pytorch_model_engine.py | 1 + .../test_trtllm_serve_config_schemas.py | 6 + 18 files changed, 1091 insertions(+), 30 deletions(-) create mode 100644 tensorrt_llm/_torch/pyexecutor/beam_compaction.py create mode 100644 tests/unittest/_torch/executor/test_beam_compaction.py diff --git a/docs/source/features/sampling.md b/docs/source/features/sampling.md index 0203d05d2f0f..92c748800cc2 100644 --- a/docs/source/features/sampling.md +++ b/docs/source/features/sampling.md @@ -323,6 +323,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 e0d5abd61b35..8465bf0ae5d3 100644 --- a/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py +++ b/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py @@ -27,6 +27,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) @@ -75,6 +76,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( @@ -149,6 +152,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, ...] = () class CUDAGraphRunner: @@ -163,6 +167,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 @@ -413,16 +419,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: @@ -681,6 +691,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] @@ -734,8 +773,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 @@ -788,6 +831,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( @@ -960,6 +1010,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( @@ -973,6 +1027,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"]) @@ -1325,6 +1383,7 @@ def clear(self): 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 2480fecb6ce2..de08ec545b34 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 @@ -408,6 +409,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=( @@ -417,6 +445,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, @@ -636,6 +665,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) @@ -716,12 +751,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 221f9145b17c..9b403895e619 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_engine.py +++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py @@ -740,6 +740,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 d567596628e1..0031e9d5f28c 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 8b268a9df3f2..1fdd2416f48f 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 }, @@ -6285,7 +6285,7 @@ "type": "string" } }, - "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc1/_static/schemas/trtllm-serve-config.schema.json", + "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc2/_static/schemas/trtllm-serve-config.schema.json", "$schema": "https://json-schema.org/draft/2020-12/schema", "additionalProperties": false, "description": "Static YAML checks only; runtime validation remains authoritative.", @@ -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.", @@ -7757,6 +7767,6 @@ } }, "required": [], - "title": "TensorRT-LLM 1.4.0rc1 trtllm-serve --config", + "title": "TensorRT-LLM 1.4.0rc2 trtllm-serve --config", "type": "object" } diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index c9a19899dff5..4be05789ecf4 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.", @@ -5682,7 +5692,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 +5732,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 }, @@ -7846,7 +7856,7 @@ "type": "string" } }, - "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc1/_static/schemas/trtllm-serve-disagg-config.schema.json", + "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc2/_static/schemas/trtllm-serve-disagg-config.schema.json", "$schema": "https://json-schema.org/draft/2020-12/schema", "additionalProperties": false, "description": "Static YAML checks only; runtime validation remains authoritative.", @@ -7997,6 +8007,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": [ { @@ -9502,6 +9522,6 @@ } }, "required": [], - "title": "TensorRT-LLM 1.4.0rc1 trtllm-serve disaggregated --config", + "title": "TensorRT-LLM 1.4.0rc2 trtllm-serve disaggregated --config", "type": "object" } diff --git a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json index cffa1974bcdc..c4f2e4b5320c 100644 --- a/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-visual-gen-config.schema.json @@ -980,7 +980,7 @@ "type": "object" } }, - "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc1/_static/schemas/trtllm-serve-visual-gen-config.schema.json", + "$id": "https://nvidia.github.io/TensorRT-LLM/1.4.0rc2/_static/schemas/trtllm-serve-visual-gen-config.schema.json", "$schema": "https://json-schema.org/draft/2020-12/schema", "additionalProperties": false, "description": "Static YAML checks only; runtime validation remains authoritative.", @@ -1104,6 +1104,6 @@ "status": "prototype" } }, - "title": "TensorRT-LLM 1.4.0rc1 trtllm-serve --visual_gen_args", + "title": "TensorRT-LLM 1.4.0rc2 trtllm-serve --visual_gen_args", "type": "object" } diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index e58f42598f2d..a70f8c2ae81c 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..e8ff5a6468ce --- /dev/null +++ b/tests/unittest/_torch/executor/test_beam_compaction.py @@ -0,0 +1,615 @@ +# 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 + + +def test_graph_runner_captures_widths_and_replays_updated_rows(monkeypatch): + 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) + 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), + ) + runner = CUDAGraphRunner(config) + inputs = make_inputs(2, 5) + 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 3a7293bd7e9d..92b7d774fb14 100644 --- a/tests/unittest/_torch/executor/test_pytorch_model_engine.py +++ b/tests/unittest/_torch/executor/test_pytorch_model_engine.py @@ -930,6 +930,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),