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/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..7c39ce3d7983 --- /dev/null +++ b/tensorrt_llm/_torch/configs/laguna.py @@ -0,0 +1,23 @@ +# 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 + + +# 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/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..03c82f67a9fb --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_laguna.py @@ -0,0 +1,683 @@ +# 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 *= 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 + + # 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, + rope_fusion=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 + + 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 + # 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, + quant_config=model_config.get_quant_config(), + ) + + @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)) + attention_factor = rp_dict.get("attention_factor") + if attention_factor is not None: + rp.mscale = float(attention_factor) + 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 *= 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, + ) + + @torch.inference_mode() + 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/_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 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..65e9d03d5823 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 = 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 + + 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/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] 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: 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/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" 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