[None][fix] Forward layer_idx for VSWA + narrow FLASHINFER+spec rejection - #15692
Draft
mihai-chiorean wants to merge 18 commits into
Draft
[None][fix] Forward layer_idx for VSWA + narrow FLASHINFER+spec rejection#15692mihai-chiorean wants to merge 18 commits into
mihai-chiorean wants to merge 18 commits into
Conversation
…ckends (VSWA) Both FlashInfer and Vanilla attention backends called kv_cache_manager.get_batch_cache_indices() without a layer_idx argument. For VSWA models (e.g. Qwen3-Next family) the manager holds multiple KV pools with independent page numbering, so the argless call raises: ValueError: layer_idx or window_size must be provided for VSWA FlashInfer fix: detect _vswa_layer_to_pool at prepare() time and pass the primary pool's representative layer_idx to the initial block-count fetch. Per-pool indices are already rebuilt in the existing VSWA block below. Vanilla fix: prepare() detects kv_cache_manager.is_vswa and defers the block-ID lookup (setting block_ids_per_seq=None). forward() then calls get_batch_cache_indices(layer_idx=self.layer_idx) so each layer reads from the correct pool's page table. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
…yle) The blanket rejection at py_executor_creator.py:409 blocked ALL FlashInfer + one-engine speculative decoding (MTP) combinations, including the non-overlap-scheduler path that is actually viable. Root cause: BatchDecodeWithPagedKVCacheWrapper assumes q_len=1 per request. With MTP, generation requests carry num_nextn+1 tokens each, breaking the decode kernel. Fix (vLLM reorder_batch_threshold pattern adapted): - Add FlashInferAttentionMetadata._plan_mtp_gen_prefill(): plans a fresh BatchPrefillWithPagedKVCacheWrapper for the generation sub-batch using the rebased generation qo_indptr and paged_kv_indptr_decode slices. - In forward_impl(): detect MTP (gen seq_lens.max() > 1) and when not is_cuda_graph, route generation tokens through the prefill wrapper (mtp_gen_forward) instead of decode_forward. - Narrow the py_executor_creator guard to only reject when CUDA graphs are also enabled (where variable q_len truly is incompatible). CUDA-graph mode with FlashInfer + MTP remains unsupported (variable q_len per request cannot be captured). Non-graph mode (disable_overlap_ scheduler=True, no cuda_graph_config) now works end-to-end. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
…llow-up) The initial VSWA fix only handled KVCacheManagerV2 (checked via layer_to_pool_mapping_dict). KVCacheManagerV1 also exposes is_vswa=True when it has multiple distinct window sizes, but lacks the per-pool dict so _vswa_layer_to_pool stays None. The argless get_batch_cache_indices() call then raises ValueError: layer_idx or window_size must be provided. Add a fallback for V1: when is_vswa is True but _vswa_layer_to_pool is None, pick the first key from kv_cache_manager.layer_offsets and pass it as layer_idx. This resolves the window_size via the existing layer_offsets → max_attention_window_vec lookup path (V1 path in get_batch_cache_indices lines 1338-1342). Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
…ow models The V1 VSWA detection checked kv_cache_manager.is_vswa, but is_vswa has an additional 'all(w > 0)' guard that excludes models with a full-attention sentinel window (value -2147483647). Such models have len(max_attention_window_vec) > 1 but is_vswa = False, so the V1 fallback never fired. Fix: use len(max_attention_window_vec) > 1 directly as the trigger, matching the exact condition in get_batch_cache_indices that raises. This covers all V1 multi-window configurations regardless of sentinel values. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Contributor
📝 WalkthroughWalkthroughAdds MTP (multi-token prediction) speculative decoding support to the FlashInfer attention backend via a new paged-prefill wrapper for generation sequences with q_len > 1. Fixes VSWA models in both FlashInfer and Vanilla backends to use layer-aware block ID lookup. Relaxes the one-engine speculative decoding compatibility check to block only CUDA-graph configurations. ChangesFlashInfer MTP and VSWA fixes + executor guard update
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested labels
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/py_executor_creator.py (1)
402-424: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winMove the CUDA-graph guard outside the overlap-scheduler branch.
As written, users who set
disable_overlap_scheduler=Trueskip this validation entirely, but FlashInfer’s MTP path still cannot run under CUDA graphs and will fall back to the decode wrapper for variableq_len.Proposed fix
- if not llm_args.disable_overlap_scheduler and spec_config is not None: + if (spec_config is not None and llm_args.attn_backend == "FLASHINFER" + and getattr(llm_args, "cuda_graph_config", None) is not None): + raise ValueError( + f"FLASHINFER attention backend is not supported with one-engine speculative " + f"decoding mode '{spec_config.spec_dec_mode.name}' when CUDA graphs are " + f"enabled (cuda_graph_config is set). The FlashInfer MTP generation path " + f"requires variable q_len per request which is incompatible with CUDA-graph " + f"capture. Either disable CUDA graphs or use 'TRTLLM' attention backend.") + + if not llm_args.disable_overlap_scheduler and spec_config is not None: if not spec_config.spec_dec_mode.support_overlap_scheduler(): logger.warning( f"Disable overlap scheduler for speculation mode {spec_config.spec_dec_mode.name}" ) llm_args.disable_overlap_scheduler = True - - # FLASHINFER + overlap scheduler + one-engine spec-dec: not supported. - # The overlap scheduler pipelines decode and KV-append across steps; - # with multi-token generation requests (MTP/EAGLE) this conflicts with - # how FlashInfer plans its decode wrapper. Disable-overlap-scheduler - # users can use FLASHINFER + MTP without the overlap scheduler — the - # generation sub-batch is routed through the paged-prefill wrapper - # (see FlashInferAttentionMetadata._plan_mtp_gen_prefill). - # CUDA-graph mode is also unsupported for variable-q_len generation. - if (llm_args.attn_backend == "FLASHINFER" - and getattr(llm_args, "cuda_graph_config", None) is not None): - raise ValueError( - f"FLASHINFER attention backend is not supported with one-engine speculative " - f"decoding mode '{spec_config.spec_dec_mode.name}' when CUDA graphs are " - f"enabled (cuda_graph_config is set). The FlashInfer MTP generation path " - f"requires variable q_len per request which is incompatible with CUDA-graph " - f"capture. Either disable CUDA graphs or use 'TRTLLM' attention backend.")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/py_executor_creator.py` around lines 402 - 424, The CUDA-graph validation in py_executor_creator.py is incorrectly nested under the overlap-scheduler check, so FlashInfer one-engine speculative decoding can bypass the guard when disable_overlap_scheduler is already true. Move the CUDA-graph guard to a separate check in the same setup flow around llm_args and spec_config, using the existing symbols llm_args.attn_backend, cuda_graph_config, and spec_config.spec_dec_mode.name, so FLASHINFER + MTP is rejected whenever CUDA graphs are enabled regardless of overlap-scheduler state.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/attention_backend/flashinfer.py`:
- Around line 1977-1981: The MTP detection path in flashinfer.py is
synchronizing the decode hot path by reading from metadata.seq_lens_cuda and
calling max().item() on a CUDA tensor. Update the generation slice in the decode
logic to use metadata.seq_lens instead of metadata.seq_lens_cuda so the
is_mtp_gen check stays on the host side, and keep the existing
num_contexts/num_generations slicing and flashinfer decode flow unchanged.
---
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor_creator.py`:
- Around line 402-424: The CUDA-graph validation in py_executor_creator.py is
incorrectly nested under the overlap-scheduler check, so FlashInfer one-engine
speculative decoding can bypass the guard when disable_overlap_scheduler is
already true. Move the CUDA-graph guard to a separate check in the same setup
flow around llm_args and spec_config, using the existing symbols
llm_args.attn_backend, cuda_graph_config, and spec_config.spec_dec_mode.name, so
FLASHINFER + MTP is rejected whenever CUDA graphs are enabled regardless of
overlap-scheduler state.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: a5c4d703-a7fe-41f9-af74-9cdbe6671f2e
📒 Files selected for processing (3)
tensorrt_llm/_torch/attention_backend/flashinfer.pytensorrt_llm/_torch/attention_backend/vanilla.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.py
Comment on lines
+1977
to
+1981
| gen_seq_lens = metadata.seq_lens_cuda[ | ||
| num_contexts:num_contexts + num_generations] | ||
| is_mtp_gen = (num_generations > 0 | ||
| and not metadata.is_cuda_graph | ||
| and gen_seq_lens.max().item() > 1) |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Locate the target code and the metadata definitions it depends on.
rg -n "gen_seq_lens|max\\(\\)\\.item\\(\\)|is_mtp_gen|seq_lens_cuda|seq_lens\\b" tensorrt_llm/_torch/attention_backend/flashinfer.py tensorrt_llm/_torch -g '!**/__pycache__/**'
echo
echo "=== flashinfer.py around the target lines ==="
sed -n '1950,2005p' tensorrt_llm/_torch/attention_backend/flashinfer.py
echo
echo "=== attention developer guide header / relevant sections ==="
fd -a ATTENTION_DEVELOPER_GUIDE.md tensorrt_llm/_torch -x sh -c 'echo "--- {}"; sed -n "1,260p" "{}"'
echo
echo "=== metadata class definitions mentioning seq_lens ==="
rg -n "class .*Metadata|seq_lens|seq_lens_cuda" tensorrt_llm/_torch -g '!**/__pycache__/**' -A 4 -B 4Repository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== flashinfer.py outline ==="
ast-grep outline tensorrt_llm/_torch/attention_backend/flashinfer.py --view expanded | sed -n '1,260p'
echo
echo "=== flashinfer.py target region ==="
sed -n '1958,1995p' tensorrt_llm/_torch/attention_backend/flashinfer.py
echo
echo "=== metadata field definitions in flashinfer.py ==="
rg -n "class .*Metadata|seq_lens|seq_lens_cuda|is_cuda_graph|num_contexts|num_generations" tensorrt_llm/_torch/attention_backend/flashinfer.py -A 4 -B 4
echo
echo "=== top-level metadata helpers that mention seq_lens on host vs cuda ==="
rg -n "self\\.(seq_lens|seq_lens_cuda)|def .*seq_lens|seq_lens_host|seq_lens_cpu|seq_lens_runtime" tensorrt_llm/_torch/attention_backend -g '!**/__pycache__/**' -A 2 -B 2Repository: NVIDIA/TensorRT-LLM
Length of output: 50375
Avoid syncing the decode hot path for MTP detection (flashinfer.py:1977-1981). metadata.seq_lens is already available here, so use it instead of seq_lens_cuda for the generation slice; gen_seq_lens.max().item() on a CUDA tensor adds an avoidable device-to-host sync on every non-graph decode step.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/attention_backend/flashinfer.py` around lines 1977 -
1981, The MTP detection path in flashinfer.py is synchronizing the decode hot
path by reading from metadata.seq_lens_cuda and calling max().item() on a CUDA
tensor. Update the generation slice in the decode logic to use metadata.seq_lens
instead of metadata.seq_lens_cuda so the is_mtp_gen check stays on the host
side, and keep the existing num_contexts/num_generations slicing and flashinfer
decode flow unchanged.
…nch (closes FIXME) The VSWA / linear-attention branch of KVCacheManager.__init__ computed blocks_per_window (window_size -> (primary, secondary)) but never set self.blocks_in_primary_pool or self.blocks_in_secondary_pool, leaving them unset for any non-estimation run on: - Pure-VSWA models (e.g. Qwen2.5-1M-style sliding-window-only) - Hybrid Mamba2 models using CppMambaHybridCacheManager Downstream, FlashInferAttentionMetadata.__post_init__ reads kv_cache_manager.blocks_in_primary_pool, raising AttributeError. Fix: after blocks_per_window is finalised (including the optional MPI allreduce), set blocks_in_primary_pool / blocks_in_secondary_pool to the maximum across all windows. FlashInfer already uses layer_idx for per-pool dispatch; the scalar attribute only needs to represent the upper bound (largest window), which is the correct semantic. Empirically validated on Spark (trtllm-smoke, Qwen3.6-35B-A3B NVFP4, attn_backend=FLASHINFER): the KV cache manager now initialises cleanly with two pools reported: primary blocks=8 [window=-2147483647], primary blocks=257 [window=4096] The fix resolves the FIXME added at this line. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
mihai-chiorean
marked this pull request as draft
June 28, 2026 14:52
The pure-decode branch (num_contexts == 0) of Mamba2Metadata.prepare() left query_start_loc as None, on the assumption that downstream callers only used query_start_loc_long. GatedDeltaNetMixer violates that: it always packs mamba_metadata.query_start_loc into kwargs and slices it, which raises TypeError when None. Reuse the existing int32 _arange_buffer (sized max_batch_size + 1) for the all-decode case. Numerically identical to what cu_seqlens[:batch_size+1] would be when every sequence contributes one token ([0, 1, ..., batch_size]). Surfaces on FlashInfer + Qwen3-Next (hybrid Mamba2) where the pure-decode / CUDA-graph capture path is first exercised. Non-FlashInfer paths reach the gdn_mixer with mixed batches where num_contexts > 0, so query_start_loc was always populated. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Spec-decoding code (eagle3._prepare_attn_metadata_for_spec_dec, drafting_loops) reads and writes spec_decoding_packed_mask and friends on attn_metadata. These fields were declared only on TrtllmAttentionMetadata, not on FlashInferAttentionMetadata or the AttentionMetadata base. This was invisible while FLASHINFER + spec_dec was blanket-rejected. Now that the gate has been narrowed to one-engine + CUDA-graph only, FlashInfer + MTP exercises the code path and AttributeErrors. Schema-only fix: add the missing Optional[torch.Tensor] = None fields to FlashInferAttentionMetadata so the code path runs. Actually plumbing the packed mask into BatchPrefillWithPagedKVCacheWrapper.plan(custom_mask=...) for correct numerics is a separate concern and not addressed here. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
For one-engine speculative decoding (MTP, Eagle3) on hybrid Mamba models such as Qwen3-Next / Qwen3.6-35B-A3B-NVFP4, the separate-draft-KV-cache layout places the MTP draft layer (layer_idx == num_hidden_layers) in its own KVCacheManager. During the draft forward pass, the worker swaps that draft manager onto attn_metadata via SpecConfig.draft_kv_cache_context (and prepare_attn_metadata_for_draft_replay). Both helpers, however, short-circuit on `isinstance(attn_metadata, TrtllmAttentionMetadata)`, so when attn_backend != "TRTLLM" the swap never happens. The draft attention call then resolves metadata.kv_cache_manager.get_buffers(layer_idx=num_hidden_layers) against the target manager, whose layer_offsets only covers the dense range 0..(num_hidden_layers-1), and crashes with `KeyError`. Symptom on Spark: Qwen3.6-35B-A3B-NVFP4 + attn_backend=FLASHINFER + MTPDecodingConfig(num_nextn_predict_layers=1) raises KeyError: 40 at resource_manager.py:1435 (`layer_offset = self.layer_offsets[layer_idx]`) inside flashinfer.py:1853. Fix: for non-TRTLLM backends, refuse to create the separate draft KV cache manager. `extract_mamba_kv_cache_params` already extends the hybrid layer masks with MTP/draft attention entries when `layer_mask` is None, and `get_pp_layers` extends pp_layers symmetrically, so the combined target+draft layout places the draft layer in the target manager's layer_offsets and get_buffers resolves correctly. TRTLLM- attn retains the separate layout it already supports. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
On Qwen3-Next-family hybrid Mamba models, layer 0 is a Mamba2 / linear- attention layer. FlashInferAttentionMetadata.prepare() was resolving _vswa_init_layer to the rep layer of pool 0 - which is the linear / recurrent pool. BlockManager::getFreeBlock returns NEGATIVE placeholder block IDs for Mamba recurrent-state slots. Those negatives populated _paged_kv_indices and leaked into _vswa_pool_buf_0 via the primary_buf copy. On the second autotuner warmup (post-estimation pool resize), swap_paged_kv_indices_for_layer for a full-attn layer surfaced stale linear-pool data, append_paged_kv_cache received negative kv_indices and crashed with an illegal memory access. Fix: in _vswa_init_layer resolution, prefer pools whose window > 0 (real SWA / full-attention) over linear pools (window == -INT_MAX). The fallback for the all-linear edge case is preserved. Same heuristic applied to the V1 KVCacheManager fallback added in 464e91b. Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
Signed-off-by: Mihai Chiorean <mihai.v.chiorean@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Two related fixes that progressively unblock non-
TRTLLMattention backends for SM120/SM121 and speculative-decoding use cases. Discovered while validating the Qwen3.6-35B-A3B-NVFP4 port (PR #15680) on DGX Spark.Fix 1 — Forward
layer_idxfor VSWA (3 commits,flashinfer.py+vanilla.py)KVCacheManager.get_batch_cache_indices()requireslayer_idx(orwindow_size) when the model uses Varying Sliding Window Attention (VSWA). TheTRTLLMattention backend forwards it correctly. TheFLASHINFERandVANILLAbackends callget_batch_cache_indices()with no per-layer args, so any model withlen(max_attention_window_vec) > 1crashes at LLM init:Repro on the model that surfaced it:
The fix uses a three-tier VSWA detection:
kv_cache_manager._vswa_layer_to_pooldict (V2 KVCacheManager)len(max_attention_window_vec) > 1on the manager (V1 path)is_vswaproperty is not reliable (itsall(w > 0)guard returns False for models withINT_MINsentinel full-attention windows)flashinfer.pyis updated to plumblayer_idxthroughprepare().vanilla.pydefers the per-layer call toVanillaAttention.forward()whereself.layer_idxis available, leavingblock_ids_per_seq = Noneat metadata-build time for VSWA models.Fix 2 — Narrow FLASHINFER + speculative decoding rejection (1 commit,
py_executor_creator.py)The previous guard at
py_executor_creator.py:409-416blanket-rejected ANYattn_backend="FLASHINFER"+spec_config != Nonecombination with an error message that advertised "one-engine speculative decoding mode 'MTP_EAGLE_ONE_MODEL'" but actually fired for two-engine modes too.vLLM solved the equivalent problem at
vllm/v1/attention/backends/flashinfer.pywithreorder_batch_threshold: route gen+spec requests withq_lens > 1through the prefill wrapper (which accepts arbitraryqo_indptr). FlashInfer also supportsBatchDecodeWithPagedKVCacheWrapper.plan(q_len_per_req=N)since 0.6.x (TRT-LLM pinsflashinfer-python==0.6.12, both options are present).This PR takes a conservative first step: narrow the rejection so it only fires when both
FLASHINFERandcuda_graph_configare set. Non-CUDA-graph speculative paths can now proceed (the CUDA-graph case still has stable-buffer concerns whenq_len_per_reqvaries between captures). The wrapper-level routing (full vLLM-style) remains as a follow-up — see the discovered blockers below.Validation
Smoke-tested on DGX Spark GB10 (SM121) against
nvidia/Qwen3.6-35B-A3B-NVFP4. Baseline path (TRTLLM backend + MTP) still works unchanged.Two additional pre-existing structural incompatibilities between FlashInfer and the Qwen3.6 / Qwen3-Next family were discovered while testing — these are NOT addressed in this PR and motivate follow-up work:
Mamba2 hybrid + FlashInfer:
CppMambaHybridCacheManager(used by Mamba2-hybrid models including Qwen3-Next) does not exposeblocks_in_primary_pool, whichFlashInferAttentionMetadata.__post_init__requires. Crash:Tracking issue to follow.
FlashInferAttentionMetadata.spec_decoding_packed_mask: missing. The MTP/eagle3 path attensorrt_llm/_torch/models/eagle3.py:606expectsTrtllmAttentionMetadata's schema. Tracking issue to follow.Both blockers are independent of (and downstream of) the fixes in this PR. With both also resolved, the
attn_backend="FLASHINFER" + MTPcombination would be unblocked end-to-end on Qwen3-Next-family hybrid models.Test plan
TRTLLMbackend + MTP — works (1.28× TPOT vs baseline)flashinfer.py+vanilla.pyno longer crash on VSWA models atprepare()py_executor_creator.py/bot runOut of scope
reorder_batch_thresholdrouting for FlashInfer + spec_dec (the deeper Mamba + spec_decoding_packed_mask gaps make a complete end-to-end measurement impossible until those land separately)FlashInferAttentionMetadata(separate issue)spec_decoding_packed_maskschema parity between attention backends (separate issue)PR Checklist
PR description clearly explains what and why.
PR follows TRT-LLM CODING GUIDELINES.
Test cases provided for new code paths (smoke-verified; full unit tests pending follow-up).
No new dependencies introduced.
[CODEOWNERS] not changed.
Documentation: no public API changed.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
Summary by CodeRabbit
New Features
Bug Fixes