From 3c4acdc98e19bedbe3fde8783e36219544ddeea2 Mon Sep 17 00:00:00 2001 From: alyosha-swamy Date: Thu, 16 Apr 2026 03:59:35 +0000 Subject: [PATCH 1/8] Add AFMoE support Signed-off-by: alyosha-swamy --- tensorrt_llm/_torch/models/__init__.py | 2 + .../_torch/models/checkpoints/__init__.py | 7 +- .../checkpoints/hf/afmoe_weight_mapper.py | 42 ++ tensorrt_llm/_torch/models/modeling_afmoe.py | 510 ++++++++++++++++++ .../_torch/modeling/test_modeling_afmoe.py | 398 ++++++++++++++ 5 files changed, 956 insertions(+), 3 deletions(-) create mode 100644 tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py create mode 100644 tensorrt_llm/_torch/models/modeling_afmoe.py create mode 100644 tests/unittest/_torch/modeling/test_modeling_afmoe.py diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index 9c3b032421b2..3c070732ce7e 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -5,6 +5,7 @@ # under transformers >= 5.5; see _torch/configs/__init__.py. import tensorrt_llm._torch.configs # noqa: F401 +from .modeling_afmoe import AfmoeForCausalLM from .modeling_auto import AutoModelForCausalLM from .modeling_bert import BertForSequenceClassification from .modeling_clip import CLIPVisionModel @@ -51,6 +52,7 @@ # Note: for better readiblity, this should have same order as imports above __all__ = [ + "AfmoeForCausalLM", "AutoModelForCausalLM", "BertForSequenceClassification", "CLIPVisionModel", diff --git a/tensorrt_llm/_torch/models/checkpoints/__init__.py b/tensorrt_llm/_torch/models/checkpoints/__init__.py index f4094417b3ca..b10a93132ca8 100644 --- a/tensorrt_llm/_torch/models/checkpoints/__init__.py +++ b/tensorrt_llm/_torch/models/checkpoints/__init__.py @@ -1,4 +1,5 @@ from .base_checkpoint_loader import BaseCheckpointLoader +from .hf.afmoe_weight_mapper import AfmoeHfWeightMapper from .hf.checkpoint_loader import HfCheckpointLoader from .hf.config_loader import HfConfigLoader from .hf.gemma3_weight_mapper import Gemma3HfWeightMapper @@ -24,9 +25,9 @@ from .mx.checkpoint_loader import MXCheckpointLoader __all__ = [ - "HfConfigLoader", "HfWeightLoader", "HfWeightMapper", "MistralConfigLoader", - "MistralWeightMapper", "MistralCheckpointLoader", "BaseCheckpointLoader", - "HfCheckpointLoader", "NemotronHHfWeightMapper", + "AfmoeHfWeightMapper", "HfConfigLoader", "HfWeightLoader", "HfWeightMapper", + "MistralConfigLoader", "MistralWeightMapper", "MistralCheckpointLoader", + "BaseCheckpointLoader", "HfCheckpointLoader", "NemotronHHfWeightMapper", "NemotronNasHfWeightMapper", "Gemma3HfWeightMapper", "MixtralHfWeightMapper", "Llama4HfWeightMapper", "Qwen2MoeHfWeightMapper", "Qwen3MoeHfWeightMapper", "Qwen2VLHfWeightMapper", diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py new file mode 100644 index 000000000000..5d045239c7b2 --- /dev/null +++ b/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py @@ -0,0 +1,42 @@ +from torch import nn + +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper +from tensorrt_llm._torch.models.modeling_utils import register_mapper +from tensorrt_llm._torch.modules.fused_moe.interface import MoE + + +@register_mapper("HF", "AfmoeForCausalLM") +class AfmoeHfWeightMapper(HfWeightMapper): + def __init__(self): + super().__init__() + + self.params_map = { + # MoE expert weights: gate_proj->w1, up_proj->w3, down_proj->w2 + r"(.*experts\.\d+\.)gate_proj(.*)": r"\1w1\2", + r"(.*experts\.\d+\.)up_proj(.*)": r"\1w3\2", + r"(.*experts\.\d+\.)down_proj(.*)": r"\1w2\2", + # HF router weight path -> TRT-LLM gate path + r"(.*)\.router\.gate\.(.*)": r"\1.gate.\2", + # expert_bias -> gate.e_score_correction_bias + r"(.*)\.mlp\.expert_bias(.*)": r"\1.mlp.gate.e_score_correction_bias\2", + } + + def preprocess_weights(self, weights: dict) -> dict: + weights = self.rename_by_params_map(self.params_map, weights) + return weights + + def is_special_instance_module(self, module: nn.Module) -> bool: + return isinstance(module, MoE) + + def handle_special_instance_module( + self, + module: nn.Module, + module_name: str, + module_weights: dict, + allow_partial_loading: bool = False, + ) -> None: + if isinstance(module, MoE): + module.load_weights( + weights=[module_weights], + allow_partial_loading=allow_partial_loading, + ) diff --git a/tensorrt_llm/_torch/models/modeling_afmoe.py b/tensorrt_llm/_torch/models/modeling_afmoe.py new file mode 100644 index 000000000000..439da67f9b77 --- /dev/null +++ b/tensorrt_llm/_torch/models/modeling_afmoe.py @@ -0,0 +1,510 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Inference-only AFMoE (Arcee Foundation MoE) for TensorRT-LLM. + +Follows the HF implementation of AfmoeForCausalLM. + +Key architectural features: + - Per-layer attention type (sliding_attention vs global) + - Q/K RMSNorm in attention + - Gated attention output (sigmoid gate) + - RoPE only on local (sliding-window) attention layers + - Dense MLP for early layers, MoE with shared experts for later layers + - 4 layer norms per decoder block (pre/post attention, pre/post MLP) + - Optional muP embedding scaling +""" + +from typing import Dict, List, Optional + +import torch +from torch import nn +from transformers import AutoConfig, PretrainedConfig + +from tensorrt_llm.functional import PositionEmbeddingType + +from ...logger import logger +from ..attention_backend import AttentionMetadata +from ..attention_backend.interface import ( + PositionalEmbeddingParams, + PredefinedAttentionMask, + RopeParams, +) +from ..distributed import AllReduce +from ..model_config import ModelConfig +from ..modules.attention import Attention +from ..modules.decoder_layer import DecoderLayer +from ..modules.embedding import Embedding +from ..modules.fused_moe import DeepSeekV3MoeRoutingMethod, create_moe +from ..modules.fused_moe.routing import Deepseekv3RoutingImpl +from ..modules.gated_mlp import GatedMLP +from ..modules.linear import Linear, TensorParallelMode +from ..modules.rms_norm import RMSNorm +from ..utils import AuxStreamType +from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model + + +class AfmoeConfig(PretrainedConfig): + model_type = "afmoe" + + +logger.warning_once( + "transformers does not natively support 'AfmoeConfig'. " + "Registering AfmoeConfig so AutoConfig can load AFMoE checkpoints.", + key="AFMOE_REGISTER_WARNING", +) +AutoConfig.register(AfmoeConfig.model_type, AfmoeConfig) + + +def _validate_routing_config(config: PretrainedConfig) -> None: + """Validate that the routing config matches our Deepseekv3RoutingImpl assumptions.""" + score_func = getattr(config, "scoring_func", getattr(config, "score_func", "sigmoid")) + if score_func != "sigmoid": + raise ValueError( + f"AFMoE implementation uses sigmoid scoring via " + f"Deepseekv3RoutingImpl, but config has " + f"scoring_func={score_func!r}. Only 'sigmoid' is supported." + ) + + norm_topk = getattr(config, "norm_topk_prob", getattr(config, "route_norm", True)) + if not norm_topk: + raise ValueError( + "AFMoE implementation assumes normalized top-k probabilities " + "(norm_topk_prob=True / route_norm=True), but config disables it." + ) + + +class AfmoeGate(nn.Module): + """Router gate for AFMoE, following the DeepSeekV3 grouped top-k pattern.""" + + def __init__( + self, + hidden_size: int, + num_experts: int, + top_k: int, + n_group: int, + topk_group: int, + route_scale: float, + dtype: Optional[torch.dtype] = None, + ): + super().__init__() + self.weight = nn.Parameter( + torch.empty((num_experts, hidden_size), dtype=dtype), + requires_grad=False, + ) + self.e_score_correction_bias = nn.Parameter( + torch.empty(num_experts, dtype=torch.float32), + requires_grad=False, + ) + self.routing_impl = Deepseekv3RoutingImpl( + top_k=top_k, + n_group=n_group, + topk_group=topk_group, + routed_scaling_factor=route_scale, + is_fused=True, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + logits = torch.ops.trtllm.dsv3_router_gemm_op( + hidden_states, + self.weight.t(), + bias=None, + out_dtype=torch.float32, + ) + return logits + + def load_weights(self, weights: List[Dict]): + assert len(weights) == 1 + self.weight.copy_(weights[0]["weight"][:]) + self.e_score_correction_bias.copy_( + weights[0]["e_score_correction_bias"][:].to(self.e_score_correction_bias.dtype) + ) + + @property + def routing_method(self) -> DeepSeekV3MoeRoutingMethod: + return DeepSeekV3MoeRoutingMethod( + top_k=self.routing_impl.top_k, + n_group=self.routing_impl.n_group, + topk_group=self.routing_impl.topk_group, + routed_scaling_factor=self.routing_impl.routed_scaling_factor, + is_fused=self.routing_impl.is_fused, + callable_e_score_correction_bias=lambda: self.e_score_correction_bias, + ) + + +class AfmoeMoE(nn.Module): + """MoE layer with shared experts for AFMoE. + + Both routed experts and shared experts produce TP-partial results + (reduce_results=False / reduce_output=False). After summing them + we perform a single AllReduce so that each rank holds the full + hidden-state, matching the DeepSeekV3 MoE pattern. + """ + + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + aux_stream: torch.cuda.Stream, + layer_idx: Optional[int] = None, + ): + 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.num_shared_experts = getattr(config, "num_shared_experts", 0) + self.enable_attention_dp = model_config.mapping.enable_attention_dp + + self.gate = AfmoeGate( + hidden_size=self.hidden_dim, + num_experts=self.num_experts, + top_k=self.top_k, + n_group=config.n_group, + topk_group=config.topk_group, + route_scale=getattr(config, "route_scale", 1.0), + dtype=config.torch_dtype, + ) + + 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={AuxStreamType.MoeChunkingOverlap: aux_stream}, + dtype=config.torch_dtype, + reduce_results=False, + model_config=model_config, + layer_idx=layer_idx, + ) + + if self.num_shared_experts > 0: + shared_intermediate = config.moe_intermediate_size * self.num_shared_experts + self.shared_experts = GatedMLP( + hidden_size=self.hidden_dim, + intermediate_size=shared_intermediate, + bias=False, + dtype=config.torch_dtype, + config=model_config, + reduce_output=False, + layer_idx=layer_idx, + ) + else: + self.shared_experts = None + + self.mapping = model_config.mapping + + self.allreduce = None + if not self.enable_attention_dp and self.mapping.tp_size > 1: + self.allreduce = AllReduce( + mapping=model_config.mapping, + strategy=model_config.allreduce_strategy, + ) + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> torch.Tensor: + 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, + use_dp_padding=False, + ) + + if self.shared_experts is not None: + shared_output = self.shared_experts(hidden_states) + final_output = shared_output.add_(routed_output) + else: + final_output = routed_output + + if self.allreduce is not None: + final_output = self.allreduce(final_output) + + return final_output + + +class AfmoeAttention(Attention): + """Attention with Q/K norm, per-layer sliding window, gated output. + + Uses a separate gate_proj linear (not fused into QKV) to gate the + attention output with sigmoid, matching the HF checkpoint layout. + """ + + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + layer_idx: Optional[int] = None, + ): + config = model_config.pretrained_config + layer_types = getattr(config, "layer_types", []) + self.is_local_attention = ( + layer_idx is not None + and layer_idx < len(layer_types) + and layer_types[layer_idx] == "sliding_attention" + ) + self._attention_window_size = config.sliding_window if self.is_local_attention else None + + rope_params = RopeParams.from_config(config) if self.is_local_attention else None + pos_embd_params = ( + PositionalEmbeddingParams( + type=PositionEmbeddingType.rope_gpt_neox, + rope=rope_params, + ) + if self.is_local_attention + else None + ) + + super().__init__( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_key_value_heads=config.num_key_value_heads, + max_position_embeddings=getattr(config, "max_position_embeddings", 131072), + bias=False, + pos_embd_params=pos_embd_params, + layer_idx=layer_idx, + dtype=config.torch_dtype, + config=model_config, + ) + + self.q_norm = RMSNorm( + hidden_size=self.head_dim, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + self.k_norm = RMSNorm( + hidden_size=self.head_dim, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + + self.gate_proj = Linear( + config.hidden_size, + config.num_attention_heads * self.head_dim, + bias=False, + dtype=config.torch_dtype, + mapping=model_config.mapping, + tensor_parallel_mode=TensorParallelMode.COLUMN, + gather_output=False, + quant_config=model_config.get_quant_config(), + skip_create_weights_in_init=model_config.skip_create_weights_in_init, + allreduce_strategy=model_config.allreduce_strategy, + force_dynamic_quantization=model_config.force_dynamic_quantization, + use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, + ) + + def apply_rope( + self, + q: torch.Tensor, + k: Optional[torch.Tensor], + v: Optional[torch.Tensor], + position_ids: torch.Tensor, + ): + q, k, v = self.split_qkv(q, k, v) + + q_shape = q.shape + k_shape = k.shape + q = self.q_norm(q.reshape(-1, self.num_heads, self.head_dim)).reshape(q_shape) + k = self.k_norm(k.reshape(-1, self.num_key_value_heads, self.head_dim)).reshape(k_shape) + + if self.is_local_attention and not self.rope_fusion and position_ids is not None: + q, k = self.rotary_emb(position_ids, [q, k]) + + return q, k, v + + def forward( + self, + position_ids: torch.IntTensor, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + gate = self.gate_proj(hidden_states) + + qkv = self.qkv_proj(hidden_states) + q, k, v = qkv, None, None + + q, k, v = self.apply_rope(q, k, v, position_ids) + q, k, v = self.convert_qkv(q, k, v) + + attn_output = self.forward_impl( + q, + k, + v, + attn_metadata, + attention_mask=PredefinedAttentionMask.CAUSAL, + attention_window_size=self._attention_window_size, + attention_mask_data=None, + mrope_config=None, + ) + + attn_output = attn_output * torch.sigmoid(gate) + + attn_output = self.o_proj(attn_output) + return attn_output + + +class AfmoeDecoderLayer(DecoderLayer): + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + layer_idx: int, + aux_stream: torch.cuda.Stream, + ): + super().__init__() + config = model_config.pretrained_config + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + + self.self_attn = AfmoeAttention(model_config, layer_idx=layer_idx) + + num_dense_layers = getattr(config, "num_dense_layers", 0) + self.moe_enabled = layer_idx >= num_dense_layers + if self.moe_enabled: + self.mlp = AfmoeMoE(model_config, aux_stream, layer_idx=layer_idx) + else: + self.mlp = GatedMLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + bias=False, + dtype=config.torch_dtype, + config=model_config, + layer_idx=layer_idx, + ) + + 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, + ) + self.pre_mlp_layernorm = RMSNorm( + hidden_size=config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.torch_dtype, + ) + self.post_mlp_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], + **kwargs, + ) -> torch.Tensor: + 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 = self.post_attention_layernorm(hidden_states) + + hidden_states, residual = self.pre_mlp_layernorm(hidden_states, residual) + + if self.moe_enabled: + hidden_states = self.mlp(hidden_states, attn_metadata) + else: + hidden_states = self.mlp(hidden_states) + + hidden_states = self.post_mlp_layernorm(hidden_states) + + return hidden_states, residual + + +class AfmoeModel(DecoderModel): + def __init__(self, model_config: ModelConfig[PretrainedConfig]): + super().__init__(model_config) + config = model_config.pretrained_config + _validate_routing_config(config) + + self.vocab_size = config.vocab_size + self.mup_enabled = getattr(config, "mup_enabled", False) + self.hidden_size = config.hidden_size + self.aux_stream = torch.cuda.Stream() + + self.embed_tokens = Embedding( + config.vocab_size, + config.hidden_size, + dtype=config.torch_dtype, + enable_torch_compile_for_embedding=model_config.enable_torch_compile_for_embedding, + ) + + self.layers = nn.ModuleList( + [ + AfmoeDecoderLayer(model_config, layer_idx, self.aux_stream) + 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, + **kwargs, + ) -> torch.Tensor: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError( + "You cannot specify both input_ids and inputs_embeds at " + "the same time, and must specify either one" + ) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if self.mup_enabled: + inputs_embeds = inputs_embeds * (self.hidden_size**0.5) + + hidden_states = inputs_embeds + + residual = None + for decoder_layer in self.layers: + hidden_states, residual = decoder_layer( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + residual=residual, + **kwargs, + ) + + hidden_states, _ = self.norm(hidden_states, residual) + return hidden_states + + +@register_auto_model("AfmoeForCausalLM") +class AfmoeForCausalLM(DecoderModelForCausalLM[AfmoeModel, PretrainedConfig]): + def __init__(self, model_config: ModelConfig[PretrainedConfig]): + super().__init__( + AfmoeModel(model_config), + config=model_config, + hidden_size=model_config.pretrained_config.hidden_size, + vocab_size=model_config.pretrained_config.vocab_size, + ) + + def load_weights(self, weights: Dict, weight_mapper, **kwargs): + weights = weight_mapper.preprocess_weights(weights) + super().load_weights(weights=weights, weight_mapper=weight_mapper, **kwargs) diff --git a/tests/unittest/_torch/modeling/test_modeling_afmoe.py b/tests/unittest/_torch/modeling/test_modeling_afmoe.py new file mode 100644 index 000000000000..c32babd49758 --- /dev/null +++ b/tests/unittest/_torch/modeling/test_modeling_afmoe.py @@ -0,0 +1,398 @@ +import unittest +from copy import deepcopy +from unittest.mock import Mock, patch + +import torch + +import tensorrt_llm +from tensorrt_llm._torch.attention_backend.utils import get_attention_backend +from tensorrt_llm._torch.metadata import KVCacheParams +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_afmoe import ( + AfmoeConfig, + AfmoeForCausalLM, + AfmoeMoE, + _validate_routing_config, +) +from tensorrt_llm._torch.models.modeling_utils import ( + MODEL_CLASS_MAPPER_MAPPING, + MODEL_CLASS_MAPPING, +) +from tensorrt_llm._torch.modules.linear import TensorParallelMode +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig + +WINDOW_SIZE = 4 +NUM_HIDDEN_LAYERS = 4 +NUM_DENSE_LAYERS = 1 + +AFMOE_CONFIG = { + "architectures": ["AfmoeForCausalLM"], + "dtype": "bfloat16", + "hidden_size": 256, + "intermediate_size": 512, + "max_position_embeddings": 2048, + "model_type": "afmoe", + "moe_intermediate_size": 128, + "n_group": 1, + "norm_topk_prob": True, + "num_attention_heads": 8, + "num_dense_layers": NUM_DENSE_LAYERS, + "num_experts": 8, + "num_experts_per_tok": 2, + "num_hidden_layers": NUM_HIDDEN_LAYERS, + "num_key_value_heads": 2, + "num_shared_experts": 1, + "rms_norm_eps": 1e-05, + "rope_theta": 10000, + "route_scale": 1.0, + "scoring_func": "sigmoid", + "sliding_window": WINDOW_SIZE, + "layer_types": [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ], + "tie_word_embeddings": False, + "topk_group": 1, + "vocab_size": 1024, + "hidden_act": "silu", + "mup_enabled": False, +} + + +class TestAfmoeRegistry(unittest.TestCase): + """Verify AfmoeForCausalLM resolves through _torch auto-model registration.""" + + def test_auto_model_registry(self): + self.assertIn("AfmoeForCausalLM", MODEL_CLASS_MAPPING) + self.assertIs(MODEL_CLASS_MAPPING["AfmoeForCausalLM"], AfmoeForCausalLM) + + def test_weight_mapper_registry(self): + self.assertIn("AfmoeForCausalLM_HF", MODEL_CLASS_MAPPER_MAPPING) + + def test_legacy_model_map_does_not_contain_afmoe(self): + from tensorrt_llm.models import MODEL_MAP + + self.assertNotIn("AfmoeForCausalLM", MODEL_MAP) + + +class TestAfmoeRoutingValidation(unittest.TestCase): + """Verify routing assumption guards.""" + + def test_valid_sigmoid_config(self): + config = AfmoeConfig.from_dict(deepcopy(AFMOE_CONFIG)) + _validate_routing_config(config) + + def test_rejects_softmax_scoring(self): + d = deepcopy(AFMOE_CONFIG) + d["scoring_func"] = "softmax" + config = AfmoeConfig.from_dict(d) + with self.assertRaisesRegex(ValueError, "Only 'sigmoid' is supported"): + _validate_routing_config(config) + + def test_rejects_disabled_norm_topk(self): + d = deepcopy(AFMOE_CONFIG) + d["norm_topk_prob"] = False + config = AfmoeConfig.from_dict(d) + with self.assertRaisesRegex(ValueError, "norm_topk_prob"): + _validate_routing_config(config) + + def test_model_init_rejects_invalid_routing(self): + d = deepcopy(AFMOE_CONFIG) + d["scoring_func"] = "softmax" + config = AfmoeConfig.from_dict(d) + model_config = ModelConfig(pretrained_config=config) + with self.assertRaisesRegex(ValueError, "Only 'sigmoid' is supported"): + AfmoeForCausalLM(model_config) + + +class TestAfmoeWeightMapper(unittest.TestCase): + """Verify AfmoeHfWeightMapper key transformations.""" + + def setUp(self): + from tensorrt_llm._torch.models.checkpoints.hf.afmoe_weight_mapper import ( + AfmoeHfWeightMapper, + ) + + self.mapper = AfmoeHfWeightMapper() + + def test_expert_key_remapping(self): + fake_weights = { + "model.layers.1.mlp.experts.0.gate_proj.weight": torch.zeros(1), + "model.layers.1.mlp.experts.0.up_proj.weight": torch.zeros(1), + "model.layers.1.mlp.experts.0.down_proj.weight": torch.zeros(1), + "model.layers.1.mlp.experts.3.gate_proj.weight": torch.zeros(1), + "model.layers.1.mlp.experts.3.up_proj.weight": torch.zeros(1), + "model.layers.1.mlp.experts.3.down_proj.weight": torch.zeros(1), + } + result = self.mapper.preprocess_weights(fake_weights) + for expert_id in [0, 3]: + prefix = f"model.layers.1.mlp.experts.{expert_id}" + self.assertIn(f"{prefix}.w1.weight", result) + self.assertIn(f"{prefix}.w3.weight", result) + self.assertIn(f"{prefix}.w2.weight", result) + self.assertNotIn(f"{prefix}.gate_proj.weight", result) + self.assertNotIn(f"{prefix}.up_proj.weight", result) + self.assertNotIn(f"{prefix}.down_proj.weight", result) + + def test_router_gate_rename(self): + fake_weights = { + "model.layers.2.mlp.router.gate.weight": torch.zeros(1), + } + result = self.mapper.preprocess_weights(fake_weights) + self.assertIn("model.layers.2.mlp.gate.weight", result) + self.assertNotIn("model.layers.2.mlp.router.gate.weight", result) + + def test_expert_bias_rename(self): + fake_weights = { + "model.layers.2.mlp.expert_bias": torch.zeros(1), + } + result = self.mapper.preprocess_weights(fake_weights) + self.assertIn("model.layers.2.mlp.gate.e_score_correction_bias", result) + self.assertNotIn("model.layers.2.mlp.expert_bias", result) + + def test_attention_gate_proj_unchanged(self): + fake_weights = { + "model.layers.0.self_attn.gate_proj.weight": torch.zeros(1), + } + result = self.mapper.preprocess_weights(fake_weights) + self.assertIn("model.layers.0.self_attn.gate_proj.weight", result) + + def test_qkv_keys_unchanged_by_preprocess(self): + fake_weights = { + "model.layers.0.self_attn.q_proj.weight": torch.zeros(1), + "model.layers.0.self_attn.k_proj.weight": torch.zeros(1), + "model.layers.0.self_attn.v_proj.weight": torch.zeros(1), + } + result = self.mapper.preprocess_weights(fake_weights) + self.assertIn("model.layers.0.self_attn.q_proj.weight", result) + self.assertIn("model.layers.0.self_attn.k_proj.weight", result) + self.assertIn("model.layers.0.self_attn.v_proj.weight", result) + + def test_dense_mlp_keys_unchanged_by_preprocess(self): + fake_weights = { + "model.layers.0.mlp.gate_proj.weight": torch.zeros(1), + "model.layers.0.mlp.up_proj.weight": torch.zeros(1), + "model.layers.0.mlp.down_proj.weight": torch.zeros(1), + } + result = self.mapper.preprocess_weights(fake_weights) + self.assertIn("model.layers.0.mlp.gate_proj.weight", result) + self.assertIn("model.layers.0.mlp.up_proj.weight", result) + self.assertIn("model.layers.0.mlp.down_proj.weight", result) + + def test_is_special_instance_module_for_moe(self): + from unittest.mock import MagicMock + + from tensorrt_llm._torch.modules.fused_moe.interface import MoE + + mock_moe = MagicMock(spec=MoE) + mock_moe.__class__ = MoE + self.assertTrue(self.mapper.is_special_instance_module(mock_moe)) + + mock_linear = MagicMock(spec=torch.nn.Linear) + self.assertFalse(self.mapper.is_special_instance_module(mock_linear)) + + +class TestAfmoeWeightLoading(unittest.TestCase): + """Verify AfmoeForCausalLM applies mapper preprocessing in the real load hook.""" + + def test_load_weights_preprocesses_mapper_weights(self): + from tensorrt_llm._torch.models.modeling_utils import DecoderModelForCausalLM + + model = object.__new__(AfmoeForCausalLM) + raw_weights = {"model.layers.1.mlp.router.gate.weight": torch.zeros(1)} + processed_weights = {"model.layers.1.mlp.gate.weight": torch.zeros(1)} + mapper = Mock() + mapper.preprocess_weights.return_value = processed_weights + + with patch.object(DecoderModelForCausalLM, "load_weights", autospec=True) as load_weights: + AfmoeForCausalLM.load_weights(model, raw_weights, mapper, allow_partial_loading=True) + + mapper.preprocess_weights.assert_called_once_with(raw_weights) + load_weights.assert_called_once() + args, kwargs = load_weights.call_args + self.assertIs(args[0], model) + self.assertIs(kwargs["weights"], processed_weights) + self.assertIs(kwargs["weight_mapper"], mapper) + self.assertTrue(kwargs["allow_partial_loading"]) + + +class TestAfmoeSanity(unittest.TestCase): + """Smoke test: build a tiny AFMoE and run a forward pass.""" + + def test_afmoe_sanity(self): + config_dict = deepcopy(AFMOE_CONFIG) + afmoe_config = AfmoeConfig.from_dict(config_dict) + + model_config = ModelConfig(pretrained_config=afmoe_config, quant_config=QuantConfig()) + dtype = afmoe_config.torch_dtype + device = torch.device("cuda") + model = AfmoeForCausalLM(model_config).to(device) + + input_ids = torch.tensor( + [100, 200, 300, 100, 200, 100, 400, 500], dtype=torch.int, device=device + ) + + context_sequence_lengths = [3, 2, 1] + sequence_lengths = context_sequence_lengths + [1, 1] + past_seen_tokens = [0, 0, 0, 62, 75] + request_ids = list(range(len(sequence_lengths))) + token_nums = (torch.tensor(past_seen_tokens) + torch.tensor(sequence_lengths)).tolist() + prompt_lens = token_nums[:3] + past_seen_tokens[3:] + + num_blocks = 100 + tokens_per_block = 128 + head_dim = afmoe_config.hidden_size // afmoe_config.num_attention_heads + num_layers = afmoe_config.num_hidden_layers + num_kv_heads = afmoe_config.num_key_value_heads + max_seq_len = num_blocks * tokens_per_block + batch_size = len(context_sequence_lengths) + 2 + + if dtype == torch.half: + kv_cache_dtype = tensorrt_llm.bindings.DataType.HALF + elif dtype == torch.bfloat16: + kv_cache_dtype = tensorrt_llm.bindings.DataType.BF16 + else: + raise ValueError("Invalid dtype") + + mapping = Mapping(world_size=1, tp_size=1, rank=0) + kv_cache_config = KvCacheConfig(max_tokens=num_blocks * tokens_per_block) + kv_cache_manager = KVCacheManager( + kv_cache_config, + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=num_layers, + num_kv_heads=num_kv_heads, + head_dim=head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=batch_size, + mapping=mapping, + dtype=kv_cache_dtype, + ) + kv_cache_manager.add_dummy_requests(request_ids, token_nums) + + metadata_cls = get_attention_backend(model_config.attn_backend).Metadata + attn_metadata = metadata_cls( + seq_lens=torch.tensor(sequence_lengths, dtype=torch.int), + num_contexts=len(context_sequence_lengths), + kv_cache_params=KVCacheParams( + use_cache=True, + num_cached_tokens_per_seq=past_seen_tokens, + ), + kv_cache_manager=kv_cache_manager, + request_ids=request_ids, + prompt_lens=prompt_lens, + max_num_requests=len(context_sequence_lengths) + 2, + max_num_tokens=8192, + ) + + position_ids = [] + for i, tokens in enumerate(past_seen_tokens): + seq_len = context_sequence_lengths[i] if i < len(context_sequence_lengths) else 1 + position_id = torch.arange(tokens, tokens + seq_len, device=input_ids.device) + position_ids.append(position_id) + position_ids = torch.cat(position_ids).unsqueeze(0) + + with torch.inference_mode(): + attn_metadata.prepare() + logits = model.forward( + input_ids=input_ids, position_ids=position_ids, attn_metadata=attn_metadata + ) + + self.assertEqual(len(past_seen_tokens), logits.shape[0]) + kv_cache_manager.shutdown() + + def test_moe_layer_config(self): + config_dict = deepcopy(AFMOE_CONFIG) + afmoe_config = AfmoeConfig.from_dict(config_dict) + + device = torch.device("cuda") + model_config = ModelConfig(pretrained_config=afmoe_config) + model = AfmoeForCausalLM(model_config).to(device) + + self.assertEqual(len(model.model.layers), NUM_HIDDEN_LAYERS) + + for i in range(NUM_HIDDEN_LAYERS): + layer = model.model.layers[i] + if i < NUM_DENSE_LAYERS: + self.assertFalse(layer.moe_enabled, f"Layer {i} should be dense") + self.assertNotIsInstance(layer.mlp, AfmoeMoE) + else: + self.assertTrue(layer.moe_enabled, f"Layer {i} should be MoE") + self.assertIsInstance(layer.mlp, AfmoeMoE) + + +class TestAfmoeTPAttributes(unittest.TestCase): + """Verify TP-related module attributes are wired correctly.""" + + def _build_model(self, tp_size): + config_dict = deepcopy(AFMOE_CONFIG) + afmoe_config = AfmoeConfig.from_dict(config_dict) + mapping = Mapping(world_size=tp_size, tp_size=tp_size, rank=0) + model_config = ModelConfig(pretrained_config=afmoe_config, mapping=mapping) + return AfmoeForCausalLM(model_config) + + def test_gate_proj_is_column_parallel(self): + model = self._build_model(tp_size=1) + for layer in model.model.layers: + gate_proj = layer.self_attn.gate_proj + self.assertEqual(gate_proj.tp_mode, TensorParallelMode.COLUMN) + + def test_moe_experts_no_reduce(self): + model = self._build_model(tp_size=1) + for layer in model.model.layers: + if layer.moe_enabled: + self.assertFalse(layer.mlp.experts.reduce_results) + + def test_allreduce_created_for_tp2(self): + model = self._build_model(tp_size=2) + for layer in model.model.layers: + if layer.moe_enabled: + self.assertIsNotNone( + layer.mlp.allreduce, "MoE layer should have allreduce for tp_size=2" + ) + + def test_no_allreduce_for_tp1(self): + model = self._build_model(tp_size=1) + for layer in model.model.layers: + if layer.moe_enabled: + self.assertIsNone( + layer.mlp.allreduce, "MoE layer should NOT have allreduce for tp_size=1" + ) + + def test_gate_proj_output_matches_q_size(self): + model = self._build_model(tp_size=2) + config = model.config + head_dim = config.hidden_size // config.num_attention_heads + expected_local_out = (config.num_attention_heads // 2) * head_dim + + for layer in model.model.layers: + gate_proj = layer.self_attn.gate_proj + actual_out = gate_proj.weight.shape[0] + self.assertEqual( + actual_out, + expected_local_out, + f"gate_proj local output should be " + f"num_heads_per_tp * head_dim = {expected_local_out}, " + f"got {actual_out}", + ) + + def test_attention_layer_types(self): + model = self._build_model(tp_size=1) + layer_types = AFMOE_CONFIG["layer_types"] + for i, layer in enumerate(model.model.layers): + attn = layer.self_attn + if layer_types[i] == "sliding_attention": + self.assertTrue(attn.is_local_attention) + self.assertEqual(attn._attention_window_size, WINDOW_SIZE) + else: + self.assertFalse(attn.is_local_attention) + self.assertIsNone(attn._attention_window_size) + + +if __name__ == "__main__": + unittest.main() From 13a2e60aa99a1c12d1650c4d432f249e9b5c30e1 Mon Sep 17 00:00:00 2001 From: Alyosha-Swamy Date: Wed, 20 May 2026 05:52:26 +0000 Subject: [PATCH 2/8] Fix AFMoE attention-DP wiring and add coverage Add missing NVIDIA headers, modernize AFMoE typing, and make AFMoE attention-DP use unsharded gate/MLP projections. Add focused regression coverage plus a dummy-weight LLM API smoke test and register the AFMoE model test in B200 CI. Signed-off-by: Alyosha-Swamy --- tensorrt_llm/_torch/models/__init__.py | 6 +- .../_torch/models/checkpoints/__init__.py | 6 +- .../checkpoints/hf/afmoe_weight_mapper.py | 15 +++ tensorrt_llm/_torch/models/modeling_afmoe.py | 42 +++++++- .../test_lists/test-db/l0_b200.yml | 1 + .../_torch/modeling/test_modeling_afmoe.py | 101 ++++++++++++++++++ 6 files changed, 160 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index 3c070732ce7e..c4360a65471e 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -34,8 +34,7 @@ from .modeling_nemotron_nas import NemotronNASForCausalLM from .modeling_phi3 import Phi3ForCausalLM from .modeling_phi4mm import Phi4MMForCausalLM -from .modeling_qwen import (Qwen2ForCausalLM, Qwen2ForProcessRewardModel, - Qwen2ForRewardModel) +from .modeling_qwen import Qwen2ForCausalLM, Qwen2ForProcessRewardModel, Qwen2ForRewardModel from .modeling_qwen2vl import Qwen2_5_VLModel, Qwen2VLModel from .modeling_qwen3 import Qwen3ForCausalLM from .modeling_qwen3_5 import Qwen3_5ForCausalLM, Qwen3_5MoeForCausalLM @@ -108,7 +107,8 @@ __all__.append("MllamaForConditionalGeneration") else: print( - f"Failed to import MllamaForConditionalGeneration as transformers.__version__ {transformers.__version__} < 4.45.1" + "Failed to import MllamaForConditionalGeneration as " + f"transformers.__version__ {transformers.__version__} < 4.45.1" ) # Gemma4 requires transformers>=5.5.0 (native Gemma4 config/model classes). diff --git a/tensorrt_llm/_torch/models/checkpoints/__init__.py b/tensorrt_llm/_torch/models/checkpoints/__init__.py index b10a93132ca8..380c5a3bbc2d 100644 --- a/tensorrt_llm/_torch/models/checkpoints/__init__.py +++ b/tensorrt_llm/_torch/models/checkpoints/__init__.py @@ -17,11 +17,9 @@ from .hf.qwen3vl_weight_mapper import Qwen3VLHfWeightMapper from .hf.weight_loader import HfWeightLoader from .hf.weight_mapper import HfWeightMapper -from .mistral.checkpoint_loader import (MistralCheckpointLoader, - MistralLarge3CheckpointLoader) +from .mistral.checkpoint_loader import MistralCheckpointLoader, MistralLarge3CheckpointLoader from .mistral.config_loader import MistralConfigLoader -from .mistral.weight_mapper import (MistralLarge3WeightMapper, - MistralWeightMapper) +from .mistral.weight_mapper import MistralLarge3WeightMapper, MistralWeightMapper from .mx.checkpoint_loader import MXCheckpointLoader __all__ = [ diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py index 5d045239c7b2..139f6e6967b2 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py @@ -1,3 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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 torch import nn from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper diff --git a/tensorrt_llm/_torch/models/modeling_afmoe.py b/tensorrt_llm/_torch/models/modeling_afmoe.py index 439da67f9b77..c5fe788a8c86 100644 --- a/tensorrt_llm/_torch/models/modeling_afmoe.py +++ b/tensorrt_llm/_torch/models/modeling_afmoe.py @@ -1,4 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + """Inference-only AFMoE (Arcee Foundation MoE) for TensorRT-LLM. Follows the HF implementation of AfmoeForCausalLM. @@ -13,13 +27,14 @@ - Optional muP embedding scaling """ -from typing import Dict, List, Optional +from typing import Optional import torch from torch import nn from transformers import AutoConfig, PretrainedConfig from tensorrt_llm.functional import PositionEmbeddingType +from tensorrt_llm.mapping import Mapping from ...logger import logger from ..attention_backend import AttentionMetadata @@ -72,6 +87,23 @@ def _validate_routing_config(config: PretrainedConfig) -> None: ) +def _get_attention_dp_mapping(mapping: Mapping) -> Mapping: + """Return the effective TP=1 mapping used by attention-DP modules.""" + if not mapping.enable_attention_dp: + return mapping + + return Mapping( + world_size=mapping.world_size, + rank=mapping.rank, + gpus_per_node=mapping.gpus_per_node, + tp_size=1, + pp_size=mapping.pp_size * mapping.tp_size, + cp_size=mapping.cp_size, + cp_config=mapping.cp_config, + enable_attention_dp=mapping.enable_attention_dp, + ) + + class AfmoeGate(nn.Module): """Router gate for AFMoE, following the DeepSeekV3 grouped top-k pattern.""" @@ -111,7 +143,7 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: ) return logits - def load_weights(self, weights: List[Dict]): + def load_weights(self, weights: list[dict]): assert len(weights) == 1 self.weight.copy_(weights[0]["weight"][:]) self.e_score_correction_bias.copy_( @@ -184,6 +216,7 @@ def __init__( bias=False, dtype=config.torch_dtype, config=model_config, + overridden_tp_size=1 if self.enable_attention_dp else None, reduce_output=False, layer_idx=layer_idx, ) @@ -285,7 +318,7 @@ def __init__( config.num_attention_heads * self.head_dim, bias=False, dtype=config.torch_dtype, - mapping=model_config.mapping, + mapping=_get_attention_dp_mapping(model_config.mapping), tensor_parallel_mode=TensorParallelMode.COLUMN, gather_output=False, quant_config=model_config.get_quant_config(), @@ -371,6 +404,7 @@ def __init__( bias=False, dtype=config.torch_dtype, config=model_config, + overridden_tp_size=1 if model_config.mapping.enable_attention_dp else None, layer_idx=layer_idx, ) @@ -505,6 +539,6 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]): vocab_size=model_config.pretrained_config.vocab_size, ) - def load_weights(self, weights: Dict, weight_mapper, **kwargs): + def load_weights(self, weights: dict, weight_mapper, **kwargs): weights = weight_mapper.preprocess_weights(weights) super().load_weights(weights=weights, weight_mapper=weight_mapper, **kwargs) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 1d1ebe9e2ae2..3ff726683242 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -133,6 +133,7 @@ l0_b200: - unittest/_torch/modeling -k "modeling_llama" - unittest/_torch/modeling -k "modeling_mixtral" - unittest/_torch/modeling -k "modeling_gpt_oss" + - unittest/_torch/modeling/test_modeling_afmoe.py - unittest/_torch/modeling/test_modeling_exaone_moe.py - unittest/_torch/modeling/test_modeling_gemma4.py - unittest/_torch/modeling/test_gemma4_multimodal.py diff --git a/tests/unittest/_torch/modeling/test_modeling_afmoe.py b/tests/unittest/_torch/modeling/test_modeling_afmoe.py index c32babd49758..e0ec2fbc960e 100644 --- a/tests/unittest/_torch/modeling/test_modeling_afmoe.py +++ b/tests/unittest/_torch/modeling/test_modeling_afmoe.py @@ -1,3 +1,20 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + +import json +import tempfile import unittest from copy import deepcopy from unittest.mock import Mock, patch @@ -5,6 +22,7 @@ import torch import tensorrt_llm +from tensorrt_llm import LLM, SamplingParams from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig @@ -21,6 +39,8 @@ from tensorrt_llm._torch.modules.linear import TensorParallelMode from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm.bindings.executor import KvCacheConfig +from tensorrt_llm.llmapi import KvCacheConfig as LlmKvCacheConfig +from tensorrt_llm.llmapi import MoeConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig @@ -326,6 +346,54 @@ def test_moe_layer_config(self): self.assertIsInstance(layer.mlp, AfmoeMoE) +class TestAfmoeEndToEnd(unittest.TestCase): + """Exercise AFMoE through the PyTorch LLM API with dummy weights.""" + + def test_llm_dummy_load_generates_from_token_ids(self): + if not torch.cuda.is_available(): + self.skipTest("AFMoE LLM API test requires CUDA") + + with tempfile.TemporaryDirectory() as tmp_model_dir: + with open(f"{tmp_model_dir}/config.json", "w", encoding="utf-8") as f: + json.dump(AFMOE_CONFIG, f, indent=2) + + prompts = [ + {"prompt_token_ids": [100, 200, 300]}, + {"prompt_token_ids": [101, 202]}, + ] + sampling_params = SamplingParams( + max_tokens=2, + end_id=AFMOE_CONFIG["vocab_size"] - 1, + pad_id=AFMOE_CONFIG["vocab_size"] - 1, + detokenize=False, + ignore_eos=True, + ) + + with LLM( + model=tmp_model_dir, + load_format="dummy", + tensor_parallel_size=1, + enable_chunked_prefill=False, + disable_overlap_scheduler=True, + attn_backend="TRTLLM", + max_batch_size=len(prompts), + max_num_tokens=16, + max_seq_len=64, + moe_config=MoeConfig(max_num_tokens=64), + moe_expert_parallel_size=-1, + moe_tensor_parallel_size=-1, + enable_attention_dp=False, + kv_cache_config=LlmKvCacheConfig(enable_block_reuse=False), + ) as llm: + outputs = llm.generate(prompts, sampling_params=sampling_params) + + self.assertEqual(len(outputs), len(prompts)) + for prompt, output in zip(prompts, outputs): + self.assertEqual(output.prompt_token_ids, prompt["prompt_token_ids"]) + self.assertEqual(len(output.outputs), 1) + self.assertEqual(len(output.outputs[0].token_ids), sampling_params.max_tokens) + + class TestAfmoeTPAttributes(unittest.TestCase): """Verify TP-related module attributes are wired correctly.""" @@ -336,6 +404,18 @@ def _build_model(self, tp_size): model_config = ModelConfig(pretrained_config=afmoe_config, mapping=mapping) return AfmoeForCausalLM(model_config) + def _build_attention_dp_model(self, tp_size): + config_dict = deepcopy(AFMOE_CONFIG) + afmoe_config = AfmoeConfig.from_dict(config_dict) + mapping = Mapping( + world_size=tp_size, + tp_size=tp_size, + rank=0, + enable_attention_dp=True, + ) + model_config = ModelConfig(pretrained_config=afmoe_config, mapping=mapping) + return AfmoeForCausalLM(model_config) + def test_gate_proj_is_column_parallel(self): model = self._build_model(tp_size=1) for layer in model.model.layers: @@ -381,6 +461,27 @@ def test_gate_proj_output_matches_q_size(self): f"got {actual_out}", ) + def test_attention_dp_uses_unsharded_gate_and_mlp_modules(self): + model = self._build_attention_dp_model(tp_size=2) + config = model.config + expected_gate_out = config.num_attention_heads * ( + config.hidden_size // config.num_attention_heads + ) + + for layer in model.model.layers: + gate_proj = layer.self_attn.gate_proj + self.assertEqual(gate_proj.tp_size, 1) + self.assertEqual(gate_proj.weight.shape[0], expected_gate_out) + + if layer.moe_enabled: + self.assertIsNone(layer.mlp.allreduce) + if layer.mlp.shared_experts is not None: + self.assertEqual(layer.mlp.shared_experts.gate_up_proj.tp_size, 1) + self.assertEqual(layer.mlp.shared_experts.down_proj.tp_size, 1) + else: + self.assertEqual(layer.mlp.gate_up_proj.tp_size, 1) + self.assertEqual(layer.mlp.down_proj.tp_size, 1) + def test_attention_layer_types(self): model = self._build_model(tp_size=1) layer_types = AFMOE_CONFIG["layer_types"] From 038cc86ca4df367639fa3db671b12d662e6e5863 Mon Sep 17 00:00:00 2001 From: Alyosha-Swamy Date: Fri, 29 May 2026 05:45:15 +0000 Subject: [PATCH 3/8] [None][refactor] AFMoE: inherit QKNormRoPEAttention with fused output gate Address PR #13148 review feedback for AFMoE (Trinity) support. - AfmoeAttention now inherits QKNormRoPEAttention instead of duplicating QK-norm / RoPE / forward logic. The attention output gate is fused into the QKV projection via attn_output_gate=True, and RoPE is applied only on sliding (local) layers via skip_rope, matching the HF AfmoeAttention reference. Removes the bespoke _get_attention_dp_mapping helper since the base Attention handles attention-DP (tp_size -> 1) natively. - AfmoeHfWeightMapper fuses the separate HF q_proj and gate_proj matrices into the per-head interleaved layout [h0_q, h0_gate, h1_q, ...] expected by the fused QKV projection. - Add a guard-skipped HF logits parity test (TestAfmoeAllCloseToHF) that loads HF weights through the mapper and compares context-phase logits; skipped automatically when transformers < 5.8 (native AFMoE unavailable). - Update TP-attribute and weight-mapper unit tests for the fused gate. - Minimize the _torch/models and checkpoints __init__.py diffs to the AFMoE entries only, and add AfmoeForCausalLM to supported-models.md. Validated: real-weight generation is coherent (Trinity-Mini), 23 AFMoE unit tests pass, and cross-checked HF logit parity passes at atol=1.0, rtol=0.5. Signed-off-by: Alyosha-Swamy --- docs/source/models/supported-models.md | 1 + tensorrt_llm/_torch/models/__init__.py | 6 +- .../_torch/models/checkpoints/__init__.py | 6 +- .../checkpoints/hf/afmoe_weight_mapper.py | 48 ++++ tensorrt_llm/_torch/models/modeling_afmoe.py | 127 ++-------- .../_torch/modeling/test_modeling_afmoe.py | 235 +++++++++++++++--- 6 files changed, 286 insertions(+), 137 deletions(-) diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 3381c41fc46b..b653e731d9e4 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -5,6 +5,7 @@ The following is a table of supported models for the PyTorch backend: | Architecture | Model | HuggingFace Example | | ------------------------------------ | ---------------------------------- | -------------------------------------------- | +| `AfmoeForCausalLM` | Arcee Foundation MoE (Trinity) | `arcee-ai/Trinity-Mini` | | `BertForSequenceClassification` | BERT-based | `textattack/bert-base-uncased-yelp-polarity` | | `Cohere2ForCausalLM` | Command A | `CohereLabs/c4ai-command-a-03-2025` | | `DeciLMForCausalLM` | Nemotron | `nvidia/Llama-3_1-Nemotron-51B-Instruct` | diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index c4360a65471e..3c070732ce7e 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -34,7 +34,8 @@ from .modeling_nemotron_nas import NemotronNASForCausalLM from .modeling_phi3 import Phi3ForCausalLM from .modeling_phi4mm import Phi4MMForCausalLM -from .modeling_qwen import Qwen2ForCausalLM, Qwen2ForProcessRewardModel, Qwen2ForRewardModel +from .modeling_qwen import (Qwen2ForCausalLM, Qwen2ForProcessRewardModel, + Qwen2ForRewardModel) from .modeling_qwen2vl import Qwen2_5_VLModel, Qwen2VLModel from .modeling_qwen3 import Qwen3ForCausalLM from .modeling_qwen3_5 import Qwen3_5ForCausalLM, Qwen3_5MoeForCausalLM @@ -107,8 +108,7 @@ __all__.append("MllamaForConditionalGeneration") else: print( - "Failed to import MllamaForConditionalGeneration as " - f"transformers.__version__ {transformers.__version__} < 4.45.1" + f"Failed to import MllamaForConditionalGeneration as transformers.__version__ {transformers.__version__} < 4.45.1" ) # Gemma4 requires transformers>=5.5.0 (native Gemma4 config/model classes). diff --git a/tensorrt_llm/_torch/models/checkpoints/__init__.py b/tensorrt_llm/_torch/models/checkpoints/__init__.py index 380c5a3bbc2d..b10a93132ca8 100644 --- a/tensorrt_llm/_torch/models/checkpoints/__init__.py +++ b/tensorrt_llm/_torch/models/checkpoints/__init__.py @@ -17,9 +17,11 @@ from .hf.qwen3vl_weight_mapper import Qwen3VLHfWeightMapper from .hf.weight_loader import HfWeightLoader from .hf.weight_mapper import HfWeightMapper -from .mistral.checkpoint_loader import MistralCheckpointLoader, MistralLarge3CheckpointLoader +from .mistral.checkpoint_loader import (MistralCheckpointLoader, + MistralLarge3CheckpointLoader) from .mistral.config_loader import MistralConfigLoader -from .mistral.weight_mapper import MistralLarge3WeightMapper, MistralWeightMapper +from .mistral.weight_mapper import (MistralLarge3WeightMapper, + MistralWeightMapper) from .mx.checkpoint_loader import MXCheckpointLoader __all__ = [ diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py index 139f6e6967b2..40bb97939175 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/afmoe_weight_mapper.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import torch from torch import nn from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper @@ -38,8 +39,55 @@ def __init__(self): def preprocess_weights(self, weights: dict) -> dict: weights = self.rename_by_params_map(self.params_map, weights) + weights = self._fuse_attention_gate(weights) return weights + def _fuse_attention_gate(self, weights: dict) -> dict: + """Fuse the separate attention ``gate_proj`` into ``q_proj``. + + AfmoeAttention uses ``attn_output_gate=True``, so the gate weights are + interleaved with the query weights per head and loaded through the fused + QKV projection. The HF checkpoint stores ``q_proj`` and ``gate_proj`` as + two separate matrices of shape ``[num_heads * head_dim, hidden]``; the + fused QKV projection expects the query slot laid out per head as + ``[head0_q, head0_gate, head1_q, head1_gate, ...]`` (see + ``Attention.forward`` where ``q_gate`` is viewed as + ``[..., num_heads, 2 * head_dim]`` and chunked into q/gate). + """ + marker = ".self_attn.gate_proj." + gate_keys = [k for k in weights if marker in k] + if not gate_keys: + return weights + + num_heads = self.model.config.num_attention_heads + for gate_key in gate_keys: + prefix, suffix = gate_key.split(marker) + q_key = f"{prefix}.self_attn.q_proj.{suffix}" + if q_key not in weights: + continue + weights[q_key] = self._interleave_per_head(weights[q_key], weights[gate_key], num_heads) + del weights[gate_key] + return weights + + @staticmethod + def _interleave_per_head(q: torch.Tensor, gate: torch.Tensor, num_heads: int) -> torch.Tensor: + """Interleave q and gate rows per head: ``[h0_q, h0_gate, h1_q, ...]``. + + Works for 2D weights ``[num_heads * per_head, hidden]`` as well as 1D + biases and FP8 block scales, since the split is always taken along the + leading (output) dimension. + """ + assert q.shape[0] % num_heads == 0, ( + f"q_proj rows {q.shape[0]} not divisible by num_heads {num_heads}" + ) + assert gate.shape == q.shape, f"gate_proj shape {gate.shape} != q_proj shape {q.shape}" + per_head = q.shape[0] // num_heads + tail = q.shape[1:] + q = q.reshape(num_heads, per_head, *tail) + gate = gate.reshape(num_heads, per_head, *tail) + fused = torch.stack([q, gate], dim=1) + return fused.reshape(num_heads * 2 * per_head, *tail).contiguous() + def is_special_instance_module(self, module: nn.Module) -> bool: return isinstance(module, MoE) diff --git a/tensorrt_llm/_torch/models/modeling_afmoe.py b/tensorrt_llm/_torch/models/modeling_afmoe.py index c5fe788a8c86..9a866997dc1c 100644 --- a/tensorrt_llm/_torch/models/modeling_afmoe.py +++ b/tensorrt_llm/_torch/models/modeling_afmoe.py @@ -34,24 +34,18 @@ from transformers import AutoConfig, PretrainedConfig from tensorrt_llm.functional import PositionEmbeddingType -from tensorrt_llm.mapping import Mapping from ...logger import logger from ..attention_backend import AttentionMetadata -from ..attention_backend.interface import ( - PositionalEmbeddingParams, - PredefinedAttentionMask, - RopeParams, -) +from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams from ..distributed import AllReduce from ..model_config import ModelConfig -from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding from ..modules.fused_moe import DeepSeekV3MoeRoutingMethod, create_moe from ..modules.fused_moe.routing import Deepseekv3RoutingImpl 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 ..utils import AuxStreamType from .modeling_utils import DecoderModel, DecoderModelForCausalLM, register_auto_model @@ -87,23 +81,6 @@ def _validate_routing_config(config: PretrainedConfig) -> None: ) -def _get_attention_dp_mapping(mapping: Mapping) -> Mapping: - """Return the effective TP=1 mapping used by attention-DP modules.""" - if not mapping.enable_attention_dp: - return mapping - - return Mapping( - world_size=mapping.world_size, - rank=mapping.rank, - gpus_per_node=mapping.gpus_per_node, - tp_size=1, - pp_size=mapping.pp_size * mapping.tp_size, - cp_size=mapping.cp_size, - cp_config=mapping.cp_config, - enable_attention_dp=mapping.enable_attention_dp, - ) - - class AfmoeGate(nn.Module): """Router gate for AFMoE, following the DeepSeekV3 grouped top-k pattern.""" @@ -259,11 +236,13 @@ def forward( return final_output -class AfmoeAttention(Attention): - """Attention with Q/K norm, per-layer sliding window, gated output. +class AfmoeAttention(QKNormRoPEAttention): + """Attention with Q/K norm, per-layer sliding window, and a sigmoid output gate. - Uses a separate gate_proj linear (not fused into QKV) to gate the - attention output with sigmoid, matching the HF checkpoint layout. + Inherits QK-norm + RoPE handling from ``QKNormRoPEAttention``. The output + gate is fused into the QKV projection (``attn_output_gate=True``), and RoPE + is applied only on local (sliding-window) layers, matching the HF + ``AfmoeAttention`` reference. """ def __init__( @@ -278,17 +257,14 @@ def __init__( and layer_idx < len(layer_types) and layer_types[layer_idx] == "sliding_attention" ) - self._attention_window_size = config.sliding_window if self.is_local_attention else None + self.attention_window_size = config.sliding_window if self.is_local_attention else None - rope_params = RopeParams.from_config(config) if self.is_local_attention else None - pos_embd_params = ( - PositionalEmbeddingParams( + pos_embd_params = None + if self.is_local_attention: + pos_embd_params = PositionalEmbeddingParams( type=PositionEmbeddingType.rope_gpt_neox, - rope=rope_params, + rope=RopeParams.from_config(config), ) - if self.is_local_attention - else None - ) super().__init__( hidden_size=config.hidden_size, @@ -297,56 +273,15 @@ def __init__( max_position_embeddings=getattr(config, "max_position_embeddings", 131072), bias=False, pos_embd_params=pos_embd_params, + fuse_qk_norm_rope=False, + skip_rope=not self.is_local_attention, + attn_output_gate=True, + is_qk_norm=True, layer_idx=layer_idx, dtype=config.torch_dtype, config=model_config, ) - self.q_norm = RMSNorm( - hidden_size=self.head_dim, - eps=config.rms_norm_eps, - dtype=config.torch_dtype, - ) - self.k_norm = RMSNorm( - hidden_size=self.head_dim, - eps=config.rms_norm_eps, - dtype=config.torch_dtype, - ) - - self.gate_proj = Linear( - config.hidden_size, - config.num_attention_heads * self.head_dim, - bias=False, - dtype=config.torch_dtype, - mapping=_get_attention_dp_mapping(model_config.mapping), - tensor_parallel_mode=TensorParallelMode.COLUMN, - gather_output=False, - quant_config=model_config.get_quant_config(), - skip_create_weights_in_init=model_config.skip_create_weights_in_init, - allreduce_strategy=model_config.allreduce_strategy, - force_dynamic_quantization=model_config.force_dynamic_quantization, - use_cute_dsl_blockscaling_mm=model_config.use_cute_dsl_blockscaling_mm, - ) - - def apply_rope( - self, - q: torch.Tensor, - k: Optional[torch.Tensor], - v: Optional[torch.Tensor], - position_ids: torch.Tensor, - ): - q, k, v = self.split_qkv(q, k, v) - - q_shape = q.shape - k_shape = k.shape - q = self.q_norm(q.reshape(-1, self.num_heads, self.head_dim)).reshape(q_shape) - k = self.k_norm(k.reshape(-1, self.num_key_value_heads, self.head_dim)).reshape(k_shape) - - if self.is_local_attention and not self.rope_fusion and position_ids is not None: - q, k = self.rotary_emb(position_ids, [q, k]) - - return q, k, v - def forward( self, position_ids: torch.IntTensor, @@ -354,30 +289,14 @@ def forward( attn_metadata: AttentionMetadata, **kwargs, ) -> torch.Tensor: - gate = self.gate_proj(hidden_states) - - qkv = self.qkv_proj(hidden_states) - q, k, v = qkv, None, None - - q, k, v = self.apply_rope(q, k, v, position_ids) - q, k, v = self.convert_qkv(q, k, v) - - attn_output = self.forward_impl( - q, - k, - v, - attn_metadata, - attention_mask=PredefinedAttentionMask.CAUSAL, - attention_window_size=self._attention_window_size, - attention_mask_data=None, - mrope_config=None, + return super().forward( + position_ids=position_ids, + hidden_states=hidden_states, + attn_metadata=attn_metadata, + attention_window_size=self.attention_window_size, + **kwargs, ) - attn_output = attn_output * torch.sigmoid(gate) - - attn_output = self.o_proj(attn_output) - return attn_output - class AfmoeDecoderLayer(DecoderLayer): def __init__( diff --git a/tests/unittest/_torch/modeling/test_modeling_afmoe.py b/tests/unittest/_torch/modeling/test_modeling_afmoe.py index e0ec2fbc960e..6b94991a2f65 100644 --- a/tests/unittest/_torch/modeling/test_modeling_afmoe.py +++ b/tests/unittest/_torch/modeling/test_modeling_afmoe.py @@ -44,6 +44,16 @@ from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig +try: + # AFMoE landed in transformers >= 5.8. When the installed transformers is + # older the parity test below is skipped (the rest of the suite still runs). + from transformers import AfmoeConfig as HFAfmoeConfig + from transformers.models.afmoe.modeling_afmoe import AfmoeForCausalLM as HFAfmoeForCausalLM + + HAS_HF_AFMOE = True +except ImportError: + HAS_HF_AFMOE = False + WINDOW_SIZE = 4 NUM_HIDDEN_LAYERS = 4 NUM_DENSE_LAYERS = 1 @@ -175,14 +185,33 @@ def test_expert_bias_rename(self): self.assertIn("model.layers.2.mlp.gate.e_score_correction_bias", result) self.assertNotIn("model.layers.2.mlp.expert_bias", result) - def test_attention_gate_proj_unchanged(self): + def test_attention_gate_fused_into_q(self): + # AfmoeAttention uses attn_output_gate=True, so the separate gate_proj + # is interleaved per head into q_proj and the gate_proj key is dropped. + num_heads, head_dim, hidden = 8, 32, 256 + self.mapper._model = Mock() + self.mapper._model.config = Mock(num_attention_heads=num_heads) + + q = torch.arange(num_heads * head_dim * hidden, dtype=torch.float32).reshape( + num_heads * head_dim, hidden + ) + gate = q + 0.5 fake_weights = { - "model.layers.0.self_attn.gate_proj.weight": torch.zeros(1), + "model.layers.0.self_attn.q_proj.weight": q, + "model.layers.0.self_attn.gate_proj.weight": gate, } result = self.mapper.preprocess_weights(fake_weights) - self.assertIn("model.layers.0.self_attn.gate_proj.weight", result) - def test_qkv_keys_unchanged_by_preprocess(self): + self.assertNotIn("model.layers.0.self_attn.gate_proj.weight", result) + fused = result["model.layers.0.self_attn.q_proj.weight"] + self.assertEqual(fused.shape, (2 * num_heads * head_dim, hidden)) + # Per head the layout is [q_head, gate_head]: first head_dim rows are q, + # next head_dim rows are gate. + torch.testing.assert_close(fused[:head_dim], q[:head_dim]) + torch.testing.assert_close(fused[head_dim : 2 * head_dim], gate[:head_dim]) + + def test_qkv_keys_unchanged_without_gate(self): + # Without a gate_proj key (e.g. partial weight dicts), q/k/v are untouched. fake_weights = { "model.layers.0.self_attn.q_proj.weight": torch.zeros(1), "model.layers.0.self_attn.k_proj.weight": torch.zeros(1), @@ -416,11 +445,12 @@ def _build_attention_dp_model(self, tp_size): model_config = ModelConfig(pretrained_config=afmoe_config, mapping=mapping) return AfmoeForCausalLM(model_config) - def test_gate_proj_is_column_parallel(self): + def test_qkv_is_column_parallel_with_output_gate(self): model = self._build_model(tp_size=1) for layer in model.model.layers: - gate_proj = layer.self_attn.gate_proj - self.assertEqual(gate_proj.tp_mode, TensorParallelMode.COLUMN) + attn = layer.self_attn + self.assertTrue(attn.attn_output_gate) + self.assertEqual(attn.qkv_proj.tp_mode, TensorParallelMode.COLUMN) def test_moe_experts_no_reduce(self): model = self._build_model(tp_size=1) @@ -444,34 +474,30 @@ def test_no_allreduce_for_tp1(self): layer.mlp.allreduce, "MoE layer should NOT have allreduce for tp_size=1" ) - def test_gate_proj_output_matches_q_size(self): + def test_qkv_output_includes_fused_gate(self): + # With attn_output_gate=True the query slot is doubled (q + gate) and + # fused into qkv_proj, so its local output is 2*q_size + 2*kv_size. model = self._build_model(tp_size=2) - config = model.config - head_dim = config.hidden_size // config.num_attention_heads - expected_local_out = (config.num_attention_heads // 2) * head_dim - for layer in model.model.layers: - gate_proj = layer.self_attn.gate_proj - actual_out = gate_proj.weight.shape[0] + attn = layer.self_attn + expected_out = attn.q_size * 2 + 2 * attn.kv_size + actual_out = attn.qkv_proj.weight.shape[0] self.assertEqual( actual_out, - expected_local_out, - f"gate_proj local output should be " - f"num_heads_per_tp * head_dim = {expected_local_out}, " - f"got {actual_out}", + expected_out, + f"qkv_proj local output should be 2*q_size + 2*kv_size = " + f"{expected_out}, got {actual_out}", ) - def test_attention_dp_uses_unsharded_gate_and_mlp_modules(self): + def test_attention_dp_uses_unsharded_qkv_and_mlp_modules(self): model = self._build_attention_dp_model(tp_size=2) - config = model.config - expected_gate_out = config.num_attention_heads * ( - config.hidden_size // config.num_attention_heads - ) for layer in model.model.layers: - gate_proj = layer.self_attn.gate_proj - self.assertEqual(gate_proj.tp_size, 1) - self.assertEqual(gate_proj.weight.shape[0], expected_gate_out) + attn = layer.self_attn + self.assertEqual(attn.qkv_proj.tp_size, 1) + # Unsharded: full heads, q slot doubled by the output gate. + expected_out = attn.q_size * 2 + 2 * attn.kv_size + self.assertEqual(attn.qkv_proj.weight.shape[0], expected_out) if layer.moe_enabled: self.assertIsNone(layer.mlp.allreduce) @@ -489,10 +515,163 @@ def test_attention_layer_types(self): attn = layer.self_attn if layer_types[i] == "sliding_attention": self.assertTrue(attn.is_local_attention) - self.assertEqual(attn._attention_window_size, WINDOW_SIZE) + self.assertEqual(attn.attention_window_size, WINDOW_SIZE) else: self.assertFalse(attn.is_local_attention) - self.assertIsNone(attn._attention_window_size) + self.assertIsNone(attn.attention_window_size) + + +@unittest.skipUnless( + HAS_HF_AFMOE, + "transformers>=5.8 with native AFMoE (transformers.models.afmoe) is required", +) +@unittest.skipUnless(torch.cuda.is_available(), "needs CUDA") +class TestAfmoeAllCloseToHF(unittest.TestCase): + """Compare TRT-LLM AFMoE context-phase logits against the HF reference. + + Loads the HF model's weights into the TRT-LLM model via AfmoeHfWeightMapper, + exercising the per-head q/gate fusion (attn_output_gate) and the HF + fused-expert -> per-expert conversion, then checks logit parity. + """ + + # Field names follow the HF AfmoeConfig schema (released-checkpoint names). + HF_CONFIG = { + "hidden_size": 256, + "intermediate_size": 512, + "moe_intermediate_size": 128, + "head_dim": 32, + "num_attention_heads": 8, + "num_key_value_heads": 2, + "num_hidden_layers": NUM_HIDDEN_LAYERS, + "num_dense_layers": NUM_DENSE_LAYERS, + "num_experts": 8, + "num_experts_per_tok": 2, + "num_shared_experts": 1, + "global_attn_every_n_layers": 4, + "sliding_window": WINDOW_SIZE, + "max_position_embeddings": 2048, + "rms_norm_eps": 1e-5, + "rope_theta": 10000, + "route_scale": 1.0, + "route_norm": True, + "score_func": "sigmoid", + "vocab_size": 1024, + "hidden_act": "silu", + "tie_word_embeddings": False, + } + + @staticmethod + def _convert_hf_experts(state_dict, moe_intermediate_size): + """Split HF fused 3D expert params into per-expert gate/up/down weights. + + HF-native AFMoE stores experts as ``experts.gate_up_proj`` + ``[num_experts, 2 * moe_inter, hidden]`` and ``experts.down_proj`` + ``[num_experts, hidden, moe_inter]``; AfmoeHfWeightMapper expects the + released-checkpoint layout with separate per-expert matrices. + """ + converted = dict(state_dict) + gate_up_keys = [k for k in state_dict if k.endswith(".experts.gate_up_proj")] + for gate_up_key in gate_up_keys: + prefix = gate_up_key[: -len(".gate_up_proj")] + gate_up = converted.pop(gate_up_key) + down = converted.pop(prefix + ".down_proj") + for expert_idx in range(gate_up.shape[0]): + converted[f"{prefix}.{expert_idx}.gate_proj.weight"] = gate_up[expert_idx][ + :moe_intermediate_size + ].contiguous() + converted[f"{prefix}.{expert_idx}.up_proj.weight"] = gate_up[expert_idx][ + moe_intermediate_size: + ].contiguous() + converted[f"{prefix}.{expert_idx}.down_proj.weight"] = down[expert_idx].contiguous() + return converted + + @torch.no_grad() + def test_afmoe_allclose_to_hf(self): + from tensorrt_llm._torch.models.checkpoints.hf.afmoe_weight_mapper import ( + AfmoeHfWeightMapper, + ) + + torch.manual_seed(0) + device = torch.device("cuda") + dtype = torch.bfloat16 + + hf_config = HFAfmoeConfig(dtype="float32", **self.HF_CONFIG) + hf_model = HFAfmoeForCausalLM(hf_config).to(dtype).to(device).eval() + + # TRT-LLM needs a couple of routing fields that the HF schema names + # differently; provide both so AfmoeConfig validates and builds. + trt_config_dict = dict(self.HF_CONFIG) + trt_config_dict.update( + architectures=["AfmoeForCausalLM"], + model_type="afmoe", + dtype="bfloat16", + n_group=1, + topk_group=1, + scoring_func=self.HF_CONFIG["score_func"], + norm_topk_prob=self.HF_CONFIG["route_norm"], + ) + afmoe_config = AfmoeConfig.from_dict(trt_config_dict) + model_config = ModelConfig(pretrained_config=afmoe_config) + model = AfmoeForCausalLM(model_config).to(dtype).to(device) + + weights = self._convert_hf_experts( + hf_model.state_dict(), self.HF_CONFIG["moe_intermediate_size"] + ) + weights = {k: v.to(dtype) for k, v in weights.items()} + + weight_mapper = AfmoeHfWeightMapper() + weight_mapper.init_model_and_config(model, model_config) + model.load_weights(weights, weight_mapper) + if hasattr(model, "post_load_weights"): + model.post_load_weights() + + # Short context: input_len < sliding_window so sliding == full attention. + input_len = WINDOW_SIZE - 1 + input_ids = torch.tensor([101, 202, 303][:input_len], dtype=torch.int32, device=device) + position_ids = torch.arange(input_len, dtype=torch.int32, device=device).unsqueeze(0) + + num_blocks, tokens_per_block = 4, 128 + kv_cache_manager = KVCacheManager( + KvCacheConfig(max_tokens=num_blocks * tokens_per_block), + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, + num_layers=hf_config.num_hidden_layers, + num_kv_heads=hf_config.num_key_value_heads, + head_dim=hf_config.head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=num_blocks * tokens_per_block, + max_batch_size=1, + mapping=Mapping(world_size=1, tp_size=1, rank=0), + dtype=tensorrt_llm.bindings.DataType.BF16, + ) + kv_cache_manager.add_dummy_requests([0], [input_len]) + metadata_cls = get_attention_backend(model_config.attn_backend).Metadata + attn_metadata = metadata_cls( + seq_lens=torch.tensor([input_len], dtype=torch.int), + num_contexts=1, + kv_cache_params=KVCacheParams(use_cache=True, num_cached_tokens_per_seq=[0]), + kv_cache_manager=kv_cache_manager, + request_ids=[0], + prompt_lens=[input_len], + max_num_requests=1, + max_num_tokens=8192, + ) + + hf_position_ids = position_ids.to(torch.long) + with torch.inference_mode(): + attn_metadata.prepare() + logits = model.forward( + input_ids=input_ids, position_ids=position_ids, attn_metadata=attn_metadata + ) + ref = hf_model.forward( + input_ids=input_ids.unsqueeze(0).long(), + position_ids=hf_position_ids, + use_cache=False, + ) + + # Loose tolerance: bf16 + token-choice MoE routing amplify per-logit + # noise (same rationale as the EXAONE-MoE parity test). + torch.testing.assert_close(logits, ref.logits[:, -1].float(), atol=1.0, rtol=0.5) + kv_cache_manager.shutdown() if __name__ == "__main__": From c425f618eb5161c07a1d10d3fac07bd459dd2776 Mon Sep 17 00:00:00 2001 From: Alyosha-Swamy Date: Sun, 31 May 2026 07:01:57 +0000 Subject: [PATCH 4/8] [None][fix] Guard AFMoE config registration Signed-off-by: Alyosha-Swamy --- tensorrt_llm/_torch/models/modeling_afmoe.py | 16 +++++++++------- .../_torch/modeling/test_modeling_afmoe.py | 16 ++-------------- 2 files changed, 11 insertions(+), 21 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_afmoe.py b/tensorrt_llm/_torch/models/modeling_afmoe.py index 9a866997dc1c..e11ad402409e 100644 --- a/tensorrt_llm/_torch/models/modeling_afmoe.py +++ b/tensorrt_llm/_torch/models/modeling_afmoe.py @@ -31,7 +31,8 @@ import torch from torch import nn -from transformers import AutoConfig, PretrainedConfig +from transformers import PretrainedConfig +from transformers.models.auto.configuration_auto import CONFIG_MAPPING from tensorrt_llm.functional import PositionEmbeddingType @@ -55,12 +56,13 @@ class AfmoeConfig(PretrainedConfig): model_type = "afmoe" -logger.warning_once( - "transformers does not natively support 'AfmoeConfig'. " - "Registering AfmoeConfig so AutoConfig can load AFMoE checkpoints.", - key="AFMOE_REGISTER_WARNING", -) -AutoConfig.register(AfmoeConfig.model_type, AfmoeConfig) +if AfmoeConfig.model_type not in CONFIG_MAPPING: + logger.warning_once( + "transformers does not natively support 'AfmoeConfig'. " + "Registering AfmoeConfig so AutoConfig can load AFMoE checkpoints.", + key="AFMOE_REGISTER_WARNING", + ) + CONFIG_MAPPING.register(AfmoeConfig.model_type, AfmoeConfig, exist_ok=True) def _validate_routing_config(config: PretrainedConfig) -> None: diff --git a/tests/unittest/_torch/modeling/test_modeling_afmoe.py b/tests/unittest/_torch/modeling/test_modeling_afmoe.py index 6b94991a2f65..58c17d0c8127 100644 --- a/tests/unittest/_torch/modeling/test_modeling_afmoe.py +++ b/tests/unittest/_torch/modeling/test_modeling_afmoe.py @@ -43,16 +43,8 @@ from tensorrt_llm.llmapi import MoeConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig - -try: - # AFMoE landed in transformers >= 5.8. When the installed transformers is - # older the parity test below is skipped (the rest of the suite still runs). - from transformers import AfmoeConfig as HFAfmoeConfig - from transformers.models.afmoe.modeling_afmoe import AfmoeForCausalLM as HFAfmoeForCausalLM - - HAS_HF_AFMOE = True -except ImportError: - HAS_HF_AFMOE = False +from transformers import AfmoeConfig as HFAfmoeConfig +from transformers.models.afmoe.modeling_afmoe import AfmoeForCausalLM as HFAfmoeForCausalLM WINDOW_SIZE = 4 NUM_HIDDEN_LAYERS = 4 @@ -521,10 +513,6 @@ def test_attention_layer_types(self): self.assertIsNone(attn.attention_window_size) -@unittest.skipUnless( - HAS_HF_AFMOE, - "transformers>=5.8 with native AFMoE (transformers.models.afmoe) is required", -) @unittest.skipUnless(torch.cuda.is_available(), "needs CUDA") class TestAfmoeAllCloseToHF(unittest.TestCase): """Compare TRT-LLM AFMoE context-phase logits against the HF reference. From 1f90a8de1f84331847d59dbdb40ac85ba264d19f Mon Sep 17 00:00:00 2001 From: Alyosha-Swamy Date: Thu, 4 Jun 2026 09:48:54 +0000 Subject: [PATCH 5/8] [None][fix] Stabilize AFMoE unit tests Signed-off-by: Alyosha-Swamy --- .../_torch/modeling/test_modeling_afmoe.py | 82 ++++++++++++------- 1 file changed, 52 insertions(+), 30 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_afmoe.py b/tests/unittest/_torch/modeling/test_modeling_afmoe.py index 58c17d0c8127..bbe1461bf9f2 100644 --- a/tests/unittest/_torch/modeling/test_modeling_afmoe.py +++ b/tests/unittest/_torch/modeling/test_modeling_afmoe.py @@ -86,6 +86,25 @@ } +def _synthetic_tp_mapping(tp_size: int, enable_attention_dp: bool = False) -> Mapping: + # These tests inspect module TP attributes in one pytest process. Force + # the lightweight MPI-topology Mapping even when TLLM_DISABLE_MPI=1 would + # otherwise require an initialized torch.distributed DeviceMesh. + with patch("tensorrt_llm.mapping.mpi_disabled", return_value=False): + return Mapping( + world_size=tp_size, + tp_size=tp_size, + rank=0, + enable_attention_dp=enable_attention_dp, + ) + + +def _shutdown_kv_cache_manager(kv_cache_manager: KVCacheManager) -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + kv_cache_manager.shutdown() + + class TestAfmoeRegistry(unittest.TestCase): """Verify AfmoeForCausalLM resolves through _torch auto-model registration.""" @@ -338,14 +357,16 @@ def test_afmoe_sanity(self): position_ids.append(position_id) position_ids = torch.cat(position_ids).unsqueeze(0) - with torch.inference_mode(): - attn_metadata.prepare() - logits = model.forward( - input_ids=input_ids, position_ids=position_ids, attn_metadata=attn_metadata - ) + try: + with torch.inference_mode(): + attn_metadata.prepare() + logits = model.forward( + input_ids=input_ids, position_ids=position_ids, attn_metadata=attn_metadata + ) - self.assertEqual(len(past_seen_tokens), logits.shape[0]) - kv_cache_manager.shutdown() + self.assertEqual(len(past_seen_tokens), logits.shape[0]) + finally: + _shutdown_kv_cache_manager(kv_cache_manager) def test_moe_layer_config(self): config_dict = deepcopy(AFMOE_CONFIG) @@ -421,20 +442,19 @@ class TestAfmoeTPAttributes(unittest.TestCase): def _build_model(self, tp_size): config_dict = deepcopy(AFMOE_CONFIG) afmoe_config = AfmoeConfig.from_dict(config_dict) - mapping = Mapping(world_size=tp_size, tp_size=tp_size, rank=0) - model_config = ModelConfig(pretrained_config=afmoe_config, mapping=mapping) + mapping = _synthetic_tp_mapping(tp_size) + model_config = ModelConfig( + pretrained_config=afmoe_config, mapping=mapping, allreduce_strategy="NCCL" + ) return AfmoeForCausalLM(model_config) def _build_attention_dp_model(self, tp_size): config_dict = deepcopy(AFMOE_CONFIG) afmoe_config = AfmoeConfig.from_dict(config_dict) - mapping = Mapping( - world_size=tp_size, - tp_size=tp_size, - rank=0, - enable_attention_dp=True, + mapping = _synthetic_tp_mapping(tp_size, enable_attention_dp=True) + model_config = ModelConfig( + pretrained_config=afmoe_config, mapping=mapping, allreduce_strategy="NCCL" ) - model_config = ModelConfig(pretrained_config=afmoe_config, mapping=mapping) return AfmoeForCausalLM(model_config) def test_qkv_is_column_parallel_with_output_gate(self): @@ -644,22 +664,24 @@ def test_afmoe_allclose_to_hf(self): max_num_tokens=8192, ) - hf_position_ids = position_ids.to(torch.long) - with torch.inference_mode(): - attn_metadata.prepare() - logits = model.forward( - input_ids=input_ids, position_ids=position_ids, attn_metadata=attn_metadata - ) - ref = hf_model.forward( - input_ids=input_ids.unsqueeze(0).long(), - position_ids=hf_position_ids, - use_cache=False, - ) + try: + hf_position_ids = position_ids.to(torch.long) + with torch.inference_mode(): + attn_metadata.prepare() + logits = model.forward( + input_ids=input_ids, position_ids=position_ids, attn_metadata=attn_metadata + ) + ref = hf_model.forward( + input_ids=input_ids.unsqueeze(0).long(), + position_ids=hf_position_ids, + use_cache=False, + ) - # Loose tolerance: bf16 + token-choice MoE routing amplify per-logit - # noise (same rationale as the EXAONE-MoE parity test). - torch.testing.assert_close(logits, ref.logits[:, -1].float(), atol=1.0, rtol=0.5) - kv_cache_manager.shutdown() + # Loose tolerance: bf16 + token-choice MoE routing amplify per-logit + # noise (same rationale as the EXAONE-MoE parity test). + torch.testing.assert_close(logits, ref.logits[:, -1].float(), atol=1.0, rtol=0.5) + finally: + _shutdown_kv_cache_manager(kv_cache_manager) if __name__ == "__main__": From 72fe8381713ee813232cf737c485735effd89d3c Mon Sep 17 00:00:00 2001 From: Alyosha-Swamy Date: Thu, 4 Jun 2026 11:18:09 +0000 Subject: [PATCH 6/8] [None][fix] Handle AFMoE tests with MPI disabled Signed-off-by: Alyosha-Swamy --- .../_torch/modeling/test_modeling_afmoe.py | 64 +++++++++++-------- 1 file changed, 38 insertions(+), 26 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_afmoe.py b/tests/unittest/_torch/modeling/test_modeling_afmoe.py index bbe1461bf9f2..acba3e8cc5f9 100644 --- a/tests/unittest/_torch/modeling/test_modeling_afmoe.py +++ b/tests/unittest/_torch/modeling/test_modeling_afmoe.py @@ -86,17 +86,15 @@ } -def _synthetic_tp_mapping(tp_size: int, enable_attention_dp: bool = False) -> Mapping: +def _force_mpi_topology_mapping(): # These tests inspect module TP attributes in one pytest process. Force # the lightweight MPI-topology Mapping even when TLLM_DISABLE_MPI=1 would # otherwise require an initialized torch.distributed DeviceMesh. - with patch("tensorrt_llm.mapping.mpi_disabled", return_value=False): - return Mapping( - world_size=tp_size, - tp_size=tp_size, - rank=0, - enable_attention_dp=enable_attention_dp, - ) + return patch("tensorrt_llm.mapping.mpi_disabled", return_value=False) + + +def _force_mpi_collectives(): + return patch("tensorrt_llm._torch.distributed.ops.mpi_disabled", return_value=False) def _shutdown_kv_cache_manager(kv_cache_manager: KVCacheManager) -> None: @@ -288,10 +286,16 @@ def test_afmoe_sanity(self): config_dict = deepcopy(AFMOE_CONFIG) afmoe_config = AfmoeConfig.from_dict(config_dict) - model_config = ModelConfig(pretrained_config=afmoe_config, quant_config=QuantConfig()) dtype = afmoe_config.torch_dtype device = torch.device("cuda") - model = AfmoeForCausalLM(model_config).to(device) + with _force_mpi_topology_mapping(): + mapping = Mapping(world_size=1, tp_size=1, rank=0) + model_config = ModelConfig( + pretrained_config=afmoe_config, + quant_config=QuantConfig(), + mapping=mapping, + ) + model = AfmoeForCausalLM(model_config).to(device) input_ids = torch.tensor( [100, 200, 300, 100, 200, 100, 400, 500], dtype=torch.int, device=device @@ -319,7 +323,6 @@ def test_afmoe_sanity(self): else: raise ValueError("Invalid dtype") - mapping = Mapping(world_size=1, tp_size=1, rank=0) kv_cache_config = KvCacheConfig(max_tokens=num_blocks * tokens_per_block) kv_cache_manager = KVCacheManager( kv_cache_config, @@ -358,7 +361,7 @@ def test_afmoe_sanity(self): position_ids = torch.cat(position_ids).unsqueeze(0) try: - with torch.inference_mode(): + with torch.inference_mode(), _force_mpi_collectives(): attn_metadata.prepare() logits = model.forward( input_ids=input_ids, position_ids=position_ids, attn_metadata=attn_metadata @@ -442,20 +445,27 @@ class TestAfmoeTPAttributes(unittest.TestCase): def _build_model(self, tp_size): config_dict = deepcopy(AFMOE_CONFIG) afmoe_config = AfmoeConfig.from_dict(config_dict) - mapping = _synthetic_tp_mapping(tp_size) - model_config = ModelConfig( - pretrained_config=afmoe_config, mapping=mapping, allreduce_strategy="NCCL" - ) - return AfmoeForCausalLM(model_config) + with _force_mpi_topology_mapping(): + mapping = Mapping(world_size=tp_size, tp_size=tp_size, rank=0) + model_config = ModelConfig( + pretrained_config=afmoe_config, mapping=mapping, allreduce_strategy="NCCL" + ) + return AfmoeForCausalLM(model_config) def _build_attention_dp_model(self, tp_size): config_dict = deepcopy(AFMOE_CONFIG) afmoe_config = AfmoeConfig.from_dict(config_dict) - mapping = _synthetic_tp_mapping(tp_size, enable_attention_dp=True) - model_config = ModelConfig( - pretrained_config=afmoe_config, mapping=mapping, allreduce_strategy="NCCL" - ) - return AfmoeForCausalLM(model_config) + with _force_mpi_topology_mapping(): + mapping = Mapping( + world_size=tp_size, + tp_size=tp_size, + rank=0, + enable_attention_dp=True, + ) + model_config = ModelConfig( + pretrained_config=afmoe_config, mapping=mapping, allreduce_strategy="NCCL" + ) + return AfmoeForCausalLM(model_config) def test_qkv_is_column_parallel_with_output_gate(self): model = self._build_model(tp_size=1) @@ -619,8 +629,10 @@ def test_afmoe_allclose_to_hf(self): norm_topk_prob=self.HF_CONFIG["route_norm"], ) afmoe_config = AfmoeConfig.from_dict(trt_config_dict) - model_config = ModelConfig(pretrained_config=afmoe_config) - model = AfmoeForCausalLM(model_config).to(dtype).to(device) + with _force_mpi_topology_mapping(): + mapping = Mapping(world_size=1, tp_size=1, rank=0) + model_config = ModelConfig(pretrained_config=afmoe_config, mapping=mapping) + model = AfmoeForCausalLM(model_config).to(dtype).to(device) weights = self._convert_hf_experts( hf_model.state_dict(), self.HF_CONFIG["moe_intermediate_size"] @@ -648,7 +660,7 @@ def test_afmoe_allclose_to_hf(self): tokens_per_block=tokens_per_block, max_seq_len=num_blocks * tokens_per_block, max_batch_size=1, - mapping=Mapping(world_size=1, tp_size=1, rank=0), + mapping=mapping, dtype=tensorrt_llm.bindings.DataType.BF16, ) kv_cache_manager.add_dummy_requests([0], [input_len]) @@ -666,7 +678,7 @@ def test_afmoe_allclose_to_hf(self): try: hf_position_ids = position_ids.to(torch.long) - with torch.inference_mode(): + with torch.inference_mode(), _force_mpi_collectives(): attn_metadata.prepare() logits = model.forward( input_ids=input_ids, position_ids=position_ids, attn_metadata=attn_metadata From 1397a7e100cc60f2e639500b38cecb55908805a6 Mon Sep 17 00:00:00 2001 From: Alyosha-Swamy Date: Thu, 4 Jun 2026 15:59:36 +0000 Subject: [PATCH 7/8] [None][fix] Guard AFMoE HF test imports Signed-off-by: Alyosha-Swamy --- .../_torch/modeling/test_modeling_afmoe.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_afmoe.py b/tests/unittest/_torch/modeling/test_modeling_afmoe.py index acba3e8cc5f9..01005600b8a5 100644 --- a/tests/unittest/_torch/modeling/test_modeling_afmoe.py +++ b/tests/unittest/_torch/modeling/test_modeling_afmoe.py @@ -43,8 +43,17 @@ from tensorrt_llm.llmapi import MoeConfig from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantConfig -from transformers import AfmoeConfig as HFAfmoeConfig -from transformers.models.afmoe.modeling_afmoe import AfmoeForCausalLM as HFAfmoeForCausalLM + +# AFMoE is a recent addition to HF transformers; older installed versions may +# not ship it. Guard the reference-model imports (matching the exaone4 test +# pattern) so the whole module still collects when HF afmoe is unavailable and +# only the HF parity test is skipped. +SKIP_AFMOE_HF_ACCURACY_TEST = False +try: + from transformers import AfmoeConfig as HFAfmoeConfig + from transformers.models.afmoe.modeling_afmoe import AfmoeForCausalLM as HFAfmoeForCausalLM +except ImportError: + SKIP_AFMOE_HF_ACCURACY_TEST = True WINDOW_SIZE = 4 NUM_HIDDEN_LAYERS = 4 @@ -544,6 +553,10 @@ def test_attention_layer_types(self): @unittest.skipUnless(torch.cuda.is_available(), "needs CUDA") +@unittest.skipIf( + SKIP_AFMOE_HF_ACCURACY_TEST, + "installed transformers does not provide the HF afmoe reference model", +) class TestAfmoeAllCloseToHF(unittest.TestCase): """Compare TRT-LLM AFMoE context-phase logits against the HF reference. From cb2645a7742cbb756de7c5c03460181cda862592 Mon Sep 17 00:00:00 2001 From: Alyosha-Swamy Date: Fri, 5 Jun 2026 17:45:37 +0000 Subject: [PATCH 8/8] [None][fix] Stabilize AFMoE sanity tests Signed-off-by: Alyosha-Swamy --- .../_torch/modeling/test_modeling_afmoe.py | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/tests/unittest/_torch/modeling/test_modeling_afmoe.py b/tests/unittest/_torch/modeling/test_modeling_afmoe.py index 01005600b8a5..934fa59b650a 100644 --- a/tests/unittest/_torch/modeling/test_modeling_afmoe.py +++ b/tests/unittest/_torch/modeling/test_modeling_afmoe.py @@ -299,10 +299,12 @@ def test_afmoe_sanity(self): device = torch.device("cuda") with _force_mpi_topology_mapping(): mapping = Mapping(world_size=1, tp_size=1, rank=0) + # Keep this model-wiring smoke test off backend-native attention kernels. model_config = ModelConfig( pretrained_config=afmoe_config, quant_config=QuantConfig(), mapping=mapping, + attn_backend="VANILLA", ) model = AfmoeForCausalLM(model_config).to(device) @@ -332,7 +334,12 @@ def test_afmoe_sanity(self): else: raise ValueError("Invalid dtype") - kv_cache_config = KvCacheConfig(max_tokens=num_blocks * tokens_per_block) + kv_cache_config = KvCacheConfig( + enable_block_reuse=False, + enable_partial_reuse=False, + copy_on_partial_reuse=False, + max_tokens=num_blocks * tokens_per_block, + ) kv_cache_manager = KVCacheManager( kv_cache_config, tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, @@ -644,7 +651,12 @@ def test_afmoe_allclose_to_hf(self): afmoe_config = AfmoeConfig.from_dict(trt_config_dict) with _force_mpi_topology_mapping(): mapping = Mapping(world_size=1, tp_size=1, rank=0) - model_config = ModelConfig(pretrained_config=afmoe_config, mapping=mapping) + # Keep HF parity focused on AFMoE weights/model math, not attention-kernel coverage. + model_config = ModelConfig( + pretrained_config=afmoe_config, + mapping=mapping, + attn_backend="VANILLA", + ) model = AfmoeForCausalLM(model_config).to(dtype).to(device) weights = self._convert_hf_experts( @@ -664,14 +676,20 @@ def test_afmoe_allclose_to_hf(self): position_ids = torch.arange(input_len, dtype=torch.int32, device=device).unsqueeze(0) num_blocks, tokens_per_block = 4, 128 + max_seq_len = num_blocks * tokens_per_block kv_cache_manager = KVCacheManager( - KvCacheConfig(max_tokens=num_blocks * tokens_per_block), + KvCacheConfig( + enable_block_reuse=False, + enable_partial_reuse=False, + copy_on_partial_reuse=False, + max_tokens=num_blocks * tokens_per_block, + ), tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF, num_layers=hf_config.num_hidden_layers, num_kv_heads=hf_config.num_key_value_heads, head_dim=hf_config.head_dim, tokens_per_block=tokens_per_block, - max_seq_len=num_blocks * tokens_per_block, + max_seq_len=max_seq_len, max_batch_size=1, mapping=mapping, dtype=tensorrt_llm.bindings.DataType.BF16,