From 9289d06912850c676c9024347e0dfabe16b20d28 Mon Sep 17 00:00:00 2001 From: Daniel Campora <961215+dcampora@users.noreply.github.com> Date: Fri, 17 Apr 2026 10:40:20 +0000 Subject: [PATCH 1/9] Rebase and squash of Laguna development commits MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit [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 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> --- docs/source/models/supported-models.md | 1 + laguna_minimal_overlay.sh | 73 ++ tensorrt_llm/_torch/model_config.py | 37 +- tensorrt_llm/_torch/models/__init__.py | 2 + tensorrt_llm/_torch/models/modeling_laguna.py | 678 ++++++++++++++++++ .../modules/fused_moe/fused_moe_cutlass.py | 50 +- .../fused_moe_triton_fp8_block_scale.py | 603 ++++++++++++++++ tensorrt_llm/llmapi/llm_utils.py | 49 +- tensorrt_llm/llmapi/reasoning_parser.py | 2 + tensorrt_llm/models/modeling_utils.py | 22 +- .../defs/accuracy/references/gsm8k.yaml | 8 + .../defs/accuracy/references/mmlu.yaml | 8 + .../defs/accuracy/test_llm_api_pytorch.py | 54 ++ .../fused_moe/test_triton_fp8_block_scale.py | 203 ++++++ 14 files changed, 1769 insertions(+), 21 deletions(-) create mode 100755 laguna_minimal_overlay.sh create mode 100644 tensorrt_llm/_torch/models/modeling_laguna.py create mode 100644 tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py create mode 100644 tests/unittest/_torch/modules/fused_moe/test_triton_fp8_block_scale.py diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 6670e8366b96..a2b4120e59a7 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -27,6 +27,7 @@ The following is a table of supported models for the PyTorch backend: | `HunYuanMoEForCausalLM` [^5] | Hunyuan MoE | `tencent/Hunyuan-A13B-Instruct` | | `InternLM3ForCausalLM` [^5] | InternLM3 | `internlm/internlm3-8b-instruct` | | `KimiK25ForConditionalGeneration` | Kimi-K2.5 | `moonshotai/Kimi-K2.5` | +| `LagunaForCausalLM` | Laguna-XS | `poolside/laguna-XS.2` | | `LlamaForCausalLM` | Llama 3.1, Llama 3, Llama 2, LLaMA | `meta-llama/Meta-Llama-3.1-70B` | | `Llama4ForConditionalGeneration` | Llama 4 | `meta-llama/Llama-4-Scout-17B-16E-Instruct` | | `MiniMaxM2ForCausalLM` [^5] | MiniMax M2/M2.1/M2.7 | `MiniMaxAI/MiniMax-M2.7` | diff --git a/laguna_minimal_overlay.sh b/laguna_minimal_overlay.sh new file mode 100755 index 000000000000..9a0bf23319f0 --- /dev/null +++ b/laguna_minimal_overlay.sh @@ -0,0 +1,73 @@ +#!/usr/bin/env bash +# Minimal transformers-4.57-compat overlay for a Laguna checkpoint directory. +# +# Laguna's `configuration_laguna.py` imports two transformers >= 4.58 symbols: +# - `transformers.configuration_utils.PreTrainedConfig` +# - `transformers.modeling_rope_utils.RopeParameters` +# Neither exists in transformers 4.57, which is what we have installed. +# +# This overlay symlinks every file in the checkpoint dir, then replaces just +# `configuration_laguna.py` with a copy that is prefixed with a compat shim +# aliasing the missing symbols. No other files are modified — additional +# fixes (config.json rewrites, etc.) are deliberately out of scope. +# +# Usage: +# laguna_minimal_overlay.sh [overlay_dir] +# Echoes the overlay directory path on stdout. +set -euo pipefail + +if [[ $# -lt 1 ]]; then + echo "Usage: $(basename "$0") [overlay_dir]" >&2 + exit 1 +fi + +SRC="$(realpath "$1")" +OVERLAY="${2:-/tmp/laguna-minimal-$(basename "$SRC")}" + +if [[ ! -d "$SRC" ]]; then + echo "ERROR: source dir does not exist: $SRC" >&2 + exit 1 +fi +if [[ ! -f "$SRC/configuration_laguna.py" ]]; then + echo "ERROR: $SRC has no configuration_laguna.py" >&2 + exit 1 +fi + +rm -rf "$OVERLAY" +mkdir -p "$OVERLAY" + +for entry in "$SRC"/*; do + name="$(basename "$entry")" + if [[ "$name" == "configuration_laguna.py" ]]; then + continue + fi + ln -s "$entry" "$OVERLAY/$name" +done + +SHIM_HEADER='# --- compat shim (added by laguna_minimal_overlay.sh) --- +try: + from transformers.configuration_utils import PreTrainedConfig # transformers >= 4.58 +except ImportError: + from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig # transformers 4.57 +try: + from transformers.modeling_rope_utils import RopeParameters # transformers >= 4.58 +except ImportError: + RopeParameters = dict # type: ignore[misc,assignment] +# Neuter `@strict` (huggingface_hub) — newer Laguna configs decorate a non-dataclass +# subclass of PretrainedConfig with it, which raises in our hf-hub version. +def strict(cls=None, **_kwargs): # type: ignore[no-redef] + return cls if cls is not None else (lambda c: c) +# Neuter `@auto_docstring` (transformers utils) — only adds docstrings, safe to skip. +def auto_docstring(*_args, **_kwargs): # type: ignore[no-redef] + return (lambda c: c) if not _args or not callable(_args[0]) else _args[0] +# --- end compat shim --- +' + +# Strip the original imports we just shimmed; everything else is kept verbatim. +{ + printf '%s' "$SHIM_HEADER" + grep -vE '^from transformers\.configuration_utils import PreTrainedConfig$|^from transformers\.modeling_rope_utils import RopeParameters$|^from huggingface_hub\.dataclasses import strict$|^from transformers\.utils import auto_docstring$' \ + "$SRC/configuration_laguna.py" +} > "$OVERLAY/configuration_laguna.py" + +echo "$OVERLAY" diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index d9e3fb8cf17a..dc4ed2259e51 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -513,12 +513,42 @@ def load_hf_quant_config(hf_quant_config, moe_backend, checkpoint_dir=None): raise ValueError( f"Unsupported weights_quant_strategy: {weights_quant_strategy}. " "Supported strategies: 'channel', 'block'.") + elif (weights_quant_config["num_bits"] == 4 + and weights_quant_config.get("type") == "float" + and weights_quant_strategy == "tensor_group"): + # llm-compressor NVFP4: weights FP4 with FP8 per-group scales + # (group_size=16), scaled by an FP32 global scale. + if inputs_quant_strategy != "tensor_group": + raise ValueError( + f"Unsupported inputs_quant_strategy for NVFP4: {inputs_quant_strategy}." + ) + group_size = weights_quant_config["group_size"] + if group_size != 16: + raise ValueError( + f"Unsupported group_size: {group_size}. Supported: 16 for NVFP4." + ) + quant_config.quant_algo = QuantAlgo.NVFP4 + quant_config.group_size = group_size else: raise ValueError( f"Unsupported quant_bits: {weights_quant_config['num_bits']}. " - "Supported: 8.") + "Supported: 8 (FP8) or 4 (NVFP4).") + + # kv_cache_scheme (llm-compressor): FP8 per-tensor KV cache. + kv_cache_scheme = hf_quant_config.get("kv_cache_scheme") + if kv_cache_scheme is not None: + if (kv_cache_scheme.get("num_bits") == 8 + and kv_cache_scheme.get("type") == "float"): + quant_config.kv_cache_quant_algo = QuantAlgo.FP8 + else: + raise ValueError( + f"Unsupported kv_cache_scheme: {kv_cache_scheme}.") - quant_config.exclude_modules = hf_quant_config.get("ignore", []) + if hf_exclude_modules is not None: + quant_config.exclude_modules = list( + set(hf_exclude_modules + hf_quant_config.get("ignore", []))) + else: + quant_config.exclude_modules = hf_quant_config.get("ignore", []) elif hf_quant_config.get("quant_method") == "nvfp4": quant_config.quant_algo = QuantAlgo.NVFP4 group_size = hf_quant_config.get("group_size", 16) @@ -740,7 +770,8 @@ def _recursive_update_config(config: transformers.PretrainedConfig, quant_config, layer_quant_config = cls._build_modelopt_quant_config( normalized, checkpoint_dir, moe_backend_hint) # quantized ckpt in other formats - elif hasattr(pretrained_config, "quantization_config"): + elif getattr(pretrained_config, "quantization_config", + None) is not None: hf_quant_config = pretrained_config.quantization_config quant_config, layer_quant_config = cls.load_hf_quant_config( hf_quant_config, diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index 55177767ac5a..25aa4ce34cce 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -20,6 +20,7 @@ from .modeling_hunyuan_moe import HunYuanMoEV1ForCausalLM from .modeling_hyperclovax import HCXVisionForCausalLM from .modeling_kimi_k25 import KimiK25ForConditionalGeneration +from .modeling_laguna import LagunaForCausalLM from .modeling_llama import LlamaForCausalLM from .modeling_llava_next import LlavaNextModel from .modeling_minimaxm2 import MiniMaxM2ForCausalLM @@ -58,6 +59,7 @@ "Gemma3ForCausalLM", "Gemma3VLM", "HCXVisionForCausalLM", + "LagunaForCausalLM", "HunYuanDenseV1ForCausalLM", "HunYuanMoEV1ForCausalLM", "KimiK25ForConditionalGeneration", diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py new file mode 100644 index 000000000000..36b75c854efe --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -0,0 +1,678 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Laguna / Laguna-XS model for TensorRT-LLM PyTorch backend.""" + +from typing import Dict, List, Optional, Type + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig + +from tensorrt_llm.functional import PositionEmbeddingType, RotaryScalingType + +from ..attention_backend import AttentionMetadata +from ..attention_backend.interface import ( + PositionalEmbeddingParams, + PredefinedAttentionMask, + RopeParams, +) +from ..distributed import AllReduce, AllReduceParams +from ..modules.attention import _helix_cp_allgather_input, _helix_cp_output_projection +from ..modules.decoder_layer import DecoderLayer +from ..modules.embedding import Embedding +from ..modules.fused_moe import MiniMaxM2MoeRoutingMethod, create_moe, get_moe_cls +from ..modules.fused_moe.interface import MoE, MoEWeightLoadingMode +from ..modules.fused_moe.interface import MoE as MoEInterface +from ..modules.gated_mlp import GatedMLP +from ..modules.linear import Linear, TensorParallelMode +from ..modules.qk_norm_attention import QKNormRoPEAttention +from ..modules.rms_norm import RMSNorm +from ..speculative import SpecMetadata +from ..utils import AuxStreamType +from .checkpoints.hf.weight_mapper import HfWeightMapper +from .modeling_speculative import SpecDecOneEngineForCausalLM +from .modeling_utils import DecoderModel, register_auto_model, register_mapper + +# --------------------------------------------------------------------------- +# Gate (MoE router) +# --------------------------------------------------------------------------- + + +class LagunaGate(nn.Module): + """Sigmoid router with e_score_correction_bias (MiniMaxM2 routing).""" + + def __init__( + self, + hidden_size: int, + num_experts: int, + top_k: int, + dtype: Optional[torch.dtype] = None, + moe_backend_cls: Type[MoE] = None, + ): + super().__init__() + self.top_k = top_k + self.moe_backend_cls = moe_backend_cls + self.weight = nn.Parameter( + torch.empty((num_experts, hidden_size), dtype=dtype), + requires_grad=False, + ) + self.e_score_correction_bias = nn.Parameter( + torch.zeros(num_experts, dtype=torch.float32), + requires_grad=False, + ) + self.out_dtype = dtype + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = torch.ops.trtllm.cublas_mm( + hidden_states, self.weight.t(), bias=None, out_dtype=self.out_dtype + ) + return logits + + def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False): + assert len(weights) == 1 + w = weights[0].get("weight") + if w is not None: + self.weight.copy_(w[:]) + b = weights[0].get("e_score_correction_bias") + if b is not None: + self.e_score_correction_bias.copy_(b[:]) + + @property + def routing_method(self): + return MiniMaxM2MoeRoutingMethod( + top_k=self.top_k, + num_experts=self.weight.shape[0], + callable_e_score_correction_bias=lambda: self.e_score_correction_bias, + output_dtype=torch.float32, + ) + + +# --------------------------------------------------------------------------- +# MoE +# --------------------------------------------------------------------------- + + +class LagunaMoE(nn.Module): + def __init__(self, model_config, layer_idx, aux_stream_dict): + super().__init__() + config = model_config.pretrained_config + self.hidden_dim = config.hidden_size + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_tok + self.mapping = model_config.mapping + + self.allreduce = None + if self.mapping.tp_size > 1: + self.allreduce = AllReduce( + mapping=model_config.mapping, + strategy=model_config.allreduce_strategy, + ) + + self.gate = LagunaGate( + hidden_size=self.hidden_dim, + num_experts=self.num_experts, + top_k=self.top_k, + dtype=config.torch_dtype, + moe_backend_cls=get_moe_cls(model_config), + ) + + self.experts = create_moe( + num_experts=self.num_experts, + routing_method=self.gate.routing_method, + hidden_size=self.hidden_dim, + intermediate_size=config.moe_intermediate_size, + aux_stream_dict=aux_stream_dict, + dtype=config.torch_dtype, + reduce_results=False, + model_config=model_config, + layer_idx=layer_idx, + weight_loading_mode=MoEWeightLoadingMode.VANILLA, + ) + + shared_size = getattr(config, "shared_expert_intermediate_size", 0) + if shared_size > 0: + # In DEP mode each rank processes different tokens, so TP-sharding the + # shared expert across ranks would produce partial outputs that cannot + # be combined with AllReduce (ranks hold different tokens). Use + # overridden_tp_size=1 so every rank holds full weights and computes + # the shared expert independently. + shared_expert_tp_size = 1 if self.mapping.enable_attention_dp else None + self.shared_expert = GatedMLP( + hidden_size=self.hidden_dim, + intermediate_size=shared_size, + bias=False, + dtype=config.torch_dtype, + config=model_config, + overridden_tp_size=shared_expert_tp_size, + is_shared_expert=True, + reduce_output=False, + ) + else: + self.shared_expert = None + + self.routed_scaling_factor = float(getattr(config, "moe_routed_scaling_factor", 1.0)) + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + all_reduce_params: Optional[AllReduceParams] = None, + **kwargs, + ) -> torch.Tensor: + orig_shape = hidden_states.shape + hidden_states = hidden_states.view(-1, self.hidden_dim) + all_rank_num_tokens = attn_metadata.all_rank_num_tokens + + router_logits = self.gate(hidden_states) + + routed_output = self.experts( + hidden_states, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + ) + + if self.routed_scaling_factor != 1.0: + routed_output = routed_output * self.routed_scaling_factor + + if self.shared_expert is not None: + shared_output = self.shared_expert(hidden_states) + final_output = routed_output + shared_output + else: + final_output = routed_output + + # In DEP mode (enable_attention_dp) AllGatherReduceScatter already + # produced complete per-rank outputs via ReduceScatter; an AllReduce + # here would be semantically wrong (ranks hold different tokens) and + # causes a NCCL deadlock. Skip it in DEP mode. + if ( + self.mapping.tp_size > 1 + and self.allreduce is not None + and not self.mapping.enable_attention_dp + ): + final_output = self.allreduce(final_output, all_reduce_params=all_reduce_params) + + return final_output.view(orig_shape) + + +# --------------------------------------------------------------------------- +# Attention +# --------------------------------------------------------------------------- + + +class LagunaAttention(QKNormRoPEAttention): + """Laguna attention with per-head softplus gating, per-layer heads, + dual RoPE, and sliding window.""" + + def __init__( + self, + model_config, + layer_idx: Optional[int] = None, + reduce_output: bool = True, + ): + config = model_config.pretrained_config + + per_layer_heads = getattr(config, "num_attention_heads_per_layer", None) + num_heads = ( + per_layer_heads[layer_idx] + if per_layer_heads is not None + else config.num_attention_heads + ) + + layer_types = getattr(config, "layer_types", None) + is_sliding = layer_types is not None and layer_types[layer_idx] == "sliding_attention" + + rope_params = self._build_rope_params(config, is_sliding) + + self.attention_window_size = config.sliding_window if is_sliding else None + + pos_embd_params = PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=rope_params, + ) + + gating = getattr(config, "gating", False) + self._use_gating = bool(gating) + self._gate_per_head = gating == "per-head" or gating is True + self._num_heads_local = None # set after super().__init__ + + # fuse_qk_norm_rope=False is required: the fused kernel reads + # partial_rotary_factor and yarn params from pretrained_config + # globally, ignoring per-layer RopeParams. Laguna has different + # RoPE per layer type (global=yarn/partial=0.5, SWA=linear/full). + super().__init__( + hidden_size=config.hidden_size, + num_attention_heads=num_heads, + num_key_value_heads=config.num_key_value_heads, + max_position_embeddings=config.max_position_embeddings, + bias=getattr(config, "qkv_bias", False) or getattr(config, "attention_bias", False), + pos_embd_params=pos_embd_params, + fuse_qk_norm_rope=False, + layer_idx=layer_idx, + dtype=config.torch_dtype, + dense_bias=False, + config=model_config, + reduce_output=reduce_output, + ) + + self._num_heads_local = self.num_heads + + if self._use_gating: + g_out = num_heads if self._gate_per_head else num_heads * self.head_dim + # In DEP mode (enable_attention_dp=True) the attention base class uses + # a local mapping with tp_size=1, so self.num_heads == num_heads (full). + # g_proj must match: no COLUMN sharding, each GPU holds the full weight. + g_tp_mode = ( + None if model_config.mapping.enable_attention_dp else TensorParallelMode.COLUMN + ) + self.g_proj = Linear( + config.hidden_size, + g_out, + bias=False, + dtype=config.torch_dtype, + mapping=model_config.mapping, + tensor_parallel_mode=g_tp_mode, + ) + + @staticmethod + def _build_rope_from_flat_dict(config, rp_dict: dict) -> RopeParams: + rp = RopeParams() + rp.theta = float(rp_dict.get("rope_theta", 10000.0)) + rp.max_positions = config.max_position_embeddings + rp.dim = int(config.head_dim * float(rp_dict.get("partial_rotary_factor", 1.0))) + rope_type = rp_dict.get("rope_type", "default") + if rope_type == "linear": + rp.scale_type = RotaryScalingType.linear + rp.scale = float(rp_dict.get("factor", 1.0)) + elif rope_type == "yarn": + rp.scale_type = RotaryScalingType.yarn + rp.scale = float(rp_dict.get("factor", 1.0)) + rp.beta_fast = float(rp_dict.get("beta_fast", 32.0)) + rp.beta_slow = float(rp_dict.get("beta_slow", 1.0)) + rp.mscale = float(rp_dict.get("attention_factor", 1.0)) + rp.original_max_positions = int( + rp_dict.get("original_max_position_embeddings", config.max_position_embeddings) + ) + else: + rp.scale_type = RotaryScalingType.none + rp.scale = 1.0 + return rp + + @staticmethod + def _build_rope_params(config, is_sliding: bool) -> RopeParams: + rope_parameters = getattr(config, "rope_parameters", None) + if isinstance(rope_parameters, dict): + layer_key = "sliding_attention" if is_sliding else "full_attention" + sub = rope_parameters.get(layer_key) or rope_parameters.get("full_attention") + if sub is not None: + return LagunaAttention._build_rope_from_flat_dict(config, sub) + if is_sliding: + swa_rp = getattr(config, "swa_rope_parameters", None) + if swa_rp is not None: + return LagunaAttention._build_rope_from_flat_dict(config, swa_rp) + return RopeParams.from_config(config) + else: + rp = RopeParams.from_config(config) + prf = getattr(config, "partial_rotary_factor", None) + if prf is not None: + rp.dim = int(config.head_dim * prf) + return rp + + def forward( + self, + position_ids, + hidden_states, + attn_metadata, + attention_mask=PredefinedAttentionMask.CAUSAL, + all_reduce_params=None, + lora_params=None, + attention_window_size=None, + attention_mask_data=None, + **kwargs, + ): + # Override the full Attention.forward so we can insert the g_proj + # gate between the attention kernel output and o_proj (the correct + # position per the HF reference implementation). + gate_input = hidden_states if self._use_gating else None + + hidden_states = _helix_cp_allgather_input( + hidden_states, attn_metadata, self.mapping, self.layer_idx + ) + + qkv = self.qkv_proj(hidden_states) + q, k, v = qkv, None, None + + # QK-norm + RoPE (inherited from QKNormRoPEAttention.apply_rope) + q, k, v = self.apply_rope(q, k, v, position_ids) + q, k, v = self.convert_qkv(q, k, v) + + window = attention_window_size or self.attention_window_size + attn_output = self.forward_impl( + q, + k, + v, + attn_metadata, + attention_mask, + window, + attention_mask_data, + mrope_config=None, + has_lora=bool(lora_params), + ) + + # Per-head softplus gating (pre-o_proj, matching HF reference) + if self._use_gating and gate_input is not None: + gate = self.g_proj(gate_input) + gate = F.softplus(gate.float()).to(attn_output.dtype) + if self._gate_per_head: + shape = attn_output.shape + attn_output = ( + attn_output.view(*shape[:-1], self._num_heads_local, self.head_dim) + * gate.unsqueeze(-1) + ).view(shape) + else: + attn_output = attn_output * gate + + attn_output = _helix_cp_output_projection( + self.o_proj, + attn_output, + attn_metadata, + all_reduce_params, + self.mapping, + self.mapping_o, + self.layer_idx, + lora_params, + ) + return attn_output + + +# --------------------------------------------------------------------------- +# Decoder layer +# --------------------------------------------------------------------------- + + +class LagunaDecoderLayer(DecoderLayer): + def __init__(self, model_config, layer_idx: int, aux_stream_dict): + super().__init__() + config = model_config.pretrained_config + self.layer_idx = layer_idx + self.mapping = model_config.mapping + + self.self_attn = LagunaAttention( + model_config, + layer_idx=layer_idx, + reduce_output=self.mapping.tp_size > 1, + ) + + is_moe_layer = config.num_experts > 0 and config.mlp_layer_types[layer_idx] == "sparse" + if is_moe_layer: + self.mlp = LagunaMoE(model_config, layer_idx, aux_stream_dict) + else: + # In DEP mode each rank holds different tokens; TP-sharding the + # dense MLP would require an AllReduce that deadlocks across ranks + # (ranks are at different positions in their token streams). + # Use overridden_tp_size=1 so every rank computes full MLP locally. + dense_mlp_tp_size = 1 if model_config.mapping.enable_attention_dp else None + self.mlp = GatedMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + bias=False, + dtype=config.torch_dtype, + config=model_config, + overridden_tp_size=dense_mlp_tp_size, + ) + + self.input_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + self.post_attention_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + + def forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + residual: Optional[torch.Tensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ): + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self.self_attn( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + **kwargs, + ) + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + + if isinstance(self.mlp, LagunaMoE): + hidden_states = self.mlp(hidden_states, attn_metadata) + else: + hidden_states = self.mlp(hidden_states) + + if spec_metadata is not None: + spec_metadata.maybe_capture_hidden_states(self.layer_idx, hidden_states, residual) + + return hidden_states, residual + + +# --------------------------------------------------------------------------- +# Model + CausalLM +# --------------------------------------------------------------------------- + + +class LagunaModel(DecoderModel): + def __init__(self, model_config): + super().__init__(model_config) + config = model_config.pretrained_config + + self.aux_stream_dict = { + AuxStreamType.MoeChunkingOverlap: torch.cuda.Stream(), + AuxStreamType.MoeBalancer: torch.cuda.Stream(), + AuxStreamType.MoeOutputMemset: torch.cuda.Stream(), + } + + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + mapping=model_config.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=True, + ) + + self.layers = nn.ModuleList( + [ + LagunaDecoderLayer(model_config, layer_idx, self.aux_stream_dict) + for layer_idx in range(config.num_hidden_layers) + ] + ) + + self.norm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + + def forward( + self, + attn_metadata: AttentionMetadata, + input_ids: Optional[torch.IntTensor] = None, + position_ids: Optional[torch.IntTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + spec_metadata: Optional[SpecMetadata] = None, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + hidden_states = inputs_embeds + residual = None + for layer in self.layers: + hidden_states, residual = layer( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=residual, + spec_metadata=spec_metadata, + **kwargs, + ) + hidden_states, _ = self.norm(hidden_states, residual) + + return hidden_states + + +@register_auto_model("LagunaForCausalLM") +class LagunaForCausalLM(SpecDecOneEngineForCausalLM[LagunaModel, PretrainedConfig]): + def __init__(self, model_config): + super().__init__(LagunaModel(model_config), model_config) + + def load_weights( + self, + weights: Dict, + weight_mapper: Optional["LagunaHfWeightMapper"] = None, + params_map: Optional[Dict[str, str]] = None, + allow_partial_loading: bool = False, + ): + if weight_mapper is not None: + weights = weight_mapper.preprocess_weights(weights) + + super().load_weights( + weights=weights, + weight_mapper=weight_mapper, + params_map=params_map, + allow_partial_loading=allow_partial_loading, + ) + + +# --------------------------------------------------------------------------- +# Weight mapper (HF checkpoint -> TRT-LLM module names) +# --------------------------------------------------------------------------- + + +@register_mapper("HF", "LagunaForCausalLM") +class LagunaHfWeightMapper(HfWeightMapper): + def preprocess_weights(self, weights: dict) -> dict: + weights = self.rename_by_params_map( + { + r"(.*)mlp\.experts\.e_score_correction_bias(.*)": r"\1mlp.gate.e_score_correction_bias\2", + }, + weights, + ) + + # Translate llm-compressor "compressed-tensors" NVFP4 naming into what + # TRT-LLM's NVFP4 linear / MoE loaders expect. llm-compressor stores + # weight_packed (uint8, packed FP4) + # weight_scale (FP8 per-group) + # weight_global_scale (FP32 scalar = (FP8_MAX*FP4_MAX)/amax_weight) + # input_global_scale (FP32 scalar = (FP8_MAX*FP4_MAX)/amax_input) + # TRT-LLM's loaders expect: + # weight (uint8 view of float4_e2m1x2) + # weight_scale (FP8, same) + # weight_scale_2 (FP32 scalar = amax_weight/(FP8_MAX*FP4_MAX)) + # input_scale (FP32 scalar = amax_input/(FP8_MAX*FP4_MAX)) + # which is the reciprocal of the *_global_scale values. + if self._config is not None and self._config.quant_config.quant_mode.has_nvfp4(): + + def _invert_global_scale(t: torch.Tensor) -> torch.Tensor: + """Convert compressed-tensors multiplier form to TRT-LLM + divisor form, protecting against uncalibrated experts. + + llm-compressor stores global_scale=(FP8_MAX*FP4_MAX)/amax; + TRT-LLM expects {weight_scale_2, input_scale}=amax/(FP8_MAX*FP4_MAX), + i.e. the reciprocal. Some checkpoints store 0 for experts with + no calibration data (amax unavailable); a naive 1/x would + produce Inf which then contaminates the MoE per-layer + max-reduction over experts and zeroes out *every* expert's + input_scale. We preserve 0 -> 0 so those entries are harmless + in the max reduction and the corresponding expert's alpha + evaluates to 0 (silencing the dead expert). + """ + t = t.to(torch.float32) + out = torch.where( + t > 0, + 1.0 / torch.clamp(t, min=torch.finfo(torch.float32).tiny), + torch.zeros_like(t), + ) + # Belt-and-braces: any residual Inf/NaN becomes 0. + return torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0).contiguous() + + renamed: dict = {} + for name, value in weights.items(): + if name.endswith(".weight_packed"): + new_name = name[: -len(".weight_packed")] + ".weight" + renamed[new_name] = value + elif name.endswith(".weight_global_scale"): + new_name = name[: -len(".weight_global_scale")] + ".weight_scale_2" + t = value[...] if not isinstance(value, torch.Tensor) else value + renamed[new_name] = _invert_global_scale(t) + elif name.endswith(".input_global_scale"): + new_name = name[: -len(".input_global_scale")] + ".input_scale" + t = value[...] if not isinstance(value, torch.Tensor) else value + renamed[new_name] = _invert_global_scale(t) + else: + renamed[name] = value + weights = renamed + + return weights + + def is_special_instance_module(self, module: nn.Module) -> bool: + return isinstance(module, MoEInterface) + + def handle_special_instance_module( + self, + module: nn.Module, + module_name: str, + module_weights: dict, + allow_partial_loading: bool = False, + ) -> None: + if isinstance(module, MoEInterface): + # DeepSeekFP8 block-scales MoE expects "weight_scale_inv" but NVFP4 + # MoE expects the original "weight_scale" (plus "weight_scale_2"). + use_fp8_block_rename = ( + self._config is not None + and self._config.quant_config.quant_mode.has_fp8_block_scales() + ) + updated = {} + for wn, wv in module_weights.items(): + new_wn = ( + wn.replace("gate_proj", "w1") + .replace("up_proj", "w3") + .replace("down_proj", "w2") + ) + if use_fp8_block_rename: + # ModelOpt uses "weight_scale"; TRT-LLM DeepSeekFP8 + # quantization expects "weight_scale_inv" (same values). + new_wn = new_wn.replace("weight_scale", "weight_scale_inv") + # filter_weights strips module prefix leaving "experts.X.w1.weight"; + # VANILLA mode expects "X.w1.weight" (integer prefix only). + if new_wn.startswith("experts."): + new_wn = new_wn[len("experts.") :] + updated[new_wn] = wv + del module_weights + module.load_weights(weights=[updated], allow_partial_loading=allow_partial_loading) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index ba4ce303fe23..5b9574a10065 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -22,7 +22,6 @@ # isort: off from .quantization import (DeepSeekFP8BlockScalesFusedMoEMethod, - DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm, FP8QDQFusedMoEMethod, MoEWeightLoadingMode, NVFP4CutlassFusedMoEMethod, INT8WoqPerChannelFusedMoEMethod, @@ -531,10 +530,7 @@ def _get_quant_method(self): if self.quant_config.layer_quant_mode.has_fp8_qdq(): return FP8QDQFusedMoEMethod() elif self.quant_config.layer_quant_mode.has_fp8_block_scales(): - if get_sm_version() == 120: - return DeepSeekFP8BlockScalesFusedMoEMethodDeepGemm() - else: - return DeepSeekFP8BlockScalesFusedMoEMethod() + return DeepSeekFP8BlockScalesFusedMoEMethod() elif self.quant_config.layer_quant_mode.has_nvfp4(): return NVFP4CutlassFusedMoEMethod() elif self.quant_config.layer_quant_mode.is_int4_weight_only_per_group( @@ -603,6 +599,50 @@ def run_moe( Returns: final_hidden_states: Output tensor from MoE computation """ + # SM120 + FP8 block scales: use Triton kernel (CUTLASS TMA fails on SM120 + # for large token counts due to cuTensorMapEncodeTiled limitations). + if self.has_deepseek_fp8_block_scales and get_sm_version() == 120: + from .fused_moe_triton_fp8_block_scale import \ + run_triton_fp8_block_scale_moe + _use_alltoall = (enable_alltoall if enable_alltoall is not None else + self.enable_alltoall) + # forward_chunk sets token_final_scales=None when + # apply_router_weight_on_input=True (weights already folded into x); + # substitute ones so the Triton kernel's per-token scaling is a no-op. + if token_final_scales is None: + token_final_scales = torch.ones_like(token_selected_experts, + dtype=torch.float32) + # token_selected_experts contains GLOBAL expert IDs in the non-alltoall + # path (slot_start .. slot_end-1 for this rank's local experts, plus + # IDs for other ranks). The Triton kernel operates on LOCAL IDs + # (0 .. expert_size_per_partition-1), so remap and zero-scale any + # non-local token-expert pairs to suppress their contribution. + local_n = self.expert_size_per_partition + if _use_alltoall: + # After alltoall dispatch, IDs are already local; padding = local_n + local_ids = token_selected_experts.clamp(0, local_n - 1) + is_local = token_selected_experts < local_n + else: + slot_start = self.slot_start + local_ids = (token_selected_experts - slot_start).clamp( + 0, local_n - 1) + is_local = ((token_selected_experts >= slot_start) + & (token_selected_experts < slot_start + local_n)) + local_scales = token_final_scales * is_local.to( + token_final_scales.dtype) + result = run_triton_fp8_block_scale_moe( + x, + local_ids, + local_scales, + self.w3_w1_weight, + self.quant_scales.fc_weight_scales, + self.w2_weight, + self.quant_scales.proj_weight_scales, + activation_type=self.activation_type, + output_dtype=output_dtype, + ) + return result + # Pad input for mxfp4 alignment (128-aligned hidden_size). # Done here rather than in quantize_input so that dispatch sends # unpadded tensors and avoids NVLink workspace overallocation. diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py new file mode 100644 index 000000000000..c043857cda05 --- /dev/null +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton_fp8_block_scale.py @@ -0,0 +1,603 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 NVIDIA Corporation. All rights reserved. +""" +Triton FP8 block-scale MoE forward pass for SM120. +""" + +import torch +import triton +import triton.language as tl + +_BLOCK_SHAPE = [128, 128] # [group_k, group_n] for FP8 block-scale +_BLOCK_SIZE_M = 16 + + +@triton.jit +def _moe_histogram_kernel( + flat_ids_ptr, + counts_ptr, + N, + BLOCK_N: tl.constexpr, +): + """Build a histogram of selected experts.""" + pid = tl.program_id(0) + offs = pid * BLOCK_N + tl.arange(0, BLOCK_N) + mask = offs < N + ids = tl.load(flat_ids_ptr + offs, mask=mask, other=0).to(tl.int32) + ones = tl.full([BLOCK_N], 1, dtype=tl.int32) + tl.atomic_add(counts_ptr + ids, ones, mask=mask) + + +@triton.jit +def _int32_add(a, b): + return a + b + + +@triton.jit +def _moe_prefix_kernel( + counts_ptr, + out_off_ptr, + num_post_pad_ptr, + BLOCK_SIZE: tl.constexpr, + NUM_EXPERTS: tl.constexpr, +): + """Compute exclusive offsets from padded expert counts.""" + e = tl.arange(0, NUM_EXPERTS) + counts = tl.load(counts_ptr + e).to(tl.int32) + padded = ((counts + BLOCK_SIZE - 1) // BLOCK_SIZE) * BLOCK_SIZE + # Inclusive scan; subtract self for exclusive prefix. + inc = tl.associative_scan(padded, 0, _int32_add) + out_off = inc - padded + tl.store(out_off_ptr + e, out_off) + tl.store(num_post_pad_ptr, tl.sum(padded).to(tl.int32)) + + +@triton.jit +def _moe_expert_ids_kernel( + counts_ptr, + out_off_ptr, + expert_ids_ptr, + num_experts, + BLOCK_SIZE: tl.constexpr, +): + """Fill the expert id for each padded token block.""" + e = tl.program_id(0) + if e >= num_experts: + return + cnt = tl.load(counts_ptr + e).to(tl.int32) + off = tl.load(out_off_ptr + e).to(tl.int32) + padded = ((cnt + BLOCK_SIZE - 1) // BLOCK_SIZE) * BLOCK_SIZE + start_blk = off // BLOCK_SIZE + n_blk = padded // BLOCK_SIZE + for i in range(n_blk): + tl.store(expert_ids_ptr + start_blk + i, e) + + +@triton.jit +def _moe_scatter_kernel( + flat_ids_ptr, + counters_ptr, + sorted_token_ids_ptr, + N, + BLOCK_N: tl.constexpr, +): + """Scatter token indices into their expert ranges.""" + pid = tl.program_id(0) + offs = pid * BLOCK_N + tl.arange(0, BLOCK_N) + mask = offs < N + ids = tl.load(flat_ids_ptr + offs, mask=mask, other=0).to(tl.int32) + ones = tl.full([BLOCK_N], 1, dtype=tl.int32) + # Atomic add returns the slot reserved for this token. + pos = tl.atomic_add(counters_ptr + ids, ones, mask=mask) + tl.store(sorted_token_ids_ptr + pos, offs.to(tl.int32), mask=mask) + + +def moe_align_block_size( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Sort token-expert pairs by expert and pad each expert's row-count to + a multiple of block_size. + + The implementation is CUDA-graph compatible: all buffer sizes are Python + integers and there are no GPU-to-CPU synchronizations. + + Args: + topk_ids: (num_tokens, top_k) int64 expert IDs. + block_size: BLOCK_SIZE_M used in the GEMM kernel. + num_experts: Total number of experts. + + Returns: + Tuple of sorted token IDs, block expert IDs, and the padded token count. + """ + device = topk_ids.device + flat_ids = topk_ids.reshape(-1).to(torch.int32) + N = flat_ids.shape[0] # Python int, no GPU sync + + max_tokens_static = N + num_experts * block_size + num_m_blocks_static = N // block_size + num_experts + + counts = torch.zeros(num_experts, dtype=torch.int32, device=device) + out_offsets = torch.empty(num_experts, dtype=torch.int32, device=device) + sorted_token_ids = torch.full((max_tokens_static,), N, dtype=torch.int32, device=device) + # `expert_ids_out` does not need to be zero-initialized: _moe_expert_ids_kernel + # writes positions [0, total_n_blocks) exhaustively (per-expert ranges are + # back-to-back via the prefix scan). The allocated tail is never read. + expert_ids_out = torch.empty((num_m_blocks_static,), dtype=torch.int32, device=device) + num_post_pad = torch.empty(1, dtype=torch.int32, device=device) + + _BLOCK_N = 32 + + if N > 0: + _moe_histogram_kernel[triton.cdiv(N, _BLOCK_N),](flat_ids, counts, N, BLOCK_N=_BLOCK_N) + + _moe_prefix_kernel[1,]( + counts, + out_offsets, + num_post_pad, + BLOCK_SIZE=block_size, + NUM_EXPERTS=num_experts, + ) + + _moe_expert_ids_kernel[num_experts,]( + counts, + out_offsets, + expert_ids_out, + num_experts=num_experts, + BLOCK_SIZE=block_size, + ) + + if N > 0: + counters = out_offsets.clone() + _moe_scatter_kernel[triton.cdiv(N, _BLOCK_N),]( + flat_ids, counters, sorted_token_ids, N, BLOCK_N=_BLOCK_N + ) + + return sorted_token_ids, expert_ids_out, num_post_pad + + +@triton.jit +def _write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + compute_type: tl.constexpr, +): + acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=compute_type) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, acc, mask=c_mask) + + +def _build_moe_gemm_autotune_configs(): + """Config grid for `_fused_moe_bf16act_fp8w_kernel`. + + Constraints: + - BLOCK_SIZE_M is FIXED at 16. The MoE alignment kernels + (`_moe_expert_ids_kernel` etc.) lay out `expert_ids_out` with one + expert per 16-token chunk (the `_BLOCK_SIZE_M` constant). The GEMM + loads `expert_ids[pid_m]` assuming the same chunk size, so + mismatched values produce wrong outputs (autotune cannot detect + this; it benchmarks runtime, not numerics). + - BLOCK_SIZE_K <= group_k (=128): the kernel loads ONE scale per + K-tile, so larger K-tiles would alias multiple block scales. + """ + configs = [] + for BN in (64, 128, 256): + for BK in (64, 128): + for GROUP_M in (1, 4, 8): + for num_stages in (3, 4): + for num_warps in (4, 8): + configs.append( + triton.Config( + { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": BN, + "BLOCK_SIZE_K": BK, + "GROUP_SIZE_M": GROUP_M, + }, + num_stages=num_stages, + num_warps=num_warps, + ) + ) + return configs + + +@triton.autotune( + configs=_build_moe_gemm_autotune_configs(), + key=["N", "K"], + # FUSE_TOPK_REDUCE uses atomic_add into c_ptr, so autotune trials must start + # from zero instead of accumulating on the previous trial's output. + reset_to_zero=["c_ptr"], +) +@triton.jit +def _fused_moe_bf16act_fp8w_kernel( + a_ptr, + b_ptr, + c_ptr, + b_scale_ptr, + topk_weights_ptr, + sorted_token_ids_ptr, + expert_ids_ptr, + num_tokens_post_padded_ptr, + N, + K, + EM, + num_valid_tokens, + stride_am, + stride_ak, + stride_be, + stride_bk, + stride_bn, + stride_cm, + stride_cn, + stride_bse, + stride_bsk, + stride_bsn, + group_n: tl.constexpr, + group_k: tl.constexpr, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + top_k: tl.constexpr, + compute_type: tl.constexpr, + MUL_ROUTED_WEIGHT: tl.constexpr, + FUSE_TOPK_REDUCE: tl.constexpr, + REDUCE_TOPK: tl.constexpr, +): + """ + Fused MoE GEMM with BF16 activations and FP8 block-scaled weights. + + For weight-only quantized models: activations never go through FP8, so + there is no per-layer quantization error. Each FP8 weight tile is + dequanted to BF16 inside the K-loop using the per-block scale: + C[m, n] += dot(A_bf16[m, :], B_fp8[:, n].to(bf16)) * B_scale[n//gn, k//gk] + """ + pid = tl.program_id(0) + num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + (pid % num_pid_in_group) % group_size_m + pid_n = (pid % num_pid_in_group) // group_size_m + + offs = tl.arange(0, BLOCK_SIZE_M).to(tl.int64) + num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr) + if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded: + return + + offs_token_id = pid_m * BLOCK_SIZE_M + offs + offs_token = tl.load(sorted_token_ids_ptr + offs_token_id).to(tl.int64) + token_mask = offs_token < num_valid_tokens + + off_experts = tl.load(expert_ids_ptr + pid_m).to(tl.int64) + if off_experts == -1: + if FUSE_TOPK_REDUCE: + return + _write_zeros_to_output( + c_ptr, + stride_cm, + stride_cn, + pid_n, + N, + offs_token, + token_mask, + BLOCK_SIZE_M, + BLOCK_SIZE_N, + compute_type, + ) + return + + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int64)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + + a_ptrs = a_ptr + offs_token[:, None] // top_k * stride_am + offs_k[None, :] * stride_ak + b_ptrs = ( + b_ptr + off_experts * stride_be + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn + ) + + offs_bsn = offs_bn // group_n + b_scale_ptrs = b_scale_ptr + off_experts * stride_bse + offs_bsn * stride_bsn + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load( + a_ptrs, + mask=token_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K), + other=0.0, + ) + b_fp8 = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + b = b_fp8.to(tl.bfloat16) + + k_start = k * BLOCK_SIZE_K + offs_ks = k_start // group_k + b_scale = tl.load(b_scale_ptrs + offs_ks * stride_bsk) + + accumulator += tl.dot(a, b, out_dtype=tl.float32) * b_scale[None, :] + + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + + if MUL_ROUTED_WEIGHT: + moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0) + accumulator *= moe_weight[:, None] + + accumulator = accumulator.to(compute_type) + + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + if FUSE_TOPK_REDUCE: + out_row = offs_token // REDUCE_TOPK + c_ptrs = c_ptr + stride_cm * out_row[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.atomic_add(c_ptrs, accumulator, mask=c_mask) + else: + c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :] + c_mask = token_mask[:, None] & (offs_cn[None, :] < N) + tl.store(c_ptrs, accumulator, mask=c_mask) + + +def _invoke_bf16act_fp8w_moe_kernel( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + B_scale: torch.Tensor, + topk_weights: torch.Tensor, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + mul_routed_weight: bool, + top_k: int, + fuse_topk_reduce: bool = False, + reduce_topk: int = 1, +) -> None: + """ + Launch _fused_moe_bf16act_fp8w_kernel. + + Tensor layouts (default, fuse_topk_reduce=False): + A : (num_tokens_or_expanded, K) BF16, contiguous + B : (E, N, K) FP8, contiguous + C : (batch, top_k_orig, N) BF16, stride(1)=N, stride(2)=1 + B_scale : (E, N//128, K//128) F32 + topk_weights: (num_tokens, top_k_orig) F32 (flat-indexable by offs_token) + + When fuse_topk_reduce=True (used for the second MoE GEMM), each program + atomically accumulates its expert-weighted contribution into output row + `offs_token // reduce_topk`, eliminating the standalone reduction kernel. + The caller is responsible for zero-initializing C, which then has shape + (num_tokens, N). + """ + EM = sorted_token_ids.shape[0] + N = B.shape[1] + K = B.shape[2] + group_n, group_k = _BLOCK_SHAPE[1], _BLOCK_SHAPE[0] + + # C is 3D by default and 2D when fusing the top-k reduce. + if fuse_topk_reduce: + stride_cm, stride_cn = C.stride(0), C.stride(1) + else: + stride_cm, stride_cn = C.stride(1), C.stride(2) + + def grid(META): + return (triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),) + + _fused_moe_bf16act_fp8w_kernel[grid]( + A, + B, + C, + B_scale, + topk_weights, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + N, + K, + EM, + A.shape[0] * top_k, + A.stride(0), + A.stride(1), + B.stride(0), + B.stride(2), + B.stride(1), + stride_cm, + stride_cn, + B_scale.stride(0) if B_scale.ndim >= 2 else 0, + B_scale.stride(2) if B_scale.ndim == 3 else 0, + B_scale.stride(1) if B_scale.ndim >= 2 else 0, + group_n=group_n, + group_k=group_k, + top_k=top_k, + compute_type=tl.bfloat16, + MUL_ROUTED_WEIGHT=mul_routed_weight, + FUSE_TOPK_REDUCE=fuse_topk_reduce, + REDUCE_TOPK=reduce_topk, + ) + + +@triton.jit +def _gated_activation_kernel( + ic1_ptr, + ic2_ptr, + M, + INTER, + BLOCK_M: tl.constexpr, + BLOCK_I: tl.constexpr, + ACTIVATION: tl.constexpr, +): + """Compute ic2[m, i] = activation(gate) * up.""" + pid_m = tl.program_id(0) + pid_i = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_i = pid_i * BLOCK_I + tl.arange(0, BLOCK_I) + mask = (offs_m[:, None] < M) & (offs_i[None, :] < INTER) + + row_stride = 2 * INTER + up_ptrs = ic1_ptr + offs_m[:, None] * row_stride + offs_i[None, :] + gate_ptrs = up_ptrs + INTER + + up = tl.load(up_ptrs, mask=mask, other=0.0).to(tl.float32) + gate = tl.load(gate_ptrs, mask=mask, other=0.0).to(tl.float32) + + if ACTIVATION == 0: + act = gate * (1.0 / (1.0 + tl.exp(-gate))) + else: + act = 0.5 * gate * (1.0 + tl.erf(gate * 0.7071067811865475)) + + out = (act * up).to(tl.bfloat16) + out_ptrs = ic2_ptr + offs_m[:, None] * INTER + offs_i[None, :] + tl.store(out_ptrs, out, mask=mask) + + +def _invoke_gated_activation_kernel( + ic1: torch.Tensor, + ic2: torch.Tensor, + activation: int, +) -> None: + """Launch `_gated_activation_kernel`. + + Args: + ic1: (M, 2*INTER) bfloat16, contiguous (up half | gate half). + ic2: (M, INTER) bfloat16, contiguous (output, may be uninitialized). + activation: 0 for silu/swiglu, 1 for gelu/geglu (exact). + """ + assert ic1.is_contiguous() and ic2.is_contiguous() + assert ic1.dtype == torch.bfloat16 and ic2.dtype == torch.bfloat16 + M, two_i = ic1.shape + M2, inter = ic2.shape + assert M == M2 and two_i == 2 * inter + + BLOCK_M = 16 + BLOCK_I = 128 if inter >= 128 else triton.next_power_of_2(inter) + grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(inter, BLOCK_I)) + _gated_activation_kernel[grid]( + ic1, + ic2, + M, + inter, + BLOCK_M=BLOCK_M, + BLOCK_I=BLOCK_I, + ACTIVATION=activation, + ) + + +def run_triton_fp8_block_scale_moe( + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + w3_w1: torch.Tensor, + w3_w1_scales: torch.Tensor, + w2: torch.Tensor, + w2_scales: torch.Tensor, + activation_type: int, + output_dtype: torch.dtype | None = None, +) -> torch.Tensor: + """ + Complete MoE forward pass using Triton FP8 block-scale GEMMs. + + This is a drop-in replacement for torch.ops.trtllm.fused_moe() for + SM120 + FP8_BLOCK_SCALES, matching the weight layout produced by + DeepSeekFP8BlockScalesFusedMoEMethod. + + Args: + x: (T, H) BF16 input activations. + token_selected_experts: (T, K) int64 expert IDs per token. + token_final_scales: (T, K) F32 routing weights. + w3_w1: (E, 2I, H) FP8 gate+up projection weights. + w3_w1_scales: (E, Nb, Kb) F32 block scales for w3_w1. + w2: (E, H, I) FP8 down-projection weights. + w2_scales: (E, Nb, Kb) F32 block scales for w2. + activation_type: ActivationType enum value. + output_dtype: Optional output dtype. + + Returns: + (T, H) output tensor. + """ + from tensorrt_llm._torch.modules.fused_moe.interface import ActivationType + + num_tokens, hidden = x.shape + num_experts, gate_up_size, _ = w3_w1.shape + intermediate = gate_up_size // 2 + top_k = token_selected_experts.shape[1] + device = x.device + + act_int = int(activation_type) + swiglu = int(ActivationType.Swiglu) + geglu = int(ActivationType.Geglu) + silu = int(ActivationType.Silu) + gelu = int(ActivationType.Gelu) + + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + token_selected_experts, _BLOCK_SIZE_M, num_experts + ) + + ic1 = torch.empty((num_tokens, top_k, gate_up_size), dtype=torch.bfloat16, device=device) + _invoke_bf16act_fp8w_moe_kernel( + x.contiguous(), + w3_w1, + ic1, + w3_w1_scales, + token_final_scales, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + mul_routed_weight=False, + top_k=top_k, + ) + + if act_int in (swiglu, silu): + activation_id = 0 + elif act_int in (geglu, gelu): + activation_id = 1 + else: + raise ValueError(f"Unsupported activation_type={act_int} in Triton FP8 block-scale MoE") + + ic2 = torch.empty((num_tokens * top_k, intermediate), dtype=torch.bfloat16, device=device) + _invoke_gated_activation_kernel( + ic1.view(num_tokens * top_k, gate_up_size), + ic2, + activation_id, + ) + + # The kernel constexpr `top_k=1` here means each ic2 row is already one + # (token, expert) pair (not to be confused with the model's `top_k`). + fuse_topk_reduce = top_k > 1 + if fuse_topk_reduce: + output = torch.zeros((num_tokens, hidden), dtype=torch.bfloat16, device=device) + gemm2_out = output + else: + # The unfused path uses the 3D C layout expected by the GEMM wrapper. + gemm2_out = torch.empty((num_tokens, 1, hidden), dtype=torch.bfloat16, device=device) + + _invoke_bf16act_fp8w_moe_kernel( + ic2, + w2, + gemm2_out, + w2_scales, + token_final_scales, + sorted_token_ids, + expert_ids, + num_tokens_post_padded, + mul_routed_weight=True, + top_k=1, + fuse_topk_reduce=fuse_topk_reduce, + reduce_topk=top_k, + ) + + if not fuse_topk_reduce: + output = gemm2_out.squeeze(1) + + if output_dtype is not None and output.dtype != output_dtype: + output = output.to(output_dtype) + return output diff --git a/tensorrt_llm/llmapi/llm_utils.py b/tensorrt_llm/llmapi/llm_utils.py index 1d362a96d63c..5061efa171de 100644 --- a/tensorrt_llm/llmapi/llm_utils.py +++ b/tensorrt_llm/llmapi/llm_utils.py @@ -507,12 +507,53 @@ def _update_from_hf_quant_config(self) -> bool: raise ValueError( f"Unsupported weights_quant_strategy: {weights_quant_strategy}. " "Supported strategies: 'channel', 'block'.") + elif (weights_quant_config["num_bits"] == 4 + and weights_quant_config.get("type") == "float" + and weights_quant_strategy == "tensor_group"): + # llm-compressor NVFP4: weights FP4 with FP8 per-group + # scales (group_size=16), scaled by an FP32 global scale. + if inputs_quant_strategy != "tensor_group": + raise ValueError( + f"Unsupported inputs_quant_strategy for NVFP4: {inputs_quant_strategy}." + ) + group_size = weights_quant_config["group_size"] + if group_size != 16: + raise ValueError( + f"Unsupported group_size: {group_size}. Supported: 16 for NVFP4." + ) + quant_config.quant_algo = QuantAlgo.NVFP4 + quant_config.group_size = group_size else: raise ValueError( f"Unsupported quant_bits: {weights_quant_config['num_bits']}. " - "Supported: 8.") + "Supported: 8 (FP8) or 4 (NVFP4).") + + # kv_cache_scheme (llm-compressor): FP8 per-tensor KV cache. + kv_cache_scheme = hf_quant_config.get("kv_cache_scheme") + if kv_cache_scheme is not None: + if (kv_cache_scheme.get("num_bits") == 8 + and kv_cache_scheme.get("type") == "float"): + if quant_config.kv_cache_quant_algo in (None, + QuantAlgo.FP8): + quant_config.kv_cache_quant_algo = QuantAlgo.FP8 + else: + raise ValueError( + f"Specified kv_cache_quant_algo={quant_config.kv_cache_quant_algo}, " + f"conflicting with FP8 KV cache from HF quant config." + ) + else: + raise ValueError( + f"Unsupported kv_cache_scheme: {kv_cache_scheme}.") - quant_config.exclude_modules = hf_quant_config.get("ignore", []) + hf_exclude_modules = hf_quant_config.get( + "modules_to_not_convert", None) + if hf_exclude_modules is not None: + quant_config.exclude_modules = list( + set(hf_exclude_modules + + hf_quant_config.get("ignore", []))) + else: + quant_config.exclude_modules = hf_quant_config.get( + "ignore", []) elif hf_quant_config.get("quant_method") == "nvfp4": quant_config.quant_algo = QuantAlgo.NVFP4 group_size = hf_quant_config.get("group_size", 16) @@ -527,6 +568,10 @@ def _update_from_hf_quant_config(self) -> bool: dict.fromkeys(hf_exclude_modules + default_exclude)) else: quant_config.exclude_modules = default_exclude + elif hf_quant_config.get("quant_method") is None: + # quantization_config present but quant_method is null → not + # pre-quantized; tell caller so calibration is not skipped. + return False else: raise NotImplementedError( f"Unsupported quantization_config: {hf_quant_config}.") diff --git a/tensorrt_llm/llmapi/reasoning_parser.py b/tensorrt_llm/llmapi/reasoning_parser.py index 7c2563e91cb4..77c3fda95fa3 100644 --- a/tensorrt_llm/llmapi/reasoning_parser.py +++ b/tensorrt_llm/llmapi/reasoning_parser.py @@ -96,6 +96,7 @@ def finish(self) -> ReasoningParserResult: @register_reasoning_parser("deepseek-r1", reasoning_at_start=True) +@register_reasoning_parser("laguna") @register_reasoning_parser("qwen3") @register_reasoning_parser("minimax_m2", reasoning_at_start=True) @register_reasoning_parser("minimax_m2_append_think", reasoning_at_start=True) @@ -196,6 +197,7 @@ def parse_delta(self, delta_text: str) -> ReasoningParserResult: "qwen3_next": "qwen3", "deepseek_v3": "deepseek-r1", "deepseek_v32": "deepseek-r1", + "laguna": "laguna", "nemotron_h": "nano-v3", "nemotron_h_puzzle": "nano-v3", "gemma4": "gemma4", diff --git a/tensorrt_llm/models/modeling_utils.py b/tensorrt_llm/models/modeling_utils.py index c9487a2a39d2..66ff49d8db4c 100644 --- a/tensorrt_llm/models/modeling_utils.py +++ b/tensorrt_llm/models/modeling_utils.py @@ -213,13 +213,9 @@ def _requires_modelopt_quantization(self): return False def _get_quant_cfg(self, module_name=None): - if self.exclude_modules is not None: - for exclude_module in self.exclude_modules: - if exclude_module == module_name or ( - exclude_module.endswith('*') - and module_name.startswith(exclude_module[:-1])): - return LayerQuantConfig(quant_algo=None, - quantized_layers={}) + if (module_name is not None + and self.is_module_excluded_from_quantization(module_name)): + return LayerQuantConfig(quant_algo=None, quantized_layers={}) return self def _get_modelopt_qformat(self): @@ -254,9 +250,10 @@ def is_module_excluded_from_quantization(self, name: str) -> bool: """Check if the module is excluded from quantization. A module is excluded if its own name or any ancestor (split on - ``.``) matches an entry in ``exclude_modules`` via ``fnmatch``. - The ancestor walk means listing a parent module (without a glob - suffix) implicitly excludes all of its children. + ``.``) matches an entry in ``exclude_modules`` via ``fnmatch`` or + a ``re:`` prefixed regex. The ancestor walk means listing a parent + module (without a glob suffix) implicitly excludes all of its + children. Args: name (str): The name of the module. @@ -269,7 +266,10 @@ def is_module_excluded_from_quantization(self, name: str) -> bool: candidate = name while True: for exclude_module in self.exclude_modules: - if fnmatch.fnmatchcase(candidate, exclude_module): + if exclude_module.startswith("re:"): + if re.fullmatch(exclude_module[3:], candidate): + return True + elif fnmatch.fnmatchcase(candidate, exclude_module): return True if '.' not in candidate: return False diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 13accb492c50..fc916c64b04b 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -415,6 +415,14 @@ bigcode/starcoder2-7b: - accuracy: 26.5 bigcode/starcoder2-15b: - accuracy: 54.5 +poolside/laguna-XS.2: + - accuracy: 85.71 + - quant_algo: FP8_BLOCK_SCALES + kv_cache_quant_algo: FP8 + accuracy: 87.15 + - quant_algo: NVFP4 + kv_cache_quant_algo: FP8 + accuracy: 85.14 mistral/Mistral-Large-3-675B: - accuracy: 86.1 - spec_dec_algo: Eagle diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml index 9851e35aa577..85ec0a7a556d 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -407,6 +407,14 @@ GPT-OSS/MXFP4: accuracy: 75.50 - quant_algo: W4A8_MXFP4_FP8 accuracy: 75.50 +poolside/laguna-XS.2: + - accuracy: 74.93 + - quant_algo: FP8_BLOCK_SCALES + kv_cache_quant_algo: FP8 + accuracy: 75.19 + - quant_algo: NVFP4 + kv_cache_quant_algo: FP8 + accuracy: 73.66 mistralai/Mistral-Nemo-12b-Base: - accuracy: 69.66 - quant_algo: FP8 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index c8e62ba7f02f..680989722674 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6338,6 +6338,60 @@ def test_dflash(self): extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) +class TestLagunaXS(LlmapiAccuracyTestHarness): + MODEL_NAME = "poolside/laguna-XS.2" + MODEL_PATH = "/tmp/laguna-minimal-Laguna-XS.2" + FP8_MODEL_PATH = "/tmp/laguna-minimal-Laguna-XS.2-FP8" + NVFP4_MODEL_PATH = "/tmp/laguna-minimal-Laguna-XS.2-NVFP4" + MAX_SEQ_LEN = 4096 + MAX_NUM_TOKENS = 4096 + MAX_BATCH_SIZE = 128 + + def _run_accuracy(self, + model_path, + expected_quant_algo=None, + expected_kv_cache_quant_algo=None): + if not os.path.exists(model_path): + pytest.skip(f"Model directory {model_path} does not exist") + + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.9) + + with LLM(model_path, + max_seq_len=self.MAX_SEQ_LEN, + max_num_tokens=self.MAX_NUM_TOKENS, + max_batch_size=self.MAX_BATCH_SIZE, + kv_cache_config=kv_cache_config) as llm: + if expected_quant_algo is not None: + assert llm.args.quant_config.quant_algo == expected_quant_algo + if expected_kv_cache_quant_algo is not None: + assert (llm.args.quant_config.kv_cache_quant_algo == + expected_kv_cache_quant_algo) + + task = MMLU(self.MODEL_NAME) + task.evaluate(llm) + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + + @skip_pre_hopper + @pytest.mark.skip_less_device_memory(80000) + def test_bf16(self): + self._run_accuracy(self.MODEL_PATH) + + @skip_pre_hopper + @pytest.mark.skip_less_device_memory(80000) + def test_fp8(self): + self._run_accuracy(self.FP8_MODEL_PATH, + expected_quant_algo=QuantAlgo.FP8_BLOCK_SCALES, + expected_kv_cache_quant_algo=QuantAlgo.FP8) + + @skip_pre_blackwell + @pytest.mark.skip_less_device_memory(80000) + def test_nvfp4(self): + self._run_accuracy(self.NVFP4_MODEL_PATH, + expected_quant_algo=QuantAlgo.NVFP4, + expected_kv_cache_quant_algo=QuantAlgo.FP8) + + @pytest.mark.skip_less_device_memory(80000) class TestQwen3_5_35B_A3B(LlmapiAccuracyTestHarness): MODEL_NAME = "Qwen/Qwen3.5-35B-A3B" diff --git a/tests/unittest/_torch/modules/fused_moe/test_triton_fp8_block_scale.py b/tests/unittest/_torch/modules/fused_moe/test_triton_fp8_block_scale.py new file mode 100644 index 000000000000..6db886fe4dbe --- /dev/null +++ b/tests/unittest/_torch/modules/fused_moe/test_triton_fp8_block_scale.py @@ -0,0 +1,203 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +""" +Unit tests for ``run_triton_fp8_block_scale_moe``. + +Run as a pytest: + pytest tests/unittest/_torch/modules/fused_moe/test_triton_fp8_block_scale.py -v +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import pytest +import torch +import torch.nn.functional as F + +from tensorrt_llm._torch.modules.fused_moe.fused_moe_triton_fp8_block_scale import ( + run_triton_fp8_block_scale_moe, +) +from tensorrt_llm._torch.modules.fused_moe.interface import ActivationType +from tensorrt_llm._utils import get_sm_version +from tests.unittest._torch.helpers import calc_diff, per_block_cast_to_fp8_e8m0 + +skip_unsupported = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() != 120, + reason="Requires CUDA SM120 for the Triton FP8 block-scale MoE path", +) + + +def _run_moe_bf16_reference( + x: torch.Tensor, + token_selected_experts: torch.Tensor, + token_final_scales: torch.Tensor, + w3_w1: torch.Tensor, + w3_w1_scales: torch.Tensor, + w2: torch.Tensor, + w2_scales: torch.Tensor, + activation_type: int, + output_dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Dequantize FP8 weights and run the MoE in BF16.""" + num_tokens, hidden = x.shape + _, gate_up_size, _ = w3_w1.shape + intermediate = gate_up_size // 2 + top_k = token_selected_experts.shape[1] + device = x.device + + tok_exp = token_selected_experts.long() + tok_wt = token_final_scales + + def dequant_weight(w_fp8: torch.Tensor, scales: torch.Tensor) -> torch.Tensor: + num_experts, out_features, in_features = w_fp8.shape + num_n_blocks, num_k_blocks = scales.shape[1], scales.shape[2] + out = torch.zeros( + num_experts, + out_features, + in_features, + dtype=torch.bfloat16, + device=device, + ) + for expert_idx in range(num_experts): + for n_block in range(num_n_blocks): + for k_block in range(num_k_blocks): + n_slice = slice(n_block * 128, (n_block + 1) * 128) + k_slice = slice(k_block * 128, (k_block + 1) * 128) + block = w_fp8[expert_idx, n_slice, k_slice].float() + out[expert_idx, n_slice, k_slice] = ( + block * scales[expert_idx, n_block, k_block] + ).bfloat16() + return out + + w31_bf16 = dequant_weight(w3_w1, w3_w1_scales) + w2_bf16 = dequant_weight(w2, w2_scales) + + output = torch.zeros(num_tokens, hidden, dtype=torch.bfloat16, device=device) + for t in range(num_tokens): + acc = torch.zeros(hidden, dtype=torch.float32, device=device) + for k in range(top_k): + e = int(tok_exp[t, k]) + gate_up = x[t].float() @ w31_bf16[e].float().T + u, g = gate_up[:intermediate], gate_up[intermediate:] + act_int = int(activation_type) + if act_int in (int(ActivationType.Swiglu), int(ActivationType.Silu)): + ic2 = F.silu(g) * u + else: + ic2 = F.gelu(g) * u + acc += (ic2 @ w2_bf16[e].float().T) * float(tok_wt[t, k]) + output[t] = acc.bfloat16() + + if output_dtype is not None and output.dtype != output_dtype: + output = output.to(output_dtype) + return output + + +@dataclass(frozen=True) +class MoeShape: + name: str + num_tokens: int + hidden: int + intermediate: int + num_experts: int + top_k: int + + +# Laguna XS.2 routed-MoE dimensions. The shared expert is outside this kernel. +LAGUNA_HIDDEN = 2048 +LAGUNA_INTER = 512 +LAGUNA_NUM_EXPERTS = 256 +LAGUNA_TOP_K = 8 + +SHAPES = [ + MoeShape("tiny", num_tokens=64, hidden=256, intermediate=512, num_experts=8, top_k=2), + MoeShape("tiny_topk1", num_tokens=64, hidden=256, intermediate=512, num_experts=8, top_k=1), + MoeShape( + "laguna_corr", + num_tokens=64, + hidden=LAGUNA_HIDDEN, + intermediate=LAGUNA_INTER, + num_experts=LAGUNA_NUM_EXPERTS, + top_k=LAGUNA_TOP_K, + ), + MoeShape( + "laguna_bench", + num_tokens=4096, + hidden=LAGUNA_HIDDEN, + intermediate=LAGUNA_INTER, + num_experts=LAGUNA_NUM_EXPERTS, + top_k=LAGUNA_TOP_K, + ), +] + +# The BF16 reference is intentionally simple and slow, so keep the large shape +# available for manual experiments but out of the correctness parametrization. +CORRECTNESS_SHAPES = [s for s in SHAPES if s.name != "laguna_bench"] + + +def _make_inputs( + shape: MoeShape, + *, + device: str = "cuda", + seed: int = 0, + activation_type: int = int(ActivationType.Swiglu), +) -> dict[str, torch.Tensor | int]: + g = torch.Generator(device=device).manual_seed(seed) + num_tokens = shape.num_tokens + hidden = shape.hidden + inter = shape.intermediate + num_experts = shape.num_experts + top_k = shape.top_k + + x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=device, generator=g) * 0.5 + + gate_logits = torch.randn((num_tokens, num_experts), device=device, generator=g) + topk_vals, topk_ids = gate_logits.topk(top_k, dim=-1) + token_selected_experts = topk_ids.to(torch.int64) + token_final_scales = torch.softmax(topk_vals.float(), dim=-1) + + w3_w1_bf16 = ( + torch.randn( + (num_experts, 2 * inter, hidden), dtype=torch.bfloat16, device=device, generator=g + ) + * 0.05 + ) + w2_bf16 = ( + torch.randn((num_experts, hidden, inter), dtype=torch.bfloat16, device=device, generator=g) + * 0.05 + ) + w3_w1_fp8, w3_w1_sf = per_block_cast_to_fp8_e8m0(w3_w1_bf16) + w2_fp8, w2_sf = per_block_cast_to_fp8_e8m0(w2_bf16) + + return dict( + x=x, + token_selected_experts=token_selected_experts, + token_final_scales=token_final_scales, + w3_w1=w3_w1_fp8, + w3_w1_scales=w3_w1_sf, + w2=w2_fp8, + w2_scales=w2_sf, + activation_type=activation_type, + ) + + +ACTIVATIONS = [ + pytest.param(int(ActivationType.Swiglu), id="swiglu"), + pytest.param(int(ActivationType.Geglu), id="geglu"), +] + + +@skip_unsupported +@pytest.mark.parametrize("shape", CORRECTNESS_SHAPES, ids=lambda s: s.name) +@pytest.mark.parametrize("activation_type", ACTIVATIONS) +def test_fp8_matches_bf16_reference(shape: MoeShape, activation_type: int) -> None: + inputs = _make_inputs(shape, activation_type=activation_type) + + out_fp8 = run_triton_fp8_block_scale_moe(**inputs) + out_ref = _run_moe_bf16_reference(**inputs) + + assert out_fp8.shape == out_ref.shape == (shape.num_tokens, shape.hidden) + assert out_fp8.dtype == torch.bfloat16 + + diff = calc_diff(out_fp8, out_ref).item() + assert diff < 1e-2, f"shape={shape.name} cosine-diff={diff:.4e} exceeds 1e-2" From ab779ce17ba7bf3a5fc80cdd26403efc123b6675 Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Mon, 11 May 2026 11:55:22 +0000 Subject: [PATCH 2/9] fix: Changes to work without overlay and with new transformers Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- laguna_minimal_overlay.sh | 73 ------------------- tensorrt_llm/_torch/configs/__init__.py | 20 +++-- tensorrt_llm/_torch/configs/laguna.py | 21 ++++++ tensorrt_llm/_torch/models/modeling_laguna.py | 20 ++++- .../_torch/pyexecutor/config_utils.py | 1 + 5 files changed, 54 insertions(+), 81 deletions(-) delete mode 100755 laguna_minimal_overlay.sh create mode 100644 tensorrt_llm/_torch/configs/laguna.py diff --git a/laguna_minimal_overlay.sh b/laguna_minimal_overlay.sh deleted file mode 100755 index 9a0bf23319f0..000000000000 --- a/laguna_minimal_overlay.sh +++ /dev/null @@ -1,73 +0,0 @@ -#!/usr/bin/env bash -# Minimal transformers-4.57-compat overlay for a Laguna checkpoint directory. -# -# Laguna's `configuration_laguna.py` imports two transformers >= 4.58 symbols: -# - `transformers.configuration_utils.PreTrainedConfig` -# - `transformers.modeling_rope_utils.RopeParameters` -# Neither exists in transformers 4.57, which is what we have installed. -# -# This overlay symlinks every file in the checkpoint dir, then replaces just -# `configuration_laguna.py` with a copy that is prefixed with a compat shim -# aliasing the missing symbols. No other files are modified — additional -# fixes (config.json rewrites, etc.) are deliberately out of scope. -# -# Usage: -# laguna_minimal_overlay.sh [overlay_dir] -# Echoes the overlay directory path on stdout. -set -euo pipefail - -if [[ $# -lt 1 ]]; then - echo "Usage: $(basename "$0") [overlay_dir]" >&2 - exit 1 -fi - -SRC="$(realpath "$1")" -OVERLAY="${2:-/tmp/laguna-minimal-$(basename "$SRC")}" - -if [[ ! -d "$SRC" ]]; then - echo "ERROR: source dir does not exist: $SRC" >&2 - exit 1 -fi -if [[ ! -f "$SRC/configuration_laguna.py" ]]; then - echo "ERROR: $SRC has no configuration_laguna.py" >&2 - exit 1 -fi - -rm -rf "$OVERLAY" -mkdir -p "$OVERLAY" - -for entry in "$SRC"/*; do - name="$(basename "$entry")" - if [[ "$name" == "configuration_laguna.py" ]]; then - continue - fi - ln -s "$entry" "$OVERLAY/$name" -done - -SHIM_HEADER='# --- compat shim (added by laguna_minimal_overlay.sh) --- -try: - from transformers.configuration_utils import PreTrainedConfig # transformers >= 4.58 -except ImportError: - from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig # transformers 4.57 -try: - from transformers.modeling_rope_utils import RopeParameters # transformers >= 4.58 -except ImportError: - RopeParameters = dict # type: ignore[misc,assignment] -# Neuter `@strict` (huggingface_hub) — newer Laguna configs decorate a non-dataclass -# subclass of PretrainedConfig with it, which raises in our hf-hub version. -def strict(cls=None, **_kwargs): # type: ignore[no-redef] - return cls if cls is not None else (lambda c: c) -# Neuter `@auto_docstring` (transformers utils) — only adds docstrings, safe to skip. -def auto_docstring(*_args, **_kwargs): # type: ignore[no-redef] - return (lambda c: c) if not _args or not callable(_args[0]) else _args[0] -# --- end compat shim --- -' - -# Strip the original imports we just shimmed; everything else is kept verbatim. -{ - printf '%s' "$SHIM_HEADER" - grep -vE '^from transformers\.configuration_utils import PreTrainedConfig$|^from transformers\.modeling_rope_utils import RopeParameters$|^from huggingface_hub\.dataclasses import strict$|^from transformers\.utils import auto_docstring$' \ - "$SRC/configuration_laguna.py" -} > "$OVERLAY/configuration_laguna.py" - -echo "$OVERLAY" diff --git a/tensorrt_llm/_torch/configs/__init__.py b/tensorrt_llm/_torch/configs/__init__.py index 6496e3283451..c4d9fb25ecfd 100644 --- a/tensorrt_llm/_torch/configs/__init__.py +++ b/tensorrt_llm/_torch/configs/__init__.py @@ -1,27 +1,33 @@ from tensorrt_llm._torch.configs.deepseek_v3 import DeepseekV3Config +from tensorrt_llm._torch.configs.laguna import LagunaConfig def _register_custom_configs_with_transformers() -> None: # Make AutoConfig.from_pretrained / AutoTokenizer.from_pretrained accept # model_types that TRT-LLM understands but upstream transformers does not - # (DeepSeek-V3.2 and Kimi K2 both ship config.json with these model_types - # and rely on TRT-LLM's local DeepseekV3Config workaround). + # (DeepSeek-V3.2, Kimi K2, and Laguna ship config.json with these + # model_types and rely on TRT-LLM's local config workarounds). # # Without this, transformers 5.5.x falls back to a bare PreTrainedConfig # that lacks attributes like `max_position_embeddings`, and # AutoTokenizer.from_pretrained then raises AttributeError before any # tokenizer can be constructed. Bypass AutoConfig.register's model_type - # consistency check (DeepseekV3Config.model_type is "deepseek_v3") by - # writing into the underlying mapping directly. + # consistency check for aliases (for example, DeepseekV3Config.model_type + # is "deepseek_v3") by writing into the underlying mapping directly. from transformers.models.auto.configuration_auto import CONFIG_MAPPING - for model_type in ("deepseek_v32", "kimi_k2"): + custom_configs = { + "deepseek_v32": DeepseekV3Config, + "kimi_k2": DeepseekV3Config, + "laguna": LagunaConfig, + } + for model_type, config_class in custom_configs.items(): if model_type in CONFIG_MAPPING: continue - CONFIG_MAPPING.register(model_type, DeepseekV3Config, exist_ok=True) + CONFIG_MAPPING.register(model_type, config_class, exist_ok=True) _register_custom_configs_with_transformers() del _register_custom_configs_with_transformers -__all__ = ["DeepseekV3Config"] +__all__ = ["DeepseekV3Config", "LagunaConfig"] diff --git a/tensorrt_llm/_torch/configs/laguna.py b/tensorrt_llm/_torch/configs/laguna.py new file mode 100644 index 000000000000..9ddaa35d79c3 --- /dev/null +++ b/tensorrt_llm/_torch/configs/laguna.py @@ -0,0 +1,21 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from transformers.configuration_utils import PretrainedConfig + + +class LagunaConfig(PretrainedConfig): + model_type = "laguna" + keys_to_ignore_at_inference = ["past_key_values"] diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py index 36b75c854efe..3cbe1e2b02af 100644 --- a/tensorrt_llm/_torch/models/modeling_laguna.py +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -14,6 +14,7 @@ # limitations under the License. """Laguna / Laguna-XS model for TensorRT-LLM PyTorch backend.""" +import math from typing import Dict, List, Optional, Type import torch @@ -216,6 +217,18 @@ class LagunaAttention(QKNormRoPEAttention): """Laguna attention with per-head softplus gating, per-layer heads, dual RoPE, and sliding window.""" + @staticmethod + def _yarn_mscale_from_attention_factor(factor: float, attention_factor: float) -> float: + if factor <= 1.0: + if attention_factor != 1.0: + raise ValueError( + "Laguna YaRN RoPE cannot encode attention_factor != 1.0 " + f"when factor <= 1.0: factor={factor}, " + f"attention_factor={attention_factor}" + ) + return 0.0 + return (attention_factor - 1.0) / (0.1 * math.log(factor)) + def __init__( self, model_config, @@ -301,7 +314,12 @@ def _build_rope_from_flat_dict(config, rp_dict: dict) -> RopeParams: rp.scale = float(rp_dict.get("factor", 1.0)) rp.beta_fast = float(rp_dict.get("beta_fast", 32.0)) rp.beta_slow = float(rp_dict.get("beta_slow", 1.0)) - rp.mscale = float(rp_dict.get("attention_factor", 1.0)) + attention_factor = rp_dict.get("attention_factor") + if attention_factor is not None: + rp.mscale = LagunaAttention._yarn_mscale_from_attention_factor( + rp.scale, float(attention_factor) + ) + rp.mscale_all_dim = 0.0 rp.original_max_positions = int( rp_dict.get("original_max_position_embeddings", config.max_position_embeddings) ) diff --git a/tensorrt_llm/_torch/pyexecutor/config_utils.py b/tensorrt_llm/_torch/pyexecutor/config_utils.py index f4ec50639715..ef270040fc26 100644 --- a/tensorrt_llm/_torch/pyexecutor/config_utils.py +++ b/tensorrt_llm/_torch/pyexecutor/config_utils.py @@ -427,6 +427,7 @@ def __getitem__(self, key): deepseek_v32="DeepseekV3Config", kimi_k2="DeepseekV3Config", glm_moe_dsa="DeepseekV3Config", + laguna="LagunaConfig", ) # NOTE: HF config.json uses deepseek_v32 as model_type but with same DSV3 config class From 62de961db618821ff584b6d0f8b3c9c6836cb9dd Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Mon, 11 May 2026 12:44:33 +0000 Subject: [PATCH 3/9] fix: CI test paths and rope un-needed changes Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_laguna.py | 18 +----------------- .../defs/accuracy/test_llm_api_pytorch.py | 6 +++--- 2 files changed, 4 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py index 3cbe1e2b02af..9b77ab1758f1 100644 --- a/tensorrt_llm/_torch/models/modeling_laguna.py +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -14,7 +14,6 @@ # limitations under the License. """Laguna / Laguna-XS model for TensorRT-LLM PyTorch backend.""" -import math from typing import Dict, List, Optional, Type import torch @@ -217,18 +216,6 @@ class LagunaAttention(QKNormRoPEAttention): """Laguna attention with per-head softplus gating, per-layer heads, dual RoPE, and sliding window.""" - @staticmethod - def _yarn_mscale_from_attention_factor(factor: float, attention_factor: float) -> float: - if factor <= 1.0: - if attention_factor != 1.0: - raise ValueError( - "Laguna YaRN RoPE cannot encode attention_factor != 1.0 " - f"when factor <= 1.0: factor={factor}, " - f"attention_factor={attention_factor}" - ) - return 0.0 - return (attention_factor - 1.0) / (0.1 * math.log(factor)) - def __init__( self, model_config, @@ -316,10 +303,7 @@ def _build_rope_from_flat_dict(config, rp_dict: dict) -> RopeParams: rp.beta_slow = float(rp_dict.get("beta_slow", 1.0)) attention_factor = rp_dict.get("attention_factor") if attention_factor is not None: - rp.mscale = LagunaAttention._yarn_mscale_from_attention_factor( - rp.scale, float(attention_factor) - ) - rp.mscale_all_dim = 0.0 + rp.mscale = float(attention_factor) rp.original_max_positions = int( rp_dict.get("original_max_position_embeddings", config.max_position_embeddings) ) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 680989722674..65e9d03d5823 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -6340,9 +6340,9 @@ def test_dflash(self): class TestLagunaXS(LlmapiAccuracyTestHarness): MODEL_NAME = "poolside/laguna-XS.2" - MODEL_PATH = "/tmp/laguna-minimal-Laguna-XS.2" - FP8_MODEL_PATH = "/tmp/laguna-minimal-Laguna-XS.2-FP8" - NVFP4_MODEL_PATH = "/tmp/laguna-minimal-Laguna-XS.2-NVFP4" + MODEL_PATH = f"{llm_models_root()}/Laguna-XS.2" + FP8_MODEL_PATH = f"{llm_models_root()}/Laguna-XS.2-FP8" + NVFP4_MODEL_PATH = f"{llm_models_root()}/Laguna-XS.2-NVFP4" MAX_SEQ_LEN = 4096 MAX_NUM_TOKENS = 4096 MAX_BATCH_SIZE = 128 From 76100a02b0135deb96f9cccb05598b27d3cc684f Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Mon, 11 May 2026 12:54:05 +0000 Subject: [PATCH 4/9] test: update CI yaml files Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml index 81afed83e2ca..0ceba9c3bc7f 100644 --- a/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml +++ b/tests/integration/test_lists/test-db/l0_rtx_pro_6000.yml @@ -33,6 +33,9 @@ l0_rtx_pro_6000: - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-30B-A3B_fp8_hf-Qwen3/saved_models_Qwen3-30B-A3B_fp8_hf] # 3mins - test_e2e.py::test_ptp_quickstart_advanced[Qwen3-30B-A3B_nvfp4_hf-Qwen3/saved_models_Qwen3-30B-A3B_nvfp4_hf] # 2mins - test_e2e.py::test_ptp_quickstart_advanced[GPT-OSS-20B-gpt_oss/gpt-oss-20b] + - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_bf16 + - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 + - accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_nvfp4 - condition: ranges: From fb7bb3735533609cf1ee7cd2b82b72ae159f9534 Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Fri, 15 May 2026 16:34:33 +0000 Subject: [PATCH 5/9] fix: apply fix from Poolside Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_laguna.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py index 9b77ab1758f1..fb7373c52050 100644 --- a/tensorrt_llm/_torch/models/modeling_laguna.py +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -260,6 +260,7 @@ def __init__( bias=getattr(config, "qkv_bias", False) or getattr(config, "attention_bias", False), pos_embd_params=pos_embd_params, fuse_qk_norm_rope=False, + rope_fusion=False, layer_idx=layer_idx, dtype=config.torch_dtype, dense_bias=False, From 53f99835659431af464178ba2ed51af91332188c Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Mon, 18 May 2026 11:52:38 +0100 Subject: [PATCH 6/9] test: Add additional test coverage Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- .../checkpoints/laguna/test_weight_mapper.py | 122 ++++++++++++++++++ tests/unittest/_torch/test_hf_quant_config.py | 60 +++++++++ .../unittest/llmapi/test_reasoning_parser.py | 43 ++++++ tests/unittest/models/test_modeling_utils.py | 41 ++++++ 4 files changed, 266 insertions(+) create mode 100644 tests/unittest/_torch/models/checkpoints/laguna/test_weight_mapper.py create mode 100644 tests/unittest/_torch/test_hf_quant_config.py create mode 100644 tests/unittest/models/test_modeling_utils.py diff --git a/tests/unittest/_torch/models/checkpoints/laguna/test_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/laguna/test_weight_mapper.py new file mode 100644 index 000000000000..8964bacab06a --- /dev/null +++ b/tests/unittest/_torch/models/checkpoints/laguna/test_weight_mapper.py @@ -0,0 +1,122 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from tensorrt_llm._torch.models.modeling_laguna import LagunaHfWeightMapper +from tensorrt_llm._torch.modules.fused_moe.interface import MoE + + +class _FakeQuantMode: + def __init__(self, *, nvfp4: bool = False, fp8_block_scales: bool = False): + self._nvfp4 = nvfp4 + self._fp8_block_scales = fp8_block_scales + + def has_nvfp4(self): + return self._nvfp4 + + def has_fp8_block_scales(self): + return self._fp8_block_scales + + +class _FakeMoE(MoE): + @classmethod + def can_implement(cls, *args, **kwargs): + return True, None + + def __init__(self): + nn.Module.__init__(self) + self.loaded_weights = None + self.allow_partial_loading = None + + def load_weights(self, weights, allow_partial_loading: bool = False): + self.loaded_weights = weights + self.allow_partial_loading = allow_partial_loading + + +def test_laguna_hf_weight_mapper_preprocesses_nvfp4_weights(): + mapper = LagunaHfWeightMapper() + mapper._config = SimpleNamespace( + quant_config=SimpleNamespace(quant_mode=_FakeQuantMode(nvfp4=True)) + ) + + packed_weight = torch.tensor([1, 2], dtype=torch.uint8) + correction_bias = torch.tensor([3.0]) + weights = { + "model.layers.0.mlp.experts.e_score_correction_bias": correction_bias, + "model.layers.0.mlp.experts.0.gate_proj.weight_packed": packed_weight, + "model.layers.0.mlp.experts.0.gate_proj.weight_global_scale": torch.tensor([2.0, 0.0]), + "model.layers.0.mlp.experts.0.gate_proj.input_global_scale": torch.tensor([4.0]), + } + + result = mapper.preprocess_weights(weights) + + assert result["model.layers.0.mlp.gate.e_score_correction_bias"] is correction_bias + assert result["model.layers.0.mlp.experts.0.gate_proj.weight"] is packed_weight + torch.testing.assert_close( + result["model.layers.0.mlp.experts.0.gate_proj.weight_scale_2"], + torch.tensor([0.5, 0.0]), + ) + torch.testing.assert_close( + result["model.layers.0.mlp.experts.0.gate_proj.input_scale"], + torch.tensor([0.25]), + ) + + +@pytest.mark.parametrize( + ("fp8_block_scales", "scale_key"), + [ + (False, "0.w1.weight_scale"), + (True, "0.w1.weight_scale_inv"), + ], +) +def test_laguna_hf_weight_mapper_handles_special_moe_module(fp8_block_scales: bool, scale_key: str): + mapper = LagunaHfWeightMapper() + mapper._config = SimpleNamespace( + quant_config=SimpleNamespace(quant_mode=_FakeQuantMode(fp8_block_scales=fp8_block_scales)) + ) + module = _FakeMoE() + + gate_weight = torch.tensor([1.0]) + up_weight = torch.tensor([2.0]) + down_weight = torch.tensor([3.0]) + weight_scale = torch.tensor([4.0]) + module_weights = { + "experts.0.gate_proj.weight": gate_weight, + "experts.0.up_proj.weight": up_weight, + "experts.0.down_proj.weight": down_weight, + "experts.0.gate_proj.weight_scale": weight_scale, + } + + mapper.handle_special_instance_module( + module, + "model.layers.0.mlp.experts", + module_weights, + allow_partial_loading=True, + ) + + assert module.loaded_weights == [ + { + "0.w1.weight": gate_weight, + "0.w3.weight": up_weight, + "0.w2.weight": down_weight, + scale_key: weight_scale, + } + ] + assert module.allow_partial_loading is True diff --git a/tests/unittest/_torch/test_hf_quant_config.py b/tests/unittest/_torch/test_hf_quant_config.py new file mode 100644 index 000000000000..958d17e3ebc9 --- /dev/null +++ b/tests/unittest/_torch/test_hf_quant_config.py @@ -0,0 +1,60 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm.models.modeling_utils import QuantAlgo + + +def _compressed_tensors_nvfp4_config(**overrides): + config = { + "quant_method": "compressed-tensors", + "config_groups": { + "group_0": { + "weights": { + "num_bits": 4, + "type": "float", + "strategy": "tensor_group", + "group_size": 16, + }, + "input_activations": { + "strategy": "tensor_group", + }, + }, + }, + } + config.update(overrides) + return config + + +def test_load_hf_quant_config_parses_nvfp4_with_kv_cache_scheme(): + gate_exclude = "re:model\\.layers\\.\\d+\\.mlp\\.gate" + hf_quant_config = _compressed_tensors_nvfp4_config( + kv_cache_scheme={ + "num_bits": 8, + "type": "float", + }, + modules_to_not_convert=[gate_exclude], + ignore=["lm_head"], + ) + + quant_config, layer_quant_config = ModelConfig.load_hf_quant_config( + hf_quant_config, moe_backend="CUTLASS" + ) + + assert layer_quant_config is None + assert quant_config.quant_algo == QuantAlgo.NVFP4 + assert quant_config.group_size == 16 + assert quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + assert set(quant_config.exclude_modules) == {gate_exclude, "lm_head"} diff --git a/tests/unittest/llmapi/test_reasoning_parser.py b/tests/unittest/llmapi/test_reasoning_parser.py index 0af8fbc326ec..05e720a4363e 100644 --- a/tests/unittest/llmapi/test_reasoning_parser.py +++ b/tests/unittest/llmapi/test_reasoning_parser.py @@ -248,6 +248,39 @@ def test_qwen3_reasoning_parser_stream(delta_texts: list, content: list, assert result.reasoning_content == reasoning_context[i] +@pytest.mark.parametrize(("text", "content", "reasoning_context"), [ + ("abc", "c", "b"), + ("ab", "b", "a"), + ("a", "", "a"), + ("a", "a", ""), + ("", "", ""), +]) +def test_laguna_reasoning_parser(text: str, content: str, + reasoning_context: str): + reasoning_parser = ReasoningParserFactory.create_reasoning_parser("laguna") + result = reasoning_parser.parse(text) + assert result.content == content + assert result.reasoning_content == reasoning_context + + +@pytest.mark.parametrize(("delta_texts", "content", "reasoning_context"), [ + (["a", "lr", "b"], ["", "r", "b"], ["a", "l", ""]), + (["ab"], ["", "b"], ["", "a"]), + (["ab"], ["", "b"], ["a", ""]), + (["", "ab"], ["", "b"], ["", "a"]), + (["a", "b"], ["", "b"], ["a", ""]), + (["ab"], ["", "", "b" + ], ["a", "", ""]), +]) +def test_laguna_reasoning_parser_stream(delta_texts: list, content: list, + reasoning_context: list): + reasoning_parser = ReasoningParserFactory.create_reasoning_parser("laguna") + for i, delta_text in enumerate(delta_texts): + result = reasoning_parser.parse_delta(delta_text) + assert result.content == content[i] + assert result.reasoning_content == reasoning_context[i] + + TOOL_CALL = "" TOOL_CALL_END = "" @@ -612,6 +645,16 @@ def test_auto_detect_gemma4(tmp_path): assert result == "gemma4" +def test_auto_detect_laguna(tmp_path): + """Laguna model → 'laguna' parser.""" + model_dir = str(tmp_path / "Laguna") + os.makedirs(model_dir) + _write_config(model_dir, "laguna") + + result = resolve_auto_reasoning_parser(model_dir) + assert result == "laguna" + + # --------------------------------------------------------------------------- # Gemma 4 reasoning parser tests # --------------------------------------------------------------------------- diff --git a/tests/unittest/models/test_modeling_utils.py b/tests/unittest/models/test_modeling_utils.py new file mode 100644 index 000000000000..cdf7f14afd9a --- /dev/null +++ b/tests/unittest/models/test_modeling_utils.py @@ -0,0 +1,41 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + +def test_quant_config_exclude_modules_regex_match(): + quant_config = QuantConfig( + quant_algo=QuantAlgo.NVFP4, + exclude_modules=[ + r"re:model\.layers\.[02468]\.mixer", + r"re:model\.layers\.\d+\.self_attn\.(q|k)_proj", + ], + ) + + for name in ( + "model.layers.2.mixer", + "model.layers.2.mixer.down_proj", + "model.layers.10.self_attn.q_proj", + "model.layers.10.self_attn.q_proj.weight", + "model.layers.10.self_attn.k_proj", + ): + assert quant_config.is_module_excluded_from_quantization(name), name + for name in ( + "model.layers.3.mixer", + "model.layers.10.self_attn.v_proj", + "model.layers.10.self_attn.o_proj", + ): + assert not quant_config.is_module_excluded_from_quantization(name), name From 378cb80e97fb1b00dfeb3eaa6cc5616ac826be97 Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Mon, 18 May 2026 15:51:56 +0000 Subject: [PATCH 7/9] fix: Another patch from Poolside Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_laguna.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py index fb7373c52050..f47b79c529fc 100644 --- a/tensorrt_llm/_torch/models/modeling_laguna.py +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -285,6 +285,7 @@ def __init__( dtype=config.torch_dtype, mapping=model_config.mapping, tensor_parallel_mode=g_tp_mode, + quant_config=model_config.get_quant_config(), ) @staticmethod From bb720c6cd4bdab963b9e58dd5b5d47ceb313b7ca Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Mon, 18 May 2026 17:05:18 +0000 Subject: [PATCH 8/9] Respond to reviewer comments Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- tensorrt_llm/_torch/configs/laguna.py | 2 ++ tensorrt_llm/_torch/models/modeling_laguna.py | 7 ++++--- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/configs/laguna.py b/tensorrt_llm/_torch/configs/laguna.py index 9ddaa35d79c3..7c39ce3d7983 100644 --- a/tensorrt_llm/_torch/configs/laguna.py +++ b/tensorrt_llm/_torch/configs/laguna.py @@ -16,6 +16,8 @@ from transformers.configuration_utils import PretrainedConfig +# This is a workaround until we use a version of transformers that supports +# Laguna. Namely: v5.7.0 or higher. class LagunaConfig(PretrainedConfig): model_type = "laguna" keys_to_ignore_at_inference = ["past_key_values"] diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py index f47b79c529fc..03c82f67a9fb 100644 --- a/tensorrt_llm/_torch/models/modeling_laguna.py +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -185,7 +185,7 @@ def forward( ) if self.routed_scaling_factor != 1.0: - routed_output = routed_output * self.routed_scaling_factor + routed_output *= self.routed_scaling_factor if self.shared_expert is not None: shared_output = self.shared_expert(hidden_states) @@ -246,7 +246,6 @@ def __init__( gating = getattr(config, "gating", False) self._use_gating = bool(gating) self._gate_per_head = gating == "per-head" or gating is True - self._num_heads_local = None # set after super().__init__ # fuse_qk_norm_rope=False is required: the fused kernel reads # partial_rotary_factor and yarn params from pretrained_config @@ -270,6 +269,7 @@ def __init__( self._num_heads_local = self.num_heads + self.g_proj = None if self._use_gating: g_out = num_heads if self._gate_per_head else num_heads * self.head_dim # In DEP mode (enable_attention_dp=True) the attention base class uses @@ -386,7 +386,7 @@ def forward( * gate.unsqueeze(-1) ).view(shape) else: - attn_output = attn_output * gate + attn_output *= gate attn_output = _helix_cp_output_projection( self.o_proj, @@ -521,6 +521,7 @@ def __init__(self, model_config): dtype=config.torch_dtype, ) + @torch.inference_mode() def forward( self, attn_metadata: AttentionMetadata, From c7b44f1e69940f5ad14c881ee25280da0160dd61 Mon Sep 17 00:00:00 2001 From: Dom Brown <3886319+DomBrown@users.noreply.github.com> Date: Mon, 18 May 2026 17:10:17 +0000 Subject: [PATCH 9/9] Update QA lists per reviewer comment Signed-off-by: Dom Brown <3886319+DomBrown@users.noreply.github.com> --- tests/integration/test_lists/qa/llm_function_core.txt | 3 +++ tests/integration/test_lists/qa/llm_function_rtx6k.txt | 3 +++ 2 files changed, 6 insertions(+) diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index 6ac7f64897b9..322d21cf6e8e 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -555,6 +555,9 @@ accuracy/test_llm_api_pytorch.py::TestKimiK25::test_nvfp4[tp8_attn_dp] accuracy/test_llm_api_pytorch.py::TestKimiK2::test_fp8_blockscale[latency] accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[4gpus] accuracy/test_llm_api_pytorch.py::TestKimiK2::test_nvfp4[8gpus] +accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_bf16 +accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 +accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_nvfp4 accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_nvfp4 accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_nvfp4_streaming[stream_interval_4] accuracy/test_llm_api_pytorch.py::TestLlama3_1_8B::test_nvfp4_streaming[stream_interval_64] diff --git a/tests/integration/test_lists/qa/llm_function_rtx6k.txt b/tests/integration/test_lists/qa/llm_function_rtx6k.txt index d31b7cf8bdd3..25e7eabd5da7 100644 --- a/tests/integration/test_lists/qa/llm_function_rtx6k.txt +++ b/tests/integration/test_lists/qa/llm_function_rtx6k.txt @@ -143,6 +143,9 @@ accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4a16[dp4-auto] accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_w4_chunked_prefill[cutlass-auto] accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_trtllm] accuracy/test_llm_api_pytorch.py::TestQwen3_235B_A22B::test_nvfp4[latency_moe_trtllm_eagle3] +accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_bf16 +accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_fp8 +accuracy/test_llm_api_pytorch.py::TestLagunaXS::test_nvfp4 accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_nvfp4_multi_gpus_sm120[throughput_tp8] accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v1_kv_cache-cutlass-one_model-overlap_scheduler] accuracy/test_llm_api_pytorch.py::TestGPTOSS::test_eagle3_4gpus[v2_kv_cache-cutlass-one_model-overlap_scheduler]