Repository navigation
[None][feat] Add Laguna model support (Poolside Laguna-XS.2) - #13559
Merged
Merged
Conversation
nekorobov
approved these changes
Apr 29, 2026
Contributor
📝 WalkthroughWalkthroughThis pull request introduces Laguna model support to TensorRT-LLM's PyTorch backend with MoE routing, specialized attention, and FP8 block-scale quantization. It adds llm-compressor NVFP4 and KV-cache quantization parsing, regex-based module exclusion matching, and a checkpoint overlay utility for configuration compatibility. Changes
Sequence Diagram(s)sequenceDiagram
participant User as User
participant Loader as ModelLoader
participant HF as HF Checkpoint
participant Mapper as LagunaHfWeightMapper
participant Model as LagunaModel
participant MoE as LagunaMoE
User->>Loader: load_hf_quant_config(checkpoint_path)
Loader->>HF: Read quantization_config
HF-->>Loader: NVFP4 weights config + KV-cache scheme
Loader->>Loader: Detect llm-compressor NVFP4 (num_bits=4, type="float", group_size=16)
Loader->>Loader: Parse KV-cache FP8 scheme
Loader-->>User: QuantAlgo.NVFP4 + QuantAlgo.FP8 (KV)
User->>Loader: load_weights(checkpoint_dir, weight_mapper)
Loader->>HF: Read packed weights & scales
Loader->>Mapper: preprocess_weights(weights)
Mapper->>Mapper: Translate NVFP4 packed to (weight, weight_scale_2, input_scale)
Mapper->>Mapper: Invert *_global_scale with zero/uncalibrated safeguards
Mapper-->>Loader: Remapped weights
Loader->>Model: Initialize LagunaModel
Loader->>Model: load_weights(remapped_weights)
Model->>Model: Instantiate decoder layers with MoE blocks
Model->>MoE: Register gate bias tensors
MoE-->>Model: MoE initialized
Model-->>Loader: Model ready
Loader-->>User: LagunaForCausalLM instance
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Tip 💬 Introducing Slack Agent: The best way for teams to turn conversations into code.Slack Agent is built on CodeRabbit's deep understanding of your code, so your team can collaborate across the entire SDLC without losing context.
Built for teams:
One agent for your entire SDLC. Right inside Slack. 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: 5
🧹 Nitpick comments (3)
tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py (2)
991-1000: Prefix unused variable with underscore to silence linter.
num_expertsis unpacked but never used in the BF16 fallback function.🔧 Suggested fix
def _run_moe_bf16_fallback( ... ) -> torch.Tensor: """BF16 path: dequant FP8 weights on-the-fly, use BF16 activations.""" from tensorrt_llm._torch.modules.fused_moe.interface import ActivationType num_tokens, hidden = x.shape - num_experts, N_gate_up, _ = w3_w1.shape + _num_experts, N_gate_up, _ = w3_w1.shape intermediate = N_gate_up // 2🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py` around lines 991 - 1000, The variable num_experts is unpacked from w3_w1.shape but never used in the BF16 fallback; to silence the linter, rename the unused unpacked variable to start with an underscore (for example _num_experts) wherever w3_w1 is unpacked (the line that currently reads "num_experts, N_gate_up, _ = w3_w1.shape") in the fused_moe_triton_fp8_block_scale code so the linter recognizes it as intentionally unused.
1042-1076: Consider guarding against empty input tensor.If
xhas zero tokens (num_tokens == 0), the function will proceed with routing and GEMM calls on empty tensors. While this may work correctly, an early return could improve clarity and avoid unnecessary kernel launches.🔧 Optional early-exit for empty input
def run_triton_fp8_block_scale_moe( ... ) -> torch.Tensor: ... num_tokens, hidden = x.shape + if num_tokens == 0: + dtype = output_dtype or torch.bfloat16 + return torch.empty((0, hidden), dtype=dtype, device=x.device) num_experts, N_gate_up, _ = w3_w1.shape🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py` around lines 1042 - 1076, Add an early-exit at the start of run_triton_fp8_block_scale_moe to guard against empty input: if x.size(0) == 0 (no tokens) return an appropriately shaped empty tensor (shape (0, x.size(1))) with the expected output dtype (respecting output_dtype if provided, otherwise default BF16) to avoid running routing/GEMM logic (token_selected_experts, token_final_scales, Triton kernels) on empty tensors.tensorrt_llm/_torch/models/modeling_laguna.py (1)
54-101: Fix type annotation and remove unused attribute.Two issues in
LagunaGate:
- Line 63:
moe_backend_cls: Type[MoE] = Noneshould use explicit| Nonesyntax per PEP 484 and coding guidelines.- Line 67:
self.moe_backend_clsis stored but never referenced anywhere in the class — this appears to be dead code.🔧 Suggested fix
def __init__( self, hidden_size: int, num_experts: int, top_k: int, dtype: Optional[torch.dtype] = None, - moe_backend_cls: Type[MoE] = None, + moe_backend_cls: Type[MoE] | None = None, ): super().__init__() self.top_k = top_k - self.moe_backend_cls = moe_backend_cls + # moe_backend_cls reserved for future use (e.g., backend-specific routing) self.weight = nn.Parameter(As per coding guidelines: "Use the
|syntax instead oftyping.Unionin Python type annotations"🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/models/modeling_laguna.py` around lines 54 - 101, In class LagunaGate update the __init__ signature to use the PEP 484/PEP 604 union syntax for the optional backend (change moe_backend_cls: Type[MoE] = None to moe_backend_cls: Type[MoE] | None = None) and remove the unused instance attribute assignment self.moe_backend_cls (delete that line in __init__), leaving the parameter if callers pass it but not storing it on the object; keep references to other attributes (weight, e_score_correction_bias, out_dtype) unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@laguna_minimal_overlay.sh`:
- Around line 25-37: The script currently does an unconditional rm -rf on the
OVERLAY path which can delete caller-provided data; change this by (a) if
OVERLAY was explicitly passed (i.e., second positional arg provided) then refuse
to remove an existing directory and instead error if it already exists or is
non-empty, and (b) if OVERLAY was not provided (using the default), create a
truly unique temp directory with mktemp -d and use that; replace the
unconditional rm -rf "$OVERLAY" and mkdir -p "$OVERLAY" with logic that checks
whether OVERLAY was user-supplied, aborts on existing/non-empty dir for
user-supplied OVERLAY, and only creates a new directory (via mktemp -d when
defaulting) to avoid wiping unrelated files.
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py`:
- Around line 614-616: The SM120 FP8 block-scale fallback unconditionally
computes local_scales = token_final_scales *
is_local.to(token_final_scales.dtype) even though forward_chunk() sets
token_final_scales = None when apply_router_weight_on_input=True; guard this
path by checking token_final_scales before using it (e.g., only compute
local_scales and call run_triton_fp8_block_scale_moe when token_final_scales is
not None and apply_router_weight_on_input is False or equivalent), otherwise
skip the SM120 fallback branch or take the alternative non-block-scale path to
avoid dereferencing None in fused_moe_cutlass.forward_chunk / the SM120 fallback
code that invokes run_triton_fp8_block_scale_moe.
In `@tensorrt_llm/llmapi/llm_utils.py`:
- Around line 525-526: The branch handling hf_quant_config when
hf_quant_config.get("quant_method") is None should explicitly indicate no
quantization: inside _update_from_hf_quant_config(), replace the no-op pass with
an explicit return False (or otherwise set the function's "updated" result to
False) so callers won't treat a null quant_method as a pre-quantized checkpoint;
reference hf_quant_config and the quant_method check in
_update_from_hf_quant_config to locate the change.
- Around line 507-524: The loader currently only assigns
hf_quant_config.get("ignore") to quant_config.exclude_modules causing modules
listed under "modules_to_not_convert" to be lost; update the block that sets
quant_config.exclude_modules to merge both hf_quant_config.get("ignore", []) and
hf_quant_config.get("modules_to_not_convert", []) (preserving existing
quant_config.exclude_modules if present), deduplicate entries, and assign the
combined list to quant_config.exclude_modules so behavior matches
tensorrt_llm/_torch/model_config.py; locate the logic around hf_quant_config and
quant_config.exclude_modules in llm_utils.py to implement this merge.
In `@tensorrt_llm/models/modeling_utils.py`:
- Around line 264-267: PretrainedConfig._get_quant_cfg currently only checks
exact matches and trailing '*' but must use the same exclude-pattern matcher as
the other code path; update QuantConfig._get_quant_cfg (and any usage in
PretrainedConfig._get_quant_cfg) to treat exclude_module entries starting with
"re:" by calling re.fullmatch(exclude_module[3:], name) and for other patterns
use fnmatch.fnmatchcase(name, exclude_module) so regex exclusions are honored;
locate the logic around variables named exclude_module and name and replace the
exact/everything-else matching with this shared matcher logic.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_laguna.py`:
- Around line 54-101: In class LagunaGate update the __init__ signature to use
the PEP 484/PEP 604 union syntax for the optional backend (change
moe_backend_cls: Type[MoE] = None to moe_backend_cls: Type[MoE] | None = None)
and remove the unused instance attribute assignment self.moe_backend_cls (delete
that line in __init__), leaving the parameter if callers pass it but not storing
it on the object; keep references to other attributes (weight,
e_score_correction_bias, out_dtype) unchanged.
In `@tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py`:
- Around line 991-1000: The variable num_experts is unpacked from w3_w1.shape
but never used in the BF16 fallback; to silence the linter, rename the unused
unpacked variable to start with an underscore (for example _num_experts)
wherever w3_w1 is unpacked (the line that currently reads "num_experts,
N_gate_up, _ = w3_w1.shape") in the fused_moe_triton_fp8_block_scale code so the
linter recognizes it as intentionally unused.
- Around line 1042-1076: Add an early-exit at the start of
run_triton_fp8_block_scale_moe to guard against empty input: if x.size(0) == 0
(no tokens) return an appropriately shaped empty tensor (shape (0, x.size(1)))
with the expected output dtype (respecting output_dtype if provided, otherwise
default BF16) to avoid running routing/GEMM logic (token_selected_experts,
token_final_scales, Triton kernels) on empty tensors.
🪄 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: 5898678f-ba19-42a0-8882-309c69db90fe
📒 Files selected for processing (9)
laguna_minimal_overlay.shtensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/modeling_laguna.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.pytensorrt_llm/llmapi/llm_utils.pytensorrt_llm/llmapi/reasoning_parser.pytensorrt_llm/models/modeling_utils.py
dcampora
force-pushed
the
user/domb/laguna-rebase
branch
from
May 4, 2026 08:11
01d8cc4 to
551a700
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #46835 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #46835 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #46867 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #46867 [ run ] completed with state
|
DomBrown
force-pushed
the
user/domb/laguna-rebase
branch
from
May 11, 2026 11:56
04ce491 to
e18f49d
Compare
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #47740 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #47740 [ run ] completed with state |
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #47971 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #47971 [ run ] completed with state |
Collaborator
|
PR_Github #49502 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49534 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49534 [ run ] completed with state
|
[None][feat] Add Laguna-XS TRT-LLM eval scripts and progress notes
Adds smoke test and GSM8K evaluation scripts for the Laguna-XS model
(out-of-tree, poolside checkpoint) using the TRT-LLM PyTorch backend.
Includes a progress.md with build notes and code review findings.
Results on RTX PRO 6000 Blackwell (TRT-LLM 1.3.0rc12):
- Smoke test: PASSED, 258.4 tok/s
- GSM8K 8-shot CoT: 93.33% (1231/1319)
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
[None][fix] SM120: dequantize FP8 block-scale MoE weights to BF16 at load time
The SM120 CUTLASS FP8 block-scale grouped-GEMM kernel fails with
TMA descriptor initialization errors (cudaError 700) on large batch
sizes (e.g. full GSM8K eval). The root cause is that cuTensorMapEncodeTiled
rejects the TMA descriptor parameters for large token counts, leaving a
sticky CUDA error that is caught by cudaGetLastError() in the kernel
launch path.
Since vLLM uses a standard Triton FP8 block-scale kernel for SM120
(TritonExperts) which has no equivalent in TRT-LLM, add a load-time
dequantization fallback instead:
- quantization.py: add DeepSeekFP8BlockScalesDequantFusedMoEMethod
Inherits from DeepSeekFP8BlockScalesFusedMoEMethod, applies parent
CUTLASS alignment padding, then dequantizes FP8 block-scale weights
to BF16 via block-scale expansion. Sets quant_scales=() and
_dequant_to_bf16=True on the module.
- fused_moe_cutlass.py: use new dequant method on SM120; gate
use_deepseek_fp8_block_scale on _dequant_to_bf16 flag; remove
TRTLLM_MOE_DEBUG debug prints.
- fused_moe_trtllm_gen.py: remove SM120 NotImplementedError guard.
- run_laguna_gsm8k.py / run_laguna_smoke_test.py: point to FP8
checkpoint and write hf_quant_config.json for FP8 auto-detection.
Validated on RTX PRO 6000 Blackwell (SM120, TRT-LLM 1.3.0rc12):
- Smoke test: PASSED, 262.5 tok/s
- GSM8K 8-shot CoT (1319 examples): 92.95% accuracy, 3517.5 tok/s
Signed-off-by: Daniel Campora <dcampora@nvidia.com>
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][fix] SM120: Triton BF16-act+FP8w MoE kernel replacing load-time dequant
The prior approach (DeepSeekFP8BlockScalesDequantFusedMoEMethod) dequantized
FP8 block-scale MoE weights to BF16 at load time to work around CUTLASS TMA
failures on SM120 (RTX PRO 6000 Blackwell). This doubled expert memory and
was a temporary workaround.
This commit replaces it with a proper Triton BF16-activation + FP8-weight
kernel (fused_moe_triton_fp8_block_scale.py) that:
- Loads FP8 weights and dequants on-the-fly inside the GEMM K-loop
- Handles FP8 block-scale (128×128) correctly
- Drops in as a replacement for torch.ops.trtllm.fused_moe on SM120
Key bug fixed: gate/up split in SwiGLU was reversed. TRT-LLM stores
w3_w1 = concat(w3/up_proj, w1/gate_proj), so SwiGLU must apply
silu() to the SECOND half (w1/gate) and multiply by the FIRST half
(w3/up). Both the Triton path and BF16 fallback had this backwards,
causing coherent-but-wrong model outputs (passed unit tests comparing
the two paths, but actual model generated Brussels instead of Paris).
Validated on Laguna-XS FP8, TP=2, SM120:
- Smoke test: PASSED (correct coherent outputs)
- GSM8K@200: 95.5% (target ≥ 92.95%)
- Throughput: 274 tok/s @ batch_size=64 (target ≥ 262.5)
Also adds cuda_graph_config=None to eval scripts since CUDA graph capture
is not compatible with the Triton fallback path on SM120.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][fix] SM120: CUDA-graph-compatible moe_align_block_size + null quant config guard
moe_align_block_size in fused_moe_triton_fp8_block_scale.py contained two
ops that abort CUDA graph capture on SM120 (RTX PRO 6000 Blackwell):
1. torch.bincount() internally syncs to the host to find the input max,
triggering cudaErrorStreamCaptureUnsupported. Replaced with
scatter_add_() on a pre-allocated torch.zeros(num_experts) buffer.
2. The previous host-side Python loop over experts used .item() calls and
dynamically-sized tensor allocations. Replaced with fully vectorised
GPU ops: argsort + prefix-sum offsets + scatter_ for sorted_token_ids,
and searchsorted on cumulative block counts for expert_ids_out. All
output buffer sizes are derived from Python-int shape attributes so no
GPU→CPU synchronisation occurs.
Also fix model_config.py: hasattr(pretrained_config, "quantization_config")
returned True even when model_kwargs had overridden the field to None,
causing load_hf_quant_config(None, ...) to crash with AttributeError.
Changed to getattr(..., None) is not None.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][feat] Add native Laguna / Laguna-XS model for PyTorch backend
Ports the out-of-tree Laguna model (LagunaForCausalLM) into the TRT-LLM
model registry so it is available without --custom_module_dirs.
Architecture: Qwen3-MoE-like with per-head softplus gating, per-layer
variable Q-head counts, mixed global/sliding attention with dual RoPE
(yarn/partial for global layers, linear/full for SWA layers), and
sigmoid MoE routing with e_score_correction_bias (MiniMaxM2 method).
Components:
- LagunaGate: sigmoid router with e_score_correction_bias
- LagunaMoE: routed experts + optional shared expert + TP allreduce
- LagunaAttention: QKNormRoPEAttention subclass with per-head softplus
gating between the attention kernel and o_proj, per-layer head counts,
and per-layer-type RoPE params
- LagunaDecoderLayer / LagunaModel / LagunaForCausalLM
- LagunaHfWeightMapper: renames gate_proj/up_proj/down_proj → w1/w3/w2,
remaps weight_scale → weight_scale_inv for FP8 block-scale checkpoints
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][fix] llm_utils: treat null quant_method in quantization_config as no quantization
When a checkpoint's config.json contains quantization_config with an
explicit null quant_method (e.g. {"quant_method": null, "mode": null}),
the existing dispatch fell through to NotImplementedError. Add a null
guard so it is treated as "no quantization" and the model loads as BF16.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][perf] SM120 FP8 MoE: replace multi-op routing with Triton histogram+scatter
Replace moe_align_block_size's 5-op PyTorch sequence (argsort + scatter_add_
+ cumsum*2 + scatter_ + searchsorted) with 4 lightweight Triton kernels:
1. _moe_histogram_kernel: vectorised atomic_add histogram (one warp/CTA)
2. _moe_prefix_kernel: exclusive prefix sum via tl.associative_scan (1 CTA)
3. _moe_expert_ids_kernel: fill expert_ids_out (one CTA/expert, runtime loop)
4. _moe_scatter_kernel: vectorised atomic scatter (one warp/CTA)
The key insight: tl.atomic_add accepts a vector of pointers, making the
histogram and scatter steps fully parallel across SIMD lanes without a
serial loop.
Measured on RTX PRO 6000 Blackwell (SM120):
N=8 (decode): 77->16 us (5.0x)
N=256 (batch): 82->16 us (5.1x)
N=8000 (prefill): 150->21 us (7.2x)
All four kernels are CUDA-graph compatible: buffer sizes are static Python-int
expressions; no .item() calls or device->host syncs.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][perf] SM120 FP8 MoE routing: add CuTe DSL implementation (~14% faster)
Add moe_align_block_size_cute alongside the existing Triton version.
Same 4-phase pipeline (histogram -> prefix -> expert_ids_fill -> scatter)
but implemented in CuTe DSL and fused into a single @cute.jit host call.
Key differences vs Triton:
- Kogge-Stone warp scan with cross-warp shared-memory merge (explicit SIMD
shuffle_sync_up) instead of tl.associative_scan
- Phase 2 writes to both offsets and counters in one pass, avoiding the
out_offsets.clone() that Triton needs for Phase 4
- All four kernel launches fold into one Python call (compiled once per device)
Measured on RTX PRO 6000 Blackwell (SM120), warmup=20 rep=200:
N=8 (decode): 15.0 -> 13.5 us (1.11x)
N=256 (batch): 15.7 -> 13.8 us (1.14x)
N=8000 (prefill): 18.3 -> 16.0 us (1.14x)
Falls back to the Triton version if nvidia-cutlass-dsl is not installed.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][fix] Use IS_CUTLASS_DSL_AVAILABLE from cute_dsl_utils instead of local try/except
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][fix] Laguna DEP mode: disable COLUMN sharding on g_proj
In DEP mode (enable_attention_dp=True) the Attention base class sets
effective tp_size=1 for all attention projections, so self.num_heads
equals the global head count on each GPU. g_proj was still using the
global mapping's tp_size with TensorParallelMode.COLUMN, producing
num_heads/tp_size outputs per GPU — mismatching the view in forward()
and causing a shape error at TP>1 with attention-DP enabled.
Fix: skip COLUMN sharding on g_proj when enable_attention_dp is set;
each GPU then holds and applies the full weight independently.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][feat] Laguna: add Eagle speculative decoding support
Wire SpecDecOneEngineForCausalLM as the base class for
LagunaForCausalLM and thread spec_metadata through
LagunaDecoderLayer.forward and LagunaModel.forward so Eagle can
capture hidden states at each layer for draft-token generation.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][fix] Laguna DEP mode: fix NCCL deadlocks from spurious AllReduce/TP ops
Three fixes to make enable_attention_dp=True (DEP mode) work correctly
with AllGatherReduceScatter (TRTLLM_CAN_USE_DEEP_EP=0):
1. Skip outer MoE AllReduce in DEP mode (LagunaMoE.forward): after
ReduceScatter, each rank already holds its own complete outputs.
An AllReduce here is semantically wrong and causes a NCCL deadlock
when ranks hold different token subsets.
2. Use overridden_tp_size=1 for the shared expert (LagunaMoE.__init__):
in DEP mode each rank computes the shared expert independently on
its own tokens, so TP-sharding would produce part
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Fix precommit
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Hacky fixes and documentation
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
[None][fix] Laguna: fix gating, rope params, and re: quant exclusion
Three upstream compatibility fixes for Laguna-XS checkpoints:
1. _gate_per_head: handle boolean gating=True (not just "per-head" string).
The Laguna-XS checkpoint uses gating=True; without this fix g_proj is
sized as num_heads*head_dim instead of num_heads, causing a weight
loading shape mismatch.
2. _build_rope_params: add _build_rope_from_flat_dict helper and handle
per-layer-type rope_parameters dict (keys "full_attention" /
"sliding_attention"). The HF v5 config stores separate RoPE configs
per attention type; the old code only handled a flat swa_rope_parameters
dict, hitting an assertion in RopeParams.from_config() on startup.
3. is_module_excluded_from_quantization: support "re:" prefix for regex
patterns. The FP8 quantization_config.ignore list uses "re:.*self_attn.*"
style patterns; without this fix fnmatch silently fails to match them
and attention layers are incorrectly quantized to FP8.
Validated on RTX PRO 6000 Blackwell (SM120, TRT-LLM 1.3.0rc13):
- BF16: smoke PASSED, GSM8K 80% (10 samples)
- FP8: smoke PASSED, GSM8K 90% (10 samples)
- NVFP4: known checkpoint export bug (missing layer-0 MLP weights)
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][cleanup] Remove Laguna-XS eval report, upstream fixes notes, and eval script
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Fix LagunaDecoderLayer constructor
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Replace with minimal shim
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
[None][cleanup] Remove vLLM attribution, dead get_sm_version call, restore torch.compile
- fused_moe_triton_fp8_block_scale.py: remove "Ports vLLM's fused_moe_kernel"
attribution from module docstring and section banner; kernel authorship
belongs in git history, not inline comments.
- fused_moe_trtllm_gen.py: drop bare get_sm_version() call with no return value
or side effect left over from an earlier SM version guard.
- routing.py: restore @torch.compile(options={"max-autotune": True}) on
Deepseekv3RoutingImpl.get_scores and MiniMaxM2MoeRoutingMethod.get_scores;
accidentally dropped in the Laguna squash commit.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][revert] Revert deep_ep_low_latency.py branch changes to main
Reverts the four branch-local changes to DeepEPLowLatency:
- mixed-batch guard in use_low_latency_mode (context+decode imbalance check)
- TMA alignment rounding of all_rank_max_num_tokens in dispatch
- storing all_rank_max_num_tokens in _dispatch_state
- reading all_rank_max_num_tokens from _dispatch_state instead of kwargs in combine
Validated that BF16 DPEP (tensor_parallel_size=2, enable_attention_dp=True)
passes the smoke test on SM120 with the upstream version of the file.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
fix: Fix pre-commit failure
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
[None][cleanup] Remove legacy Laguna MLP layer fallback
Require Laguna configs to use the current mlp_layer_types schema now that supported Poolside checkpoints provide it directly.
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
[None][cleanup] Remove Laguna diagnostic hooks
Drop bring-up-only tensor diagnostics and environment-gated debug prints from the release-ready Laguna model implementation.
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
[None][cleanup] Remove unused FP8 dequant MoE fallback
Drop the load-time BF16 dequant fallback now that SM120 FP8 block-scale MoE uses the Triton runtime path instead.
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
[None][cleanup] Keep TRTLLMGen SM120 guard
Leave SM120 TRTLLMGenFusedMoE unsupported since Laguna uses the validated Cutlass/Triton MoE paths instead.
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
[None][cleanup] Prune dead Triton MoE FP8 paths and CuTe DSL routing kernels
In tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py:
- Remove the CuTe DSL routing kernels block (four cute.kernel functions, the
fused cute.jit host, the compiled-kernel cache helper, moe_align_block_size_cute,
and the IS_CUTLASS_DSL_AVAILABLE import).
- Remove the per-token-group FP8 quantization section
(_per_token_group_quant_fp8_kernel, per_token_group_quant_fp8, _FP8_MAX);
no callers anywhere in the tree.
- Remove the unused FP8 block-scale GEMM path
(_fused_moe_fp8_block_scale_kernel, _invoke_fp8_block_scale_moe_kernel);
the live MoE path is the BF16-activation + FP8-weight kernel.
- Remove the BF16 debug fallback (_run_moe_bf16_fallback) and the env-var-gated
TRTLLM_MOE_TRITON_DEBUG / TRTLLM_MOE_TRITON_BF16 paths in
run_triton_fp8_block_scale_moe.
Verified via the Laguna-XS smoke test on bf16, fp8, and nvfp4 (single-GPU,
SM120) plus the GSM8K 10-sample integration check; all three variants pass.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][cleanup] Apply ruff-format to fused_moe_triton_fp8_block_scale.py
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][fix] Guard SM120 fp8-block-scale MoE branch against None token_final_scales
forward_chunk sets `token_final_scales = None` when
`apply_router_weight_on_input=True` (the routing weights are folded into x
upstream). The SM120 + FP8 block-scale early-return branch in run_moe then
dereferenced `token_final_scales.dtype` and multiplied through it,
crashing in that configuration.
Substitute ones-of-shape `token_selected_experts` (float32) when None,
mirroring the existing defensive code in the alltoall path of
forward_chunk. The Triton kernel's per-token scaling becomes a no-op,
which is correct because the routing weights are already applied to x.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][doc] Add LagunaForCausalLM to supported-models matrix
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][fix] Merge modules_to_not_convert with ignore in llm_utils compressed-tensors loader
The compressed-tensors quant-config branch in llm_utils.py only assigned
hf_quant_config["ignore"] to quant_config.exclude_modules, dropping any
modules listed under HF's `modules_to_not_convert` field. The PyTorch
backend equivalent in _torch/model_config.py already merges both lists
(set-deduplicated) for the same quant_method. Mirror that behavior here so
both loaders produce identical exclude_modules from the same checkpoint.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][fix] llm_utils: return False when HF quant_method is null
_update_from_hf_quant_config falls through to `return True` after the
elif chain, so the null-quant_method branch (added to gracefully handle
quantization_config blocks without an active quant_method) was being
reported to callers as if the checkpoint were already quantized.
`_load_model_from_hf` then skipped calibration (line 582:
`elif self.llm_args.quant_config._requires_calibration and not prequantized:`),
which is wrong for a non-quantized checkpoint. Return False explicitly so
calibration is not skipped.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][fix] QuantConfig._get_quant_cfg: use shared exclude-pattern matcher
QuantConfig._get_quant_cfg only honored exact matches and trailing-* prefix
matches, while is_module_excluded_from_quantization (the matcher used
elsewhere in the same class) supports the full set of exclude_modules
patterns: 're:'-prefixed regexes via re.fullmatch and arbitrary fnmatch
glob patterns via fnmatch.fnmatchcase. As a result, exclusions like
'*.mlp.gate' or 're:layers\\.\\d+\\.qkv' were skipped by _get_quant_cfg.
Delegate to is_module_excluded_from_quantization so both code paths agree.
PretrainedConfig._get_quant_cfg already forwards to this method, so this
fixes both call sites.
Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
[None][perf] Optimize Triton FP8 block-scale MoE
Consolidate the Triton FP8 block-scale MoE optimization work into a single change covering fused activation/top-k handling, autotuning cleanup, and focused test coverage.
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
[None][test] Add LagunaXS accuracy coverage
Add PyTorch backend accuracy coverage for LagunaXS BF16, FP8, and NVFP4 checkpoints on MMLU and GSM8K to catch regressions across quantization formats.
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com>
DomBrown
force-pushed
the
user/domb/laguna-rebase
branch
from
May 21, 2026 08:09
719c223 to
c7b44f1
Compare
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #49648 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49648 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49712 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49712 [ run ] completed with state
|
Collaborator
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #49750 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #49750 [ run ] completed with state |
1 task done
Contributor
Build issues hit when targeting sm_120 (RTX PRO 6000 Blackwell workstation)Built the PR against 1)
|
KleinBlueC
pushed a commit
to KleinBlueC/TensorRT-LLM
that referenced
this pull request
May 26, 2026
…13559) Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> Co-authored-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
bmarimuthu-nv
pushed a commit
to nv-auto-deploy/TensorRT-LLM
that referenced
this pull request
May 28, 2026
…13559) Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> Co-authored-by: Daniel Campora <961215+dcampora@users.noreply.github.com>
Merged
1 task done
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
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 by CodeRabbit
Release Notes
New Features
Improvements
Description
Adds TensorRT-LLM PyTorch-backend support for Poolside's Laguna model family, including the current generation Laguna-XS.2 in bf16, FP8 block-scale, and NVFP4 variants.
Laguna-specific bits: per-head softplus output gate (
g_projpre-o_proj), per-layer attention-head counts, per-layer dual RoPE (yarn for full-attention, default linear for sliding-window), MiniMaxM2-style sigmoid routing with per-layer routed scaling, and both legacy (mlp_only_layers+decoder_sparse_step) and current (mlp_layer_types) MoE-vs-dense schemas.The branch also lands several generally useful fixes that Laguna depends on:
kv_cache_schemeparsing inModelConfig.load_hf_quant_config/llm_utils.ModelLoader.re:-prefixed regex patterns supported inQuantConfig.is_module_excluded_from_quantization(compressed-tensors writes regex-shaped exclusion lists).lagunamodel_type registered with the DeepSeek-R1-style reasoning parser.Validation: full GSM8K (1319 samples)
Hardware: NVIDIA RTX PRO 6000 Blackwell (SM120).
FP8 and NVFP4 stay within ~1 point of bf16 (well inside
±0.97stderr) and execute ~1.4× / ~2.2× faster respectively. All three variants load via the same overlay script and the same modeling code, with no per-checkpoint or per-quant-format patches.Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.