[None][feat] Add Laguna DFlash drafter support

[joerowell]

Description

Adds DFlash speculative-decoding support for poolside Laguna drafters (XS.2 / XS-2.1 / M.1): DFlashForCausalLM (generic block decode — target-hidden projection, input_layernorm-folded fused KV, per-layer causal, output-gate hook) and DFlashLagunaForCausalLM (per-head softplus gating, per-aux fc_norm). get_draft_model() selects the Laguna drafter from the draft config's architectures — no new flag.

Test Coverage

Validated on poolside/Laguna-XS-2.1 (tensorrt-llm 1.3.0rc20): loads with trust_remote_code=True (no overlay, transformers 5.5.4) and generates; BFCL via trtllm-serve --tool_parser poolside_v1 --reasoning_parser laguna = 94.0% / 96.5% (simple_python / multiple), on par with the vLLM / SGLang / llama.cpp references.

PR Checklist

• PR description clearly explains what and why.
• Follows the TRT-LLM coding guidelines.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds Laguna-specific DFlash speculative decoding support, introduces a DFlash architecture selector, and updates hidden-state projection, fused KV preparation, and attention gating paths.

Changes

DFlash Laguna speculative decoding

Layer / File(s)	Summary
Projection hook and architecture flag tensorrt_llm/llmapi/llm_args.py, tensorrt_llm/_torch/models/modeling_speculative.py, tensorrt_llm/_torch/speculative/dflash.py	DFlashDecodingConfig adds dflash_model_arch; DFlashForCausalLM adds project_target_hidden(); speculative prefill and generation paths call that hook for hidden-state projection.
Fused KV preparation tensorrt_llm/_torch/models/modeling_speculative.py	load_weights() splits fused self_attn.qkv_proj.weight; _build_fused_kv_buffers() folds input_layernorm.weight into KV weights and records _input_ln_eps; precompute_context_kv() applies the matching normalization.
Attention gating and causality tensorrt_llm/_torch/models/modeling_speculative.py	dflash_forward() derives a per-layer causal flag from config.layer_types, passes it to flash_attn_with_kvcache, and routes attention output through _post_attention_gate() before o_proj.
Laguna drafter selection tensorrt_llm/_torch/models/modeling_speculative.py	DFlashLagunaForCausalLM validates Laguna-specific gating, loads aux_hidden_norms.*, applies per-aux projection norms, and get_draft_model() returns it when dflash_model_arch is "laguna".

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 52.94% 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 main change: adding Laguna DFlash drafter support.
Description check	✅ Passed	The PR description includes Description, Test Coverage, and PR Checklist sections and clearly explains the Laguna DFlash changes.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

Actionable comments posted: 3

🤖 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/models/modeling_speculative.py`:
- Around line 1517-1531: The Laguna aux-norm loading in model loading should
fail fast when expected per-aux norms are missing instead of silently falling
back to a generic path. Update the weight handling in
DFlashLagunaForCausalLM.load_weights to detect the Laguna case and verify
aux_hidden_norms.* are present before building self.fc_norm; if none are found,
raise an error rather than proceeding. Keep the existing aux_keys sorting and
ModuleList construction, but make the incomplete-check explicit so
project_target_hidden() cannot run with a partially loaded Laguna checkpoint.
- Around line 981-982: The fused-QKV setup in modeling_speculative uses
self.config.head_dim directly, which can fail for pretrained configs that do not
define it. Update the q and kv dimension calculation in the relevant method to
derive head_dim using the same fallback logic used elsewhere in the model code,
so the speculative loading path works even when head_dim is missing. Use the
existing configuration access pattern in the surrounding class/method to keep
the behavior consistent.
- Around line 1183-1187: The KV folding in modeling_speculative.py’s speculative
context path is using only input_layernorm.weight, but Gemma-style RMSNorm
requires folding the actual scale used by the forward path, namely the effective
weight plus the base offset. Update the folding logic in the block guarded by
hasattr(dl, 'input_layernorm') so the precomputed KV weights match the query
path for use_gemma-style RMSNorm. While making that change, verify the single
_input_ln_eps assumption across all dlayers before storing _input_ln_eps, and
use the existing input_layernorm/config symbols to keep the behavior consistent
with the rest of the model code.

🪄 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: cec08356-5c36-4f48-893a-e0d8f8a72da0

📥 Commits

Reviewing files that changed from the base of the PR and between 64ac565 and b338758.

📒 Files selected for processing (3)

• tensorrt_llm/_torch/models/modeling_speculative.py
• tensorrt_llm/_torch/speculative/dflash.py
• tensorrt_llm/llmapi/llm_args.py

[DomBrown]

/bot run

[tensorrt-cicd]

PR_Github #56121 [ run ] triggered by Bot. Commit: f5c0284 Link to invocation

[tensorrt-cicd]

PR_Github #56121 [ run ] completed with state FAILURE. Commit: f5c0284
/LLM/main/L0_MergeRequest_PR pipeline #44986 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

[DomBrown]

/bot run

[tensorrt-cicd]

PR_Github #56174 [ run ] triggered by Bot. Commit: 95c2442 Link to invocation

[tensorrt-cicd]

PR_Github #56174 [ run ] completed with state SUCCESS. Commit: 95c2442
/LLM/main/L0_MergeRequest_PR pipeline #45037 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

[DomBrown]

/bot run

[tensorrt-cicd]

PR_Github #56192 [ run ] triggered by Bot. Commit: 95c2442 Link to invocation

[tensorrt-cicd]

PR_Github #56192 [ run ] completed with state SUCCESS. Commit: 95c2442
/LLM/main/L0_MergeRequest_PR pipeline #45055 completed with status: 'SUCCESS'

CI Report

Link to invocation

[QiJune]

/bot run

[tensorrt-cicd]

PR_Github #59152 [ run ] triggered by Bot. Commit: ac9980f Link to invocation

[tensorrt-cicd]

PR_Github #59152 [ run ] completed with state FAILURE. Commit: ac9980f
/LLM/main/L0_MergeRequest_PR pipeline #47660 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

[DomBrown]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59160 [ run ] triggered by Bot. Commit: ac9980f Link to invocation

[DomBrown]

/bot kill

[tensorrt-cicd]

PR_Github #59191 [ kill ] triggered by Bot. Commit: f844ae3 Link to invocation

[tensorrt-cicd]

PR_Github #59160 [ run ] completed with state ABORTED. Commit: ac9980f

Link to invocation

[tensorrt-cicd]

PR_Github #59191 [ kill ] completed with state SUCCESS. Commit: f844ae3
Successfully killed previous jobs for commit f844ae3

Link to invocation

[DomBrown]

/bot run

[tensorrt-cicd]

PR_Github #59193 [ run ] triggered by Bot. Commit: c5b95cc Link to invocation

[tensorrt-cicd]

PR_Github #59193 [ run ] completed with state FAILURE. Commit: c5b95cc
/LLM/main/L0_MergeRequest_PR pipeline #47692 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

[DomBrown]

/bot run

[tensorrt-cicd]

PR_Github #59196 [ run ] triggered by Bot. Commit: c5b95cc Link to invocation

[tensorrt-cicd]

PR_Github #59196 [ run ] completed with state SUCCESS. Commit: c5b95cc
/LLM/main/L0_MergeRequest_PR pipeline #47694 completed with status: 'SUCCESS'

CI Report

Link to invocation