[None][feat] Support DFlash RoPE, sliding-window configuration, and TRTLLM-gen attention backend

[jhaotingc]

Description

Qwen3.6-35B-A3B-DFlash mixes sliding-window and full-attention draft layers. The DFlash path previously assumed simpler layer semantics: it shared one RoPE cache without validating the effective per-layer RoPE configuration, did not pass each layer's local-window arguments to FlashAttention, and could size a separate V1 draft KV cache from an unsplit budget or a target-clamped sequence length.

This PR adds end-to-end support for the combination of z-lab/Qwen3.6-35B-A3B-DFlash and nvidia/Qwen3.6-35B-A3B-NVFP4:

• Validate that all DFlash draft layers sharing the flattened RoPE cache have the same effective RoPE parameters. Mixed layer types are allowed when their RoPE configuration is uniform; incompatible layers fail early with a clear error.
• Derive FlashAttention causal and window_size arguments per draft layer. Qwen-style sliding-attention layers use their configured causal local window, full-attention layers remain global, and Laguna uses its checkpoint-defined uniform causal window.
• Derive the draft KV-cache attention windows from the draft model configuration and preserve the original pre-clamping max_seq_len when building the separate draft manager.
• Split max_gpu_total_bytes between target and draft before constructing the final V1 variable-window managers. The temporary estimation manager deliberately avoids installing the VSWA layout before that split is available.
• Centralize variable sliding-window KV-cache detection in uses_vswa_kv_cache_layout() and exclude negative recurrent-state sentinels used by hybrid linear-attention models.
• Add unit coverage for RoPE validation, per-layer attention masks, original-sequence-length propagation, budget splitting, and VSWA detection, plus Qwen3.6 NVFP4 + DFlash GSM8K accuracy coverage.

TRTLLM-Gen speculative attention backend

The second commit adds an optional TRTLLM-Gen backend for DFlash drafter attention:

• Add speculative_config.attention_backend with the main-model-style values VANILLA (default) and TRTLLM.
• Forward the existing speculative_config.advanced_sampling_mode into DFlash metadata so DFlash honors the configured sampling specialization instead of silently using the default.
• TRTLLM stores the pooled DFlash context K/V in a private HND paged cache shared by FlashInfer's paged append and the TRTLLM-Gen launchers. It dispatches the Blackwell generated context/decode FMHA kernels with CUDA-graph-compatible workspace and counter buffers.
• VANILLA preserves the existing FlashAttention implementation and contiguous context K/V cache.
• Add configuration validation, paged-cache construction coverage, and Qwen3.6 accuracy coverage for both backends.

DFlash 7 trtllm-serve example

The TRTLLM-backend benchmark used the following equivalent standalone serving configuration. attention_backend uses the same public value as the main-model attention backend:

BASE_MODEL=${BASE_MODEL:-nvidia/Qwen3.6-35B-A3B-NVFP4}
DRAFT_MODEL=${DRAFT_MODEL:-z-lab/Qwen3.6-35B-A3B-DFlash}
CONFIG=${CONFIG:-dflash7-trtllm.yml}

cat > "${CONFIG}" <<YAML
disable_overlap_scheduler: false
print_iter_log: false
enable_chunked_prefill: true
enable_attention_dp: false
scheduler_config:
  capacity_scheduler_policy: MAX_UTILIZATION
cuda_graph_config:
  enable_padding: true
  max_batch_size: 128
kv_cache_config:
  enable_block_reuse: false
  free_gpu_memory_fraction: 0.7
  dtype: fp8
  mamba_ssm_cache_dtype: bfloat16
  use_kv_cache_manager_v2: false
speculative_config:
  decoding_type: DFlash
  use_rejection_sampling: false
  max_draft_len: 7
  speculative_model: ${DRAFT_MODEL}
  attention_backend: TRTLLM
moe_config:
  backend: TRTLLM
num_postprocess_workers: 4
stream_interval: 20
return_perf_metrics: true
enable_iter_perf_stats: true
YAML

export TRTLLM_USE_GDN_REPLAY=1
export TRTLLM_WINDOW_SIZE_SHARES=0.70,0.30

trtllm-serve "${BASE_MODEL}" \
  --backend pytorch \
  --tp_size 1 \
  --ep_size 1 \
  --port 8001 \
  --host 0.0.0.0 \
  --max_batch_size 128 \
  --max_num_tokens 8192 \
  --max_seq_len 16384 \
  --extra_llm_api_options "${CONFIG}"

Benchmark results

Speed-Bench 2K throughput

Speed-Bench 2K throughput low-entropy dataset, 512 requests, average ISL 2210, requested OSL unset, and greedy decoding (top_k=1). MTP uses draft length 3 and DFlash uses draft length 7. Generation stops at EOS; max_tokens=6144 is only the safety cap.

Concurrency	MTP=3	DFlash=7	DFlash=7 (attention_backend=TRTLLM)
1	922.1	1,127.9	1,212.3
2	820.5	953.9	1,051.6
4	704.9	836.3	913.9
8	592.2	684.5	746.6
16	473.7	520.1	581.5
32	362.6	359.8	404.1
64	256.2	239.0	278.7
128	184.5	154.0	189.0
256	134.7	—¹	—¹

Values are measured OTPS/user (output tokens/s/user).

¹ DFlash at concurrency 256 exceeds the available KV-cache capacity, so no measurement is reported.

Another run with TopK=0, TopP=1, Temp=0.7 (disable TopK, TopP).

Benchmark configuration notes:

• The DFlash high-concurrency run used TRTLLM_WINDOW_SIZE_SHARES=0.7,0.3 to distribute the KV-cache budget between the sliding-window and full-attention pools.
• The scheduler used MAX_UTILIZATION to admit more concurrent requests.

MT-Bench acceptance length by category

DFlash draft length 7 (maximum acceptance length 8), greedy decoding (temperature=0), TP1/EP1, B200. Each concurrency run contains 80 MT-Bench requests. The measured average actual OSL across all seven runs was 2,377.1 tokens per request; the per-run means ranged from 2,315.8 to 2,442.6 tokens, while 4,096 was the configured maximum output length.

Category	c=1	c=2	c=4	c=8	c=16	c=32	c=64
Writing	4.162	4.208	4.124	4.253	4.194	4.274	4.211
Roleplay	3.458	3.371	3.406	3.474	3.314	3.512	3.285
Reasoning	4.190	4.139	4.143	4.186	4.135	4.224	4.163
Math	5.504	5.527	5.432	5.370	5.376	5.453	5.455
Coding	4.759	4.702	4.732	4.703	4.749	4.624	4.648
Extraction	5.421	5.358	5.395	5.348	5.362	5.452	5.304
STEM	3.500	3.596	3.563	3.593	3.575	3.579	3.608
Humanities	3.144	3.192	3.203	3.191	3.161	3.204	3.199

Overall average acceptance length remains stable across concurrency: c1 4.267, c2 4.262, c4 4.250, c8 4.265, c16 4.233, c32 4.290, and c64 4.234.

These are greedy-decoding results. MT-Bench shows no systematic acceptance-length drop as concurrency increases: c64 remains at 4.234 versus 4.267 at c1, and the intermediate measurements are similarly stable.

Review follow-up

• Preserve the original target max_seq_len before target SWA construction can clamp it, then reuse that value for the budget-split decision, draft-window derivation, and final draft-manager construction.
• Use one per-layer attention-window derivation for DFlash attention masks and draft KV-cache layout, including explicit full-attention semantics and legacy fallback behavior. Laguna follows its uniform checkpoint window.
• Propagate the derived draft KV-cache configuration through both one-model token-cost paths so external DFlash and EAGLE3/MTP estimates do not inherit the target window and overallocate memory.
• Keep VSWA classification separate from V2 estimation pool counting: recurrent-state sentinels remain non-VSWA, while hybrid linear/attention targets retain both estimation pools and unsupported target-window metadata falls back safely.
• Harden TRTLLM-Gen workspace and counter-buffer management with device checks, capacity growth, CUDA-graph capture guards, cached SM count, and int32 cumulative sequence offsets.
• Document the FlashInfer and SM100/SM103 requirements for the optional TRTLLM DFlash backend.

Refactoring

• Added get_layer_attention_window() as the canonical interpreter for use_sliding_window, sliding_window, and layer_types:
  • use_sliding_window=False explicitly disables the window even when a scalar window remains in the HF config.
  • use_sliding_window=True requires a positive integer sliding_window and fails early when it is absent or invalid.
  • An unset flag preserves legacy inference from sliding_window and layer_types.
  • Repeating layer_types patterns are supported through modulo indexing; full-attention layers return None.
  • Qwen2-style configs without layer_types honor max_window_layers as the first sliding-layer index.
  • Multiple per-layer sliding-window sizes remain unsupported and fail explicitly instead of being silently collapsed.
• Relocated the existing _is_sliding_attention_layer() predicate to config_utils.py without changing its private naming or matching semantics, allowing model and executor code to share it.
• Added draft_config_defines_attention_layout() with explicit documentation of ownership semantics:
  • When the draft config provides use_sliding_window, sliding_window, or layer_types, the draft layout is authoritative, including explicit full attention.
  • When the draft provides none of those fields, the legacy uniform-target fallback remains available.
• Updated _normalize_attention_windows() to materialize full-attention entries as max_seq_len, collapse all-full layouts to None, collapse uniform layouts to one entry, and retain heterogeneous layouts for separate pools.
• Moved uses_vswa_kv_cache_layout() to config_utils.py, removed the obsolete is_vswa_enabled() copy, and reused it in V1, V2, and KV-connector guards. V2 estimation now counts actual pools separately, including hybrid linear/attention layouts.
• Made _get_draft_max_attention_window() require an explicit KvCacheConfig instead of implicitly reading the target configuration.
• Added _get_one_model_draft_kv_cache_config() to clone the input config with the derived draft layout without mutating the target config.
• Reused the cloned draft config for external DFlash and EAGLE3/MTP token-cost calculations as well as final manager construction.
• Kept the temporary estimation manager uniform until the measured GPU budget is available; the final VSWA layout is installed only after target/draft budget splitting.
• Renamed worker-side backend state to _dflash_attention_backend and the model-side comparison value to draft_model_dflash_attention_backend to distinguish them from unrelated attention or MoE backends.
• Annotated every fallback cache-allocation/update branch as the VANILLA DFlash FlashAttention path.
• Converted context-cache slot and position indices to int32 once before the per-layer paged-append loop and hoisted the loop-invariant flattened page table, avoiding repeated conversion/allocation work.
• Added TRTLLM-Gen buffer device tracking, resize-on-growth behavior, capture-time allocation errors, and cached multiprocessor count.
• Changed TRTLLM-Gen cumulative sequence offsets to be produced directly as int32.
• Updated the Gemma 4 hybrid KV-estimation fixtures to use the real Gemma 4 E2B layout (28 sliding + 7 full layers) and its actual 512-token sliding window.
• Clarified the public DFlash backend documentation and configuration description: VANILLA uses FlashAttention with contiguous context K/V; TRTLLM requires FlashInfer and SM100/SM103 and uses a private paged context cache.

Test Coverage

• tests/unittest/_torch/modeling/test_modeling_speculative.py
  • Mixed sliding/full layer types with uniform effective RoPE are accepted.
  • Different effective RoPE configurations are rejected for both supported RoPE-parameter sources.
  • Per-layer causal/window arguments cover repeated mixed layouts, explicit SWA disable, invalid enabled-without-window configuration, and Laguna causal uniform-window behavior from the real checkpoint shape.
  • TRTLLM-Gen workspace/counter buffers are reused, grow when required, cache device properties, and reject allocation or device changes during CUDA graph capture.
• tests/unittest/_torch/executor/test_kv_cache_budget_split.py
  • Derived draft KV configuration is used by both external-DFlash and EAGLE3/MTP one-model token-cost paths without mutating the target config.
  • build_managers() preserves the original max_seq_len and constructs the final mixed draft layout from it.
  • VSWA detection excludes recurrent-state sentinels.
  • Existing proportional target/draft GPU-budget split behavior remains covered.
• tests/unittest/_torch/executor/test_kv_cache_estimation.py
  • Gemma 4 E2B mixed-window pool scaling uses the model's actual sliding_window=512.
  • Pool scaling retains sufficient per-pool capacity for the MMMU Pro workload shape.
  • Hybrid linear-attention targets retain recurrent and attention estimation pools for explicit mixed-layer layouts.
  • Unsupported target-window metadata falls back without failing startup, and Qwen2 max_window_layers is covered by a focused helper test.
• tests/unittest/_torch/speculative/hw_agnostic/test_dflash_worker.py
  • TRTLLM-Gen constructs the private HND paged context cache with the expected page allocation and leaves VANILLA contiguous buffers unused.
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
  • Qwen3.6-35B-A3B NVFP4 + DFlash runs on SM100/SM103 with FP8 KV cache, GDN replay, and both VANILLA and TRTLLM DFlash attention backends.
• tests/integration/defs/accuracy/references/gsm8k.yaml
  • DFlash GSM8K accuracy reference remains 95.0.
• Registered the accuracy coverage in llm_function_core.txt and l0_b200.yml.
• Removed low-value implementation-detail tests that duplicated backend enum validation, direct model-to-worker backend pass-through/mismatch wiring, and a contrived EAGLE3 explicit-SWA-disable configuration; retained coverage focuses on distinct behavioral contracts.
• Local validation after pruning to focused reviewer coverage: 5 passed; the full repository pre-commit suite, including compiled-binding mypy, passed.

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why.

• PR follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions).

• If the PR introduces API changes, an appropriate api-compatible or api-breaking label is added.

• New dependencies, if any, have been scanned for license and vulnerabilities.

• CODEOWNERS is updated if ownership changes.

• Documentation and the TAVA architecture diagram are updated if needed.

• The assigned reviewers are appropriate for the PR.

• I reviewed the checklist above as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, comment /bot help.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

Walkthrough

DFlash now validates uniform effective RoPE settings, derives per-layer FlashAttention masking windows, and passes them to attention. Draft KV-cache window derivation and V1 budget splitting are updated, with unit and Qwen3.6 NVFP4 GSM8K coverage added.

Changes

DFlash attention behavior

Layer / File(s)	Summary
RoPE validation and attention masking tensorrt_llm/_torch/models/modeling_speculative.py, tests/unittest/_torch/modeling/test_modeling_speculative.py	DFlash validates effective RoPE uniformity, derives causal and sliding-window arguments per layer, and passes the computed window to FlashAttention with unit coverage.
Draft KV-cache windows and budget splitting tensorrt_llm/_torch/pyexecutor/_util.py, tensorrt_llm/_torch/pyexecutor/resource_manager.py, tests/unittest/_torch/executor/test_kv_cache_budget_split.py	VSWA detection is centralized, draft attention windows are derived for estimation and final construction, and V1 budget splitting accounts for heterogeneous draft windows and the original sequence length.
DFlash accuracy coverage tests/integration/defs/accuracy/test_llm_api_pytorch.py, tests/integration/defs/accuracy/references/gsm8k.yaml, tests/integration/test_lists/qa/llm_function_core.txt, tests/integration/test_lists/test-db/l0_b200.yml	Adds the Qwen3.6 NVFP4 DFlash GSM8K test, its reference result, and test-list registrations.

Estimated code review effort: 4 (Complex) | ~45 minutes

Suggested reviewers: vallis-neria, thorjohnsen

Sequence Diagram(s)

sequenceDiagram
  participant DFlashForCausalLM
  participant get_attention_mask_args
  participant flash_attn_with_kvcache
  DFlashForCausalLM->>get_attention_mask_args: request arguments for layer_idx
  get_attention_mask_args-->>DFlashForCausalLM: return causal and window_size
  DFlashForCausalLM->>flash_attn_with_kvcache: pass causal and window_size

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 37.50% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly summarizes the PR's main changes: DFlash RoPE, sliding-window support, and the TRTLLM-Gen attention backend.
Description check	✅ Passed	The description explains the motivation, implementation, tests, benchmarks, refactoring, and checklist status in substantial detail.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (2)

tests/unittest/_torch/modeling/test_modeling_speculative.py (1)

16-25: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Test coverage summary.

Added tests: test_dflash_allows_mixed_layer_types_with_uniform_rope, test_dflash_rejects_different_effective_rope (both parametrized over rotary_emb/pos_embd_params sources), and test_dflash_attention_mask_args, exercising DFlashForCausalLM._validate_uniform_rope and _get_attention_mask_args. These are unit tests (not integration tests), so they aren't expected to appear in tests/integration/test_lists/**.

Coverage verdict: sufficient for the primary contract, but the following new error-raising branches in _get_attention_mask_args are untested: layer_types length mismatch vs num_hidden_layers, an unsupported layer_type value, and use_sliding_window=True with a missing/non-positive sliding_window. Consider adding a couple of pytest.raises cases for these branches for completeness.

As per path instructions for tests/**: "Always produce a test coverage summary, even if no issues are found."

Also applies to: 150-234

🤖 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 `@tests/unittest/_torch/modeling/test_modeling_speculative.py` around lines 16
- 25, Add pytest.raises coverage for the validation branches in
DFlashForCausalLM._get_attention_mask_args: layer_types length differing from
num_hidden_layers, an unsupported layer_type, and use_sliding_window=True with a
missing or non-positive sliding_window. Extend the existing
test_dflash_attention_mask_args cases while preserving the current valid-path
assertions.

Source: Path instructions

tensorrt_llm/_torch/models/modeling_speculative.py (1)

1345-1379: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Consider caching per-layer mask args instead of recomputing every forward call.

_get_attention_mask_args re-derives (causal, window_size) from config.layer_types/sliding_window on every call, but these values are static per layer once the config is loaded. Since dflash_forward calls this once per layer on every decode step (a hot path), precomputing a list of (causal, window_size) tuples once (e.g. alongside _build_fused_kv_buffers) would avoid repeated attribute lookups and re-validation.

🤖 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/models/modeling_speculative.py` around lines 1345 - 1379,
Cache the per-layer (causal, window_size) results during initialization,
alongside _build_fused_kv_buffers, after validating layer_types and
sliding-window configuration. Update _get_attention_mask_args to return the
cached tuple for layer_idx without re-reading or re-validating config values,
and ensure the cache preserves the existing defaults and ValueError behavior.

🤖 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.

Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_speculative.py`:
- Around line 1345-1379: Cache the per-layer (causal, window_size) results
during initialization, alongside _build_fused_kv_buffers, after validating
layer_types and sliding-window configuration. Update _get_attention_mask_args to
return the cached tuple for layer_idx without re-reading or re-validating config
values, and ensure the cache preserves the existing defaults and ValueError
behavior.

In `@tests/unittest/_torch/modeling/test_modeling_speculative.py`:
- Around line 16-25: Add pytest.raises coverage for the validation branches in
DFlashForCausalLM._get_attention_mask_args: layer_types length differing from
num_hidden_layers, an unsupported layer_type, and use_sliding_window=True with a
missing or non-positive sliding_window. Extend the existing
test_dflash_attention_mask_args cases while preserving the current valid-path
assertions.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6bcbaa68-7fba-463e-9846-63e910c6d84b

📥 Commits

Reviewing files that changed from the base of the PR and between cfeca00 and 86c84bd.

📒 Files selected for processing (6)

• tensorrt_llm/_torch/models/modeling_speculative.py
• tests/integration/defs/accuracy/references/gsm8k.yaml
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
• tests/integration/test_lists/qa/llm_function_core.txt
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/unittest/_torch/modeling/test_modeling_speculative.py

[coderabbitai]

Actionable comments posted: 1

🧹 Nitpick comments (1)

tests/unittest/_torch/executor/test_kv_cache_budget_split.py (1)

68-68: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Annotate the new test method.

Add -> None to test_v1_separate_vswa_draft_requires_gpu_budget_split.

As per coding guidelines, “Annotate every function.”

🤖 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 `@tests/unittest/_torch/executor/test_kv_cache_budget_split.py` at line 68, Add
the explicit return annotation -> None to the
test_v1_separate_vswa_draft_requires_gpu_budget_split test method, preserving
its existing behavior and body.

Source: Coding guidelines

🤖 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/pyexecutor/_util.py`:
- Around line 1085-1098: Preserve the pre-target-manager sequence length in
build_managers before target SWA construction mutates self._max_seq_len, and
reuse it for both draft-window derivation sites, including
_get_draft_max_attention_window and the corresponding budget-split calculation.
Add a regression test covering target SWA with mixed draft layers such as
sliding followed by full attention, verifying the decision and final windows use
the original length.

---

Nitpick comments:
In `@tests/unittest/_torch/executor/test_kv_cache_budget_split.py`:
- Line 68: Add the explicit return annotation -> None to the
test_v1_separate_vswa_draft_requires_gpu_budget_split test method, preserving
its existing behavior and body.

🪄 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: 8b5baa98-3a9c-4578-974c-6020ea300f12

📥 Commits

Reviewing files that changed from the base of the PR and between 86c84bd and 282fc22.

📒 Files selected for processing (9)

• tensorrt_llm/_torch/models/modeling_speculative.py
• tensorrt_llm/_torch/pyexecutor/_util.py
• tests/integration/defs/accuracy/references/gsm8k.yaml
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
• tests/integration/test_lists/qa/llm_function_core.txt
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/unittest/_torch/executor/test_kv_cache_budget_split.py
• tests/unittest/_torch/modeling/test_modeling_speculative.py
• tests/unittest/_torch/speculative/test_eagle3.py

🚧 Files skipped from review as they are similar to previous changes (6)

• tests/integration/test_lists/test-db/l0_b200.yml
• tests/integration/defs/accuracy/references/gsm8k.yaml
• tests/integration/test_lists/qa/llm_function_core.txt
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
• tests/unittest/_torch/modeling/test_modeling_speculative.py
• tensorrt_llm/_torch/models/modeling_speculative.py

[coderabbitai]

Actionable comments posted: 1

🤖 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 `@tests/unittest/_torch/executor/test_kv_cache_budget_split.py`:
- Around line 69-89: Update
test_v1_mixed_draft_uses_original_max_seq_len_for_budget_split to exercise
build_managers() rather than only calling _needs_gpu_kv_cache_budget_split() and
_get_draft_max_attention_window() directly. Mock the required manager-building
dependencies and verify the build path passes original_max_seq_len (16384) to
the relevant helper, while preserving coverage of the mixed V1 draft
configuration.

🪄 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: c02e6e9c-6f90-4239-a447-759e5c84c452

📥 Commits

Reviewing files that changed from the base of the PR and between 282fc22 and 54f1ced.

📒 Files selected for processing (8)

• tensorrt_llm/_torch/models/modeling_speculative.py
• tensorrt_llm/_torch/pyexecutor/_util.py
• tests/integration/defs/accuracy/references/gsm8k.yaml
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
• tests/integration/test_lists/qa/llm_function_core.txt
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/unittest/_torch/executor/test_kv_cache_budget_split.py
• tests/unittest/_torch/modeling/test_modeling_speculative.py

🚧 Files skipped from review as they are similar to previous changes (6)

• tests/integration/test_lists/qa/llm_function_core.txt
• tests/integration/defs/accuracy/references/gsm8k.yaml
• tensorrt_llm/_torch/models/modeling_speculative.py
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
• tests/unittest/_torch/modeling/test_modeling_speculative.py
• tensorrt_llm/_torch/pyexecutor/_util.py

[coderabbitai]

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/models/modeling_speculative.py (1)

680-701: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win

Unused residual from single-layer call.

Ruff flags the residual unpacked at Line 692 as unused. Since self.layers contains exactly one layer here, discard it explicitly.

🧹 Proposed fix

-            hidden_states, residual = self.layers[0](
+            hidden_states, _ = self.layers[0](

🤖 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/models/modeling_speculative.py` around lines 680 - 701,
Update the single-layer call in the model forward path to explicitly discard the
unused residual value instead of assigning it to residual. Preserve the existing
hidden_states result and cleanup behavior in the surrounding try/finally block.

Source: Linters/SAST tools

🤖 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.

Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_speculative.py`:
- Around line 680-701: Update the single-layer call in the model forward path to
explicitly discard the unused residual value instead of assigning it to
residual. Preserve the existing hidden_states result and cleanup behavior in the
surrounding try/finally block.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: c503fd15-4869-4073-81db-3ad7d90540ae

📥 Commits

Reviewing files that changed from the base of the PR and between 54f1ced and cc69d53.

📒 Files selected for processing (8)

• tensorrt_llm/_torch/models/modeling_speculative.py
• tensorrt_llm/_torch/pyexecutor/_util.py
• tests/integration/defs/accuracy/references/gsm8k.yaml
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
• tests/integration/test_lists/qa/llm_function_core.txt
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/unittest/_torch/executor/test_kv_cache_budget_split.py
• tests/unittest/_torch/modeling/test_modeling_speculative.py

🚧 Files skipped from review as they are similar to previous changes (7)

• tests/integration/test_lists/qa/llm_function_core.txt
• tests/integration/test_lists/test-db/l0_b200.yml
• tests/unittest/_torch/executor/test_kv_cache_budget_split.py
• tests/integration/defs/accuracy/references/gsm8k.yaml
• tests/integration/defs/accuracy/test_llm_api_pytorch.py
• tests/unittest/_torch/modeling/test_modeling_speculative.py
• tensorrt_llm/_torch/pyexecutor/_util.py

[tensorrt-cicd]

PR_Github #66338 [ run ] triggered by Bot. Commit: 9f67546 Link to invocation

[tensorrt-cicd]

PR_Github #66338 [ run ] completed with state FAILURE. Commit: 9f67546
/LLM/main/L0_MergeRequest_PR pipeline #53983 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[jhaotingc]

/bot run

[tensorrt-cicd]

PR_Github #66415 [ run ] triggered by Bot. Commit: 2bb56d5 Link to invocation

[tensorrt-cicd]

PR_Github #66415 [ run ] completed with state FAILURE. Commit: 2bb56d5
/LLM/main/L0_MergeRequest_PR pipeline #54057 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[jhaotingc]

/bot run

[tensorrt-cicd]

PR_Github #66445 [ run ] triggered by Bot. Commit: 2bb56d5 Link to invocation

[tensorrt-cicd]

PR_Github #66445 [ run ] completed with state FAILURE. Commit: 2bb56d5
/LLM/main/L0_MergeRequest_PR pipeline #54087 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[jhaotingc]

/bot run

[tensorrt-cicd]

PR_Github #66507 [ ] completed with state FAILURE. Commit: ``

Link to invocation

[jhaotingc]

/bot run

[tensorrt-cicd]

PR_Github #66510 [ run ] triggered by Bot. Commit: 2bb56d5 Link to invocation

[tensorrt-cicd]

PR_Github #66510 [ run ] completed with state FAILURE. Commit: 2bb56d5
/LLM/main/L0_MergeRequest_PR pipeline #54148 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[jhaotingc]

/bot run

[jhaotingc]

/bot run

[tensorrt-cicd]

PR_Github #66517 [ run ] triggered by Bot. Commit: 2bb56d5 Link to invocation

[jhaotingc]

/bot run

[tensorrt-cicd]

PR_Github #66590 [ run ] triggered by Bot. Commit: 2bb56d5 Link to invocation

[tensorrt-cicd]

PR_Github #66517 [ run ] completed with state ABORTED. Commit: 2bb56d5

Link to invocation

[tensorrt-cicd]

PR_Github #66590 [ run ] completed with state SUCCESS. Commit: 2bb56d5
/LLM/main/L0_MergeRequest_PR pipeline #54216 completed with status: 'SUCCESS'

CI Report

Link to invocation