diff --git a/3rdparty/fetch_content.json b/3rdparty/fetch_content.json index 31820f223a18..85a153113a25 100644 --- a/3rdparty/fetch_content.json +++ b/3rdparty/fetch_content.json @@ -32,7 +32,7 @@ { "name": "deepgemm", "git_repository": "https://github.com/deepseek-ai/DeepGEMM", - "git_tag": "f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca", + "git_tag": "8b1392b978f5a03c828dd1711090d7fb50958b8a", "git_submodules_recurse": true, "source_subdir": "dont-add-this-project-with-add-subdirectory" }, diff --git a/examples/kimi_k3/eval_extra_llm_options.yaml b/examples/kimi_k3/eval_extra_llm_options.yaml index 813fe8e0dbe8..d464e70f4c56 100644 --- a/examples/kimi_k3/eval_extra_llm_options.yaml +++ b/examples/kimi_k3/eval_extra_llm_options.yaml @@ -11,6 +11,8 @@ cuda_graph_config: enable_padding: true max_batch_size: 32 moe_config: + # TRTLLM chunking bound. Kimi's MegaMoE path privately raises this to + # max_num_tokens * dp_size for per-rank SymmBuffer capacity. max_num_tokens: 33024 use_low_precision_moe_combine: true kv_cache_config: diff --git a/scripts/attribution/data/dependency_metadata.yml b/scripts/attribution/data/dependency_metadata.yml index 2b8485cb1383..ef044adffc11 100644 --- a/scripts/attribution/data/dependency_metadata.yml +++ b/scripts/attribution/data/dependency_metadata.yml @@ -23,9 +23,9 @@ cutlass/v4.3.0: deep_ep/5be51b228a7c82dbdb213ea58e77bffd12b38af8: license: 452b3ef002dc6ec283fb723f0dd84997 source: https://github.com/deepseek-ai/DeepEP/tree/5be51b228a7c82dbdb213ea58e77bffd12b38af8 -deepgemm/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca: +deepgemm/8b1392b978f5a03c828dd1711090d7fb50958b8a: license: 452b3ef002dc6ec283fb723f0dd84997 - source: https://github.com/deepseek-ai/DeepGEMM/tree/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca + source: https://github.com/deepseek-ai/DeepGEMM/tree/8b1392b978f5a03c828dd1711090d7fb50958b8a dlpack/v1.0: license: cd9881918c97ec7b4962691660bb733e flashmla/1408756a88e52a25196b759eaf8db89d2b51b5a1: diff --git a/scripts/attribution/data/files_to_dependency.yml b/scripts/attribution/data/files_to_dependency.yml index 7a5c25002735..59682bd47f17 100644 --- a/scripts/attribution/data/files_to_dependency.yml +++ b/scripts/attribution/data/files_to_dependency.yml @@ -2093,7 +2093,7 @@ deep_ep/5be51b228a7c82dbdb213ea58e77bffd12b38af8: - dbc9ea8cf83b20e6ce4c6f383b700f29 - e4024308dd534f83de92752e1d7cd9a8 - f41ae95dbbafe6107dd98bf66af018ea -deepgemm/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca: +deepgemm/8b1392b978f5a03c828dd1711090d7fb50958b8a: - 07488ed395a262f652b63d0b1c1bb3a8 - 1101099cb0a9c0489f9e4e49a719941a - 111cfba37978c3bad26617cb42407970 @@ -2133,7 +2133,7 @@ deepgemm/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca: - 95b8e1175d420f5a05bc9fb0d0c55214 - 9bf6bbf8da71d31836279a87ffafe108 - 9e16e23b6894db65b377f93673dc733f -- a359f41b0ec5d67c1d9fda93dc65c5b5 +- 13ab873c8d6a0b7b2857bbfee317f0a4 - a3645795bcf4bad7c333094975c558b6 - a775f6a60d47cd428f4c04289d4e8cb3 - a9f87d66fb89c05e1ed20a9459c54f68 @@ -2157,7 +2157,7 @@ deepgemm/f8e8fb5830fa5cda6e4ea73d360bb3f21f87a3ca: - ea07df16c9a083277f55ae219c85f39d - ef5544cabdf0490063f2b2959f62a8cf - ef712bc72e01afbfdbceef82c6b49174 -- f43084b87a14bcf56920e1e3b1ad2dba +- cfcffb4170384e7699c65ed604fbd1d8 - f5a3009221d096c818cf26fbdb4d9693 - f5d68cc5860baa1de18f4aee4ccb0cc2 - f7f27b18dae31db0aad429bfaaf4615e diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 2d2bc9c314f2..9e0fa9a506ff 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -61,6 +61,11 @@ _DEEPSEEK_V4_ARCHITECTURES = {"DeepseekV4ForCausalLM"} _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT = "layers.0.ffn.experts.0.w1.weight" +_KIMI_K3_ARCHITECTURES = { + "KimiK3ForConditionalGeneration", + "KimiLinearForCausalLM", +} + _MINIMAX_M3_ARCHITECTURES = { "MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration", @@ -367,6 +372,9 @@ def resolve_moe_backend(moe_backend: str, if moe_backend.upper() != "AUTO": return moe_backend + if architecture in _KIMI_K3_ARCHITECTURES: + return "TRTLLM" + if architecture in _DEEPSEEK_V4_ARCHITECTURES: sm_version = get_sm_version() if 100 <= sm_version < 120: diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 913e1f618339..9f90bfc07b77 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -87,7 +87,7 @@ from ..attention_backend import AttentionMetadata from ..distributed import AllReduce, AllReduceStrategy from ..model_config import ModelConfig -from ..modules.fused_moe import ConfigurableMoE, create_moe +from ..modules.fused_moe import ConfigurableMoE, TRTLLMGenFusedMoE, create_moe from ..modules.kimi_k3_moe._mlp import KimiK3MLP, KimiK3RMSNorm from ..modules.kimi_k3_moe.kimi_k3_moe_gate import KimiK3MoEGate from ..modules.linear import Linear as TrtllmLinear @@ -710,7 +710,7 @@ def _swap(parent: nn.Module, attr: str) -> None: class KimiK3MoERuntime(nn.Module): - """Kimi K3 latent MoE block backed by ConfigurableMoE/TRTLLM-Gen.""" + """Kimi K3 latent MoE block backed by ConfigurableMoE.""" def __init__( self, @@ -751,7 +751,7 @@ def __init__( routed_moe_model_config = self._routed_moe_model_config(model_config) routed_quant_config = QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8) - self.routed_experts = create_moe( + routed_moe_kwargs = dict( routing_method=self.gate.routing_method, num_experts=self.num_experts, hidden_size=self.moe_hidden_size, @@ -765,20 +765,38 @@ def __init__( model_config=routed_moe_model_config, override_quant_config=routed_quant_config, layer_idx=layer_idx, - trtllm_gen_activation_type=ActType_TrtllmGen.SiTu, - # Cubin alpha is the gate-side SiTU beta; cubin beta is the - # linear-side SiTU beta. - trtllm_gen_activation_alpha=float(situ_beta), - trtllm_gen_activation_beta=float( - situ_linear_beta if situ_linear_beta is not None else 1.0 - ), # Let CommunicationFactory select the best available strategy. communication_method=None, ) + if routed_moe_model_config.moe_backend == "TRTLLM": + routed_moe_kwargs.update( + trtllm_gen_activation_type=ActType_TrtllmGen.SiTu, + # Cubin alpha is the gate-side SiTU beta; cubin beta is the + # linear-side SiTU beta. + trtllm_gen_activation_alpha=float(situ_beta), + trtllm_gen_activation_beta=float( + situ_linear_beta if situ_linear_beta is not None else 1.0 + ), + ) + elif routed_moe_model_config.moe_backend == "MEGAMOE_DEEPGEMM": + routed_moe_kwargs.update( + activation="situ", + situ_beta=float(situ_beta), + situ_linear_beta=float(situ_linear_beta if situ_linear_beta is not None else 1.0), + ) + self.routed_experts = create_moe(**routed_moe_kwargs) if not isinstance(self.routed_experts, ConfigurableMoE): raise RuntimeError( "Kimi K3 requires ConfigurableMoE; ENABLE_CONFIGURABLE_MOE must not be disabled." ) + if routed_moe_model_config.moe_backend == "MEGAMOE_DEEPGEMM": + from ..modules.fused_moe.mega_moe import MegaMoEDeepGemm + + if not isinstance(self.routed_experts.backend, MegaMoEDeepGemm): + raise RuntimeError( + "Kimi K3 explicitly requested MEGAMOE_DEEPGEMM, but the " + f"MoE factory selected {type(self.routed_experts.backend).__name__}." + ) if self.routed_experts.layer_load_balancer is not None: raise NotImplementedError( "Kimi K3 packed-checkpoint streaming does not yet support " @@ -801,6 +819,7 @@ def __init__( not _K3_DISABLE_MIN_LATENCY_LATENT_PROJ and not _K3_DISABLE_FUSED_LATENT_DOWN_MXFP8 and routed_comm is None + and isinstance(self.routed_experts.backend, TRTLLMGenFusedMoE) ) # Shared experts stay replicated (DeepSeek's attention-DP @@ -878,6 +897,13 @@ def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: """Build a private routed-expert mapping without mutating the shared config. Default split is EP-only; see ``_select_moe_tp_ep``.""" + supported_backends = {"TRTLLM", "MEGAMOE_DEEPGEMM"} + if model_config.moe_backend not in supported_backends: + raise ValueError( + "Kimi K3 SiTU routed experts only support the TRTLLM and " + "MEGAMOE_DEEPGEMM backends; " + f"got {model_config.moe_backend!r}." + ) if model_config.moe_load_balancer is not None: raise NotImplementedError( "Kimi K3 packed-checkpoint streaming does not yet support " @@ -915,7 +941,17 @@ def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: routed_model_config._frozen = False routed_model_config.extra_attrs = copy.copy(model_config.extra_attrs) routed_model_config.mapping = routed_mapping - routed_model_config.moe_backend = "TRTLLM" + routed_model_config.moe_backend = model_config.moe_backend + # MegaMoE uses this value as global DP SymmBuffer capacity, then + # divides it by EP size for the per-rank allocation. Other backends + # keep the user-configured value as their MoE chunking bound. + # Preserve an explicitly larger capacity. + if routed_model_config.moe_backend == "MEGAMOE_DEEPGEMM": + default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size + routed_model_config.moe_max_num_tokens = max( + int(routed_model_config.moe_max_num_tokens or 0), + default_moe_max_num_tokens, + ) routed_model_config._frozen = True return routed_model_config diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py index 23b82228d196..3c6a190d158d 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py +++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py @@ -118,7 +118,7 @@ class MegaMoEDeepGemm(MoE): _SUPPORTED_ACTIVATION_DTYPES = frozenset({torch.bfloat16}) - # Kernel owns dispatch + GEMM1 + SwiGLU + GEMM2 + combine via NVLink + # Kernel owns dispatch + GEMM1 + gated activation + GEMM2 + combine via NVLink # SymmBuffer; ConfigurableMoE must NOT layer host-side comm on top. scheduler_kind = MoESchedulerKind.FUSED_COMM @@ -203,11 +203,13 @@ def __init__( layer_idx: Optional[int] = None, activation_type: ActivationType = ActivationType.Swiglu, init_load_balancer: bool = True, - # DG tunables. ``swiglu_limit_scalar`` mirrors the upstream MoE - # kwarg; bridged to DG's ``activation_clamp`` at the call site. - activation: str = "swiglu", + # DG tunables. ``activation=None`` infers Kimi K3 SiTU from the + # pretrained config and otherwise defaults to SwiGLU. + activation: Optional[str] = None, swiglu_limit_scalar: Optional[float] = None, fast_math: bool = True, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, **kwargs, ) -> None: super().__init__( @@ -276,18 +278,27 @@ def __init__( "not equivalent. Use a different MoE backend for models that " "require pre-scaling, or extend the kernel call." ) - # DG's fp8_fp4_mega_moe currently only ships a fused SwiGLU - # activation path. Reject other ActivationType values explicitly so - # ``create_moe_backend`` callers do not silently get SwiGLU when - # they asked for GELU / etc. + # ``ActivationType.Swiglu`` describes the gated FC1 tensor geometry + # shared by SwiGLU and SiTU. The DeepGEMM-specific activation selects + # the actual elementwise function below. if activation_type != ActivationType.Swiglu: raise ValueError( f"MegaMoEDeepGemm only supports ActivationType.Swiglu (got {activation_type})." ) + activation, situ_beta, situ_linear_beta = self._resolve_activation_config( + model_config, + activation=activation, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, + ) + if activation == "situ" and swiglu_limit_scalar is not None: + raise ValueError("MegaMoEDeepGemm SiTU does not support activation_clamp.") self.apply_router_weight_on_input = apply_router_weight_on_input self.activation = activation self.swiglu_limit_scalar = swiglu_limit_scalar self.fast_math = fast_math + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta # Buffer sizing. MoE layers execute serially per forward; a single # process-level pool sized to worst-case per-rank tokens serves all. @@ -347,6 +358,46 @@ def __init__( if not model_config.skip_create_weights_in_init: self.create_weights() + @staticmethod + def _resolve_activation_config( + model_config: ModelConfig, + *, + activation: Optional[str], + situ_beta: Optional[float], + situ_linear_beta: Optional[float], + ) -> Tuple[str, Optional[float], Optional[float]]: + pretrained_config = model_config.pretrained_config + text_config = getattr(pretrained_config, "text_config", None) + config_situ_beta = getattr(pretrained_config, "activation_situ_beta", None) + config_situ_linear_beta = getattr(pretrained_config, "activation_situ_linear_beta", None) + if config_situ_beta is None: + config_situ_beta = getattr(text_config, "activation_situ_beta", None) + if config_situ_linear_beta is None: + config_situ_linear_beta = getattr(text_config, "activation_situ_linear_beta", None) + if activation is None: + activation = "situ" if config_situ_beta is not None else "swiglu" + activation = activation.lower() + if activation not in ("swiglu", "situ"): + raise ValueError( + f"MegaMoEDeepGemm activation must be 'swiglu' or 'situ'; got {activation!r}." + ) + if activation == "swiglu": + if situ_beta is not None or situ_linear_beta is not None: + raise ValueError("SiTU beta parameters require activation='situ'.") + return activation, None, None + + situ_beta = config_situ_beta if situ_beta is None else situ_beta + situ_linear_beta = config_situ_linear_beta if situ_linear_beta is None else situ_linear_beta + if situ_beta is None or situ_linear_beta is None: + raise ValueError( + "MegaMoEDeepGemm SiTU requires activation_situ_beta and " + "activation_situ_linear_beta in the pretrained config, or " + "explicit situ_beta and situ_linear_beta arguments." + ) + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError("MegaMoEDeepGemm SiTU beta parameters must be positive.") + return activation, float(situ_beta), float(situ_linear_beta) + def _supports_load_balancer(self) -> bool: # The DeepGEMM mega kernel routes by `topk_idx` interpreted as slot id # (range [0, num_slots)) once the SymmBuffer is sized to num_slots. @@ -711,5 +762,7 @@ def run_moe( activation=self.activation, activation_clamp=self.swiglu_limit_scalar, fast_math=self.fast_math, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, ) return y.to(output_dtype) diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index e134fd26b546..9dc728fe44a6 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -6427,6 +6427,7 @@ def create_weights(self, module: torch.nn.Module) -> None: # Downstream reload/EPLB metadata path; populated lazily when parameter # replacement records tensors that need rebuilding before reload. module.rebuild_tensor_metadata = {} + module._packed_mxfp4_loaded_slots = set() self.setup_quant_scales(module) def setup_quant_scales(self, module: torch.nn.Module): @@ -6511,6 +6512,109 @@ def _load_expert_weights_to_dst( w2_scale, dst_w2_weight_scale), non_blocking=True) + def load_packed_mxfp4_expert( + self, + module: torch.nn.Module, + *, + global_expert_id: int, + local_slot_id: int, + w1_weight: torch.Tensor, + w1_weight_scale: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale: torch.Tensor, + w3_weight: torch.Tensor, + w3_weight_scale: torch.Tensor, + ) -> None: + """Load one group-32 packed MXFP4 checkpoint expert into a local slot. + + Kimi K3 streams its packed checkpoint expert-by-expert so that each + safetensors mapping stays short-lived, and therefore calls this adapter + instead of ``load_weights``. The bytes written here are exactly what + ``_load_expert_weights_to_dst`` writes; only the source differs + (explicit tensors rather than a weight dict). + + ``transform_weights`` is not run here. The Kimi loader clears + ``_weights_transformed`` once every slot is filled, so the DG-native + layout is rebuilt lazily from the raw bytes staged below. + """ + if not 0 <= local_slot_id < module.expert_size_per_partition: + raise IndexError(f"local_slot_id={local_slot_id} is outside " + f"[0, {module.expert_size_per_partition}).") + expected_expert_id = module.initial_local_expert_ids[local_slot_id] + if global_expert_id != expected_expert_id: + raise ValueError( + f"local slot {local_slot_id} expects global expert " + f"{expected_expert_id}, got {global_expert_id}.") + + for name, value in ( + ("w1_weight", w1_weight), + ("w1_weight_scale", w1_weight_scale), + ("w2_weight", w2_weight), + ("w2_weight_scale", w2_weight_scale), + ("w3_weight", w3_weight), + ("w3_weight_scale", w3_weight_scale), + ): + if value.dtype != torch.uint8: + raise TypeError( + f"{name} must contain packed MXFP4 uint8 data, got " + f"{value.dtype}.") + + loaded_slots = module._packed_mxfp4_loaded_slots + if local_slot_id in loaded_slots: + raise ValueError( + f"Packed MXFP4 local slot {local_slot_id} was loaded twice.") + if not loaded_slots: + # Raw weights are about to change, so any DG-derived tensors left + # from an earlier load must not survive into transform_weights. + self._clear_transformed_weight_cache(module) + + dst_w3_w1_weight = module.w3_w1_weight.data + dst_w3_w1_weight_scale = module.w3_w1_weight_scale.data + dst_w2_weight = module.w2_weight.data + dst_w2_weight_scale = module.w2_weight_scale.data + + # DeepGEMM expects L1 in [gate | up] order before + # transform_weights_for_mega_moe interleaves gate/up rows, and TRT-LLM + # checkpoints map gate_proj -> w1 and up_proj -> w3. So despite the + # parameter being named w3_w1_*, the concatenation is [w1 | w3] -- + # the same order _load_expert_weights_to_dst uses. Reversing it stays + # shape-compatible and fails silently in the numerics, not loudly. + dst_w3_w1_weight[local_slot_id].copy_( + torch.cat( + [ + self._to_weight_device_uint8(w1_weight, dst_w3_w1_weight), + self._to_weight_device_uint8(w3_weight, dst_w3_w1_weight), + ], + dim=0, + ), + non_blocking=True, + ) + dst_w3_w1_weight_scale[local_slot_id].copy_( + torch.cat( + [ + self._to_weight_device_uint8(w1_weight_scale, + dst_w3_w1_weight_scale), + self._to_weight_device_uint8(w3_weight_scale, + dst_w3_w1_weight_scale), + ], + dim=0, + ), + non_blocking=True, + ) + dst_w2_weight[local_slot_id].copy_( + self._to_weight_device_uint8(w2_weight, dst_w2_weight), + non_blocking=True, + ) + dst_w2_weight_scale[local_slot_id].copy_( + self._to_weight_device_uint8(w2_weight_scale, dst_w2_weight_scale), + non_blocking=True, + ) + + loaded_slots.add(local_slot_id) + # transform_weights asserts this; the streaming path never reaches + # load_weights, which is where it would otherwise be set. + module._weights_loaded = True + def load_weights( self, module: torch.nn.Module, diff --git a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py index 6c716f2bccba..e17497be77f4 100644 --- a/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py +++ b/tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py @@ -26,9 +26,19 @@ import pytest import torch +import torch.distributed as dist from utils.util import check_accuracy +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.modeling_kimi_linear import ( + _K3_MOE_EP_ENV, + _K3_MOE_TP_ENV, + KimiK3MoERuntime, +) from tensorrt_llm._torch.modules.fused_moe.communication import CommunicationFactory +from tensorrt_llm._torch.modules.fused_moe.mega_moe.mega_moe_deepgemm import ( + _MEGA_MOE_SYMM_BUFFER_CACHE, +) from tensorrt_llm._torch.modules.kimi_k3_moe import KimiK3SparseMoeBlock from tensorrt_llm._torch.modules.kimi_k3_moe._moe_kernels import ( is_native_situ_supported, @@ -37,6 +47,8 @@ ) from tensorrt_llm._torch.modules.kimi_k3_moe.kimi_k3_moe_gate import KimiK3MoEGate from tensorrt_llm._torch.utils import ActType_TrtllmGen +from tensorrt_llm._utils import get_free_port +from tensorrt_llm.mapping import Mapping situ_supported = pytest.mark.skipif( not is_native_situ_supported(), @@ -44,6 +56,26 @@ ) +@pytest.fixture +def _single_rank_nccl_process_group(monkeypatch): + if dist.is_initialized(): + yield + return + monkeypatch.setenv("MASTER_ADDR", "127.0.0.1") + monkeypatch.setenv("MASTER_PORT", str(get_free_port())) + monkeypatch.setenv("RANK", "0") + monkeypatch.setenv("WORLD_SIZE", "1") + monkeypatch.setenv("LOCAL_RANK", "0") + torch.cuda.set_device(0) + dist.init_process_group(backend="nccl", rank=0, world_size=1) + try: + yield + finally: + _MEGA_MOE_SYMM_BUFFER_CACHE.clear() + if dist.is_initialized(): + dist.destroy_process_group() + + @dataclasses.dataclass class _K3Config: """Minimal config carrying the fields KimiK3SparseMoeBlock reads.""" @@ -419,8 +451,6 @@ def test_fused_forward_without_weights_raises(): def test_mapping_records_moe_tp_ep_user_specified(): - from tensorrt_llm.mapping import Mapping - # Auto default: -1 sentinels resolve to (moe_tp=tp, moe_ep=1) but must # NOT be flagged as a user request. auto = Mapping(world_size=8, tp_size=8) @@ -437,13 +467,6 @@ def test_mapping_records_moe_tp_ep_user_specified(): def test_kimi_k3_moe_split_selection(monkeypatch): - from tensorrt_llm._torch.models.modeling_kimi_linear import ( - _K3_MOE_EP_ENV, - _K3_MOE_TP_ENV, - KimiK3MoERuntime, - ) - from tensorrt_llm.mapping import Mapping - monkeypatch.delenv(_K3_MOE_TP_ENV, raising=False) monkeypatch.delenv(_K3_MOE_EP_ENV, raising=False) @@ -466,6 +489,65 @@ def test_kimi_k3_moe_split_selection(monkeypatch): assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (4, 2) +@pytest.mark.parametrize("backend", ["TRTLLM", "MEGAMOE_DEEPGEMM"]) +def test_kimi_k3_routed_config_preserves_explicit_backend(backend): + model_config = ModelConfig( + mapping=Mapping(world_size=1, rank=0, tp_size=1), + moe_backend=backend, + ) + + routed_model_config = KimiK3MoERuntime._routed_moe_model_config(model_config) + + assert routed_model_config.moe_backend == backend + assert model_config.moe_backend == backend + + +@pytest.mark.parametrize( + "backend,expected_moe_max_num_tokens", + [ + pytest.param("TRTLLM", 33024, id="trtllm"), + pytest.param("MEGAMOE_DEEPGEMM", 131072, id="megamoe"), + ], +) +def test_kimi_k3_routed_config_scopes_megamoe_capacity(backend, expected_moe_max_num_tokens): + model_config = ModelConfig( + mapping=Mapping( + world_size=16, + rank=0, + tp_size=16, + moe_ep_size=16, + enable_attention_dp=True, + ), + max_num_tokens=8192, + moe_max_num_tokens=33024, + moe_backend=backend, + ) + + routed_model_config = KimiK3MoERuntime._routed_moe_model_config(model_config) + + assert routed_model_config.moe_max_num_tokens == expected_moe_max_num_tokens + assert model_config.moe_max_num_tokens == 33024 + + +def test_kimi_k3_routed_config_rejects_backend_without_situ_support(): + model_config = ModelConfig( + mapping=Mapping(world_size=1, rank=0, tp_size=1), + moe_backend="CUTLASS", + ) + + with pytest.raises(ValueError, match="SiTU routed experts only support"): + KimiK3MoERuntime._routed_moe_model_config(model_config) + + assert model_config.moe_backend == "CUTLASS" + + +@pytest.mark.parametrize( + "architecture", ["KimiK3ForConditionalGeneration", "KimiLinearForCausalLM"] +) +def test_kimi_k3_moe_auto_backend_defaults_to_trtllm(architecture): + assert ModelConfig.resolve_moe_backend("AUTO", architecture) == "TRTLLM" + + # --------------------------------------------------------------------------- # MoE tensor-parallel shard parity (ConfigurableMoE / TRTLLM-Gen, GPU). # @@ -496,18 +578,20 @@ def scales(*shape): bank = [] for _ in range(num_experts): - bank.append({ - "w1": nibbles(intermediate, hidden // 2), - "w1_sf": scales(intermediate, hidden // 32), - "w3": nibbles(intermediate, hidden // 2), - "w3_sf": scales(intermediate, hidden // 32), - "w2": nibbles(hidden, intermediate // 2), - "w2_sf": scales(hidden, intermediate // 32), - }) + bank.append( + { + "w1": nibbles(intermediate, hidden // 2), + "w1_sf": scales(intermediate, hidden // 32), + "w3": nibbles(intermediate, hidden // 2), + "w3_sf": scales(intermediate, hidden // 32), + "w2": nibbles(hidden, intermediate // 2), + "w2_sf": scales(hidden, intermediate // 32), + } + ) return bank -def _make_test_gate(num_experts=_TP_EXPERTS, seed=71): +def _make_test_gate(num_experts=_TP_EXPERTS, top_k=_TP_TOPK, seed=71): """One deterministically-initialized gate SHARED by all modules under comparison: the fused routing kernel applies the gate's e_score_correction_bias per module, so a per-module `torch.empty` @@ -516,27 +600,29 @@ def _make_test_gate(num_experts=_TP_EXPERTS, seed=71): cfg = _K3Config( hidden_size=_TP_HIDDEN, num_experts=num_experts, - num_experts_per_token=min(_TP_TOPK, num_experts), + num_experts_per_token=min(top_k, num_experts), ) gate = KimiK3MoEGate(cfg) gen = torch.Generator().manual_seed(seed) with torch.no_grad(): - gate.weight.copy_( - torch.randn(gate.weight.shape, generator=gen, - dtype=torch.float32) * 0.05) + gate.weight.copy_(torch.randn(gate.weight.shape, generator=gen, dtype=torch.float32) * 0.05) gate.e_score_correction_bias.copy_( - torch.randn(gate.e_score_correction_bias.shape, generator=gen, - dtype=torch.float32) * 0.1) + torch.randn(gate.e_score_correction_bias.shape, generator=gen, dtype=torch.float32) + * 0.1 + ) return gate.cuda() -def _make_routed_moe(intermediate_size, gate, num_experts=_TP_EXPERTS): +def _make_routed_moe( + intermediate_size, + gate, + num_experts=_TP_EXPERTS, + moe_backend="TRTLLM", +): """Mirror KimiK3MoERuntime's create_moe call on a single-rank mapping.""" from transformers.configuration_utils import PretrainedConfig - from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.modules.fused_moe import ConfigurableMoE, create_moe - from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig pretrained_config = PretrainedConfig() @@ -544,12 +630,14 @@ def _make_routed_moe(intermediate_size, gate, num_experts=_TP_EXPERTS): pretrained_config.hidden_size = _TP_HIDDEN pretrained_config.intermediate_size = intermediate_size pretrained_config.torch_dtype = torch.bfloat16 + pretrained_config.activation_situ_beta = 4.0 + pretrained_config.activation_situ_linear_beta = 25.0 model_config = ModelConfig( pretrained_config=pretrained_config, mapping=Mapping(), - moe_backend="TRTLLM", + moe_backend=moe_backend, ) - moe = create_moe( + moe_kwargs = dict( routing_method=gate.routing_method, num_experts=num_experts, hidden_size=_TP_HIDDEN, @@ -559,11 +647,15 @@ def _make_routed_moe(intermediate_size, gate, num_experts=_TP_EXPERTS): model_config=model_config, override_quant_config=QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8), layer_idx=0, - trtllm_gen_activation_type=ActType_TrtllmGen.SiTu, - trtllm_gen_activation_alpha=4.0, - trtllm_gen_activation_beta=25.0, communication_method=None, - ).cuda() + ) + if moe_backend == "TRTLLM": + moe_kwargs.update( + trtllm_gen_activation_type=ActType_TrtllmGen.SiTu, + trtllm_gen_activation_alpha=4.0, + trtllm_gen_activation_beta=25.0, + ) + moe = create_moe(**moe_kwargs).cuda() assert isinstance(moe, ConfigurableMoE) return moe @@ -577,29 +669,47 @@ def _load_bank(moe, bank, tp_size=1, tp_rank=0): shard-sized parameters. """ backend = moe.backend - proxy = SimpleNamespace( - expert_size_per_partition=backend.expert_size_per_partition, - initial_local_expert_ids=backend.initial_local_expert_ids, - scaling_vector_size=backend.scaling_vector_size, - tp_size=tp_size, - tp_rank=tp_rank, - w3_w1_weight=backend.w3_w1_weight, - w2_weight=backend.w2_weight, - w3_w1_weight_scale=backend.w3_w1_weight_scale, - w2_weight_scale=backend.w2_weight_scale, - ) - for expert_id, tensors in enumerate(bank): - backend.quant_method.load_packed_mxfp4_expert( - proxy, - global_expert_id=expert_id, - local_slot_id=expert_id, - w1_weight=tensors["w1"], - w1_weight_scale=tensors["w1_sf"], - w2_weight=tensors["w2"], - w2_weight_scale=tensors["w2_sf"], - w3_weight=tensors["w3"], - w3_weight_scale=tensors["w3_sf"], - ) + if hasattr(backend.quant_method, "load_packed_mxfp4_expert"): + loader_module = backend + if hasattr(backend, "scaling_vector_size"): + loader_module = SimpleNamespace( + expert_size_per_partition=backend.expert_size_per_partition, + initial_local_expert_ids=backend.initial_local_expert_ids, + scaling_vector_size=backend.scaling_vector_size, + tp_size=tp_size, + tp_rank=tp_rank, + w3_w1_weight=backend.w3_w1_weight, + w2_weight=backend.w2_weight, + w3_w1_weight_scale=backend.w3_w1_weight_scale, + w2_weight_scale=backend.w2_weight_scale, + ) + for expert_id, tensors in enumerate(bank): + backend.quant_method.load_packed_mxfp4_expert( + loader_module, + global_expert_id=expert_id, + local_slot_id=expert_id, + w1_weight=tensors["w1"], + w1_weight_scale=tensors["w1_sf"], + w2_weight=tensors["w2"], + w2_weight_scale=tensors["w2_sf"], + w3_weight=tensors["w3"], + w3_weight_scale=tensors["w3_sf"], + ) + else: + weights = {} + for expert_id, tensors in enumerate(bank): + prefix = f"{expert_id}." + weights.update( + { + prefix + "w1.weight": tensors["w1"], + prefix + "w1.weight_scale": tensors["w1_sf"], + prefix + "w2.weight": tensors["w2"], + prefix + "w2.weight_scale": tensors["w2_sf"], + prefix + "w3.weight": tensors["w3"], + prefix + "w3.weight_scale": tensors["w3_sf"], + } + ) + backend.load_weights([weights]) backend._weights_transformed = False moe.post_load_weights() return moe @@ -626,23 +736,24 @@ def test_tp_shard_loader_matches_manual_slice(tp_size): rows = slice(tp_rank * ipp, (tp_rank + 1) * ipp) cols_packed = slice(tp_rank * (ipp // 2), (tp_rank + 1) * (ipp // 2)) cols_sf = slice(tp_rank * (ipp // 32), (tp_rank + 1) * (ipp // 32)) - manual_bank = [{ - "w1": e["w1"][rows].contiguous(), - "w1_sf": e["w1_sf"][rows].contiguous(), - "w3": e["w3"][rows].contiguous(), - "w3_sf": e["w3_sf"][rows].contiguous(), - "w2": e["w2"][:, cols_packed].contiguous(), - "w2_sf": e["w2_sf"][:, cols_sf].contiguous(), - } for e in bank] + manual_bank = [ + { + "w1": e["w1"][rows].contiguous(), + "w1_sf": e["w1_sf"][rows].contiguous(), + "w3": e["w3"][rows].contiguous(), + "w3_sf": e["w3_sf"][rows].contiguous(), + "w2": e["w2"][:, cols_packed].contiguous(), + "w2_sf": e["w2_sf"][:, cols_sf].contiguous(), + } + for e in bank + ] via_manual = _make_routed_moe(ipp, gate, num_experts=num_experts) _load_bank(via_manual, manual_bank) - for name in ("w3_w1_weight", "w2_weight", "w3_w1_weight_scale", - "w2_weight_scale"): + for name in ("w3_w1_weight", "w2_weight", "w3_w1_weight_scale", "w2_weight_scale"): a = getattr(via_shard.backend, name).data b = getattr(via_manual.backend, name).data - assert torch.equal(a, b), ( - f"{name} mismatch for tp_size={tp_size} tp_rank={tp_rank}") + assert torch.equal(a, b), f"{name} mismatch for tp_size={tp_size} tp_rank={tp_rank}" @situ_supported @@ -663,14 +774,12 @@ def test_tp8_sharded_forward_matches_whole_expert(num_tokens): _load_bank(whole, bank) torch.manual_seed(3) - x = torch.randn( - num_tokens, _TP_HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.5 + x = torch.randn(num_tokens, _TP_HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.5 router_logits = gate.compute_logits(x) out_whole = whole.forward(x, router_logits, all_rank_num_tokens=None) - partial_sum = torch.zeros(num_tokens, _TP_HIDDEN, dtype=torch.float32, - device="cuda") + partial_sum = torch.zeros(num_tokens, _TP_HIDDEN, dtype=torch.float32, device="cuda") for tp_rank in range(tp_size): shard = _make_routed_moe(ipp, gate) _load_bank(shard, bank, tp_size=tp_size, tp_rank=tp_rank) @@ -679,5 +788,70 @@ def test_tp8_sharded_forward_matches_whole_expert(num_tokens): del shard torch.cuda.empty_cache() - check_accuracy(partial_sum.to(torch.bfloat16), out_whole, - atol=0.08, rtol=0.08, percent=0.98) + check_accuracy(partial_sum.to(torch.bfloat16), out_whole, atol=0.08, rtol=0.08, percent=0.98) + + +@situ_supported +@pytest.mark.parametrize("num_tokens", [1, 16, 128], ids=lambda n: f"tokens{n}") +@pytest.mark.parametrize( + "num_experts,top_k", + [ + pytest.param(8, 1, id="experts8-top1"), + pytest.param(8, 2, id="experts8-top2"), + pytest.param(32, 16, id="experts32-top16"), + ], +) +def test_megamoe_deepgemm_situ_matches_trtllm_gen( + num_tokens, num_experts, top_k, _single_rank_nccl_process_group +): + """Compare SiTU kernels with identical packed MXFP4 weights and routing. + + MegaMoE folds routing weights into the FC1 activation before its MXFP8 + requantization, while TRTLLM-Gen combines after the expert output. The + quantized graphs therefore need semantic, rather than elementwise, + parity: high cosine similarity and bounded relative L2 error. + """ + bank = _make_packed_expert_bank(num_experts, _TP_INTERMEDIATE, _TP_HIDDEN) + gate = _make_test_gate(num_experts=num_experts, top_k=top_k) + + trtllm_gen = _load_bank( + _make_routed_moe(_TP_INTERMEDIATE, gate, num_experts=num_experts), + bank, + ) + mega_moe = _load_bank( + _make_routed_moe( + _TP_INTERMEDIATE, + gate, + num_experts=num_experts, + moe_backend="MEGAMOE_DEEPGEMM", + ), + bank, + ) + + torch.manual_seed(37) + x = torch.randn(num_tokens, _TP_HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.5 + router_logits = gate.compute_logits(x) + + with torch.inference_mode(): + trtllm_gen_output = trtllm_gen(x, router_logits) + mega_moe_output = mega_moe(x, router_logits) + + assert torch.isfinite(mega_moe_output).all() + diff = (mega_moe_output.float() - trtllm_gen_output.float()).abs() + ref = trtllm_gen_output.float() + cosine = torch.nn.functional.cosine_similarity( + mega_moe_output.float().flatten(), + ref.flatten(), + dim=0, + ) + relative_l2 = torch.linalg.vector_norm(diff) / torch.linalg.vector_norm(ref) + print( + f"M={num_tokens}, experts={num_experts}, top_k={top_k}: " + f"cosine={cosine.item():.8f}, " + f"relative_l2={relative_l2.item():.8f}, " + f"mean_abs={diff.mean().item():.8f}, " + f"p99_abs={torch.quantile(diff, 0.99).item():.8f}, " + f"max_abs={diff.max().item():.8f}" + ) + assert cosine > 0.998 + assert relative_l2 < 0.06 diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index 8223e08d67bb..2d9840e993f6 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -467,6 +467,43 @@ def test_megamoe_deepgemm_cache_derived_state_allocates_symm_buffer(): quant_method.cache_derived_state.assert_called_once_with(moe) +def test_megamoe_deepgemm_infers_kimi_situ_from_pretrained_config(): + model_config = ModelConfig( + pretrained_config=SimpleNamespace( + text_config=SimpleNamespace( + activation_situ_beta=4.0, + activation_situ_linear_beta=25.0, + ) + ) + ) + + activation, situ_beta, situ_linear_beta = MegaMoEDeepGemm._resolve_activation_config( + model_config, + activation=None, + situ_beta=None, + situ_linear_beta=None, + ) + + assert activation == "situ" + assert situ_beta == 4.0 + assert situ_linear_beta == 25.0 + + +def test_megamoe_deepgemm_defaults_to_swiglu_without_situ_config(): + model_config = ModelConfig(pretrained_config=SimpleNamespace()) + + activation, situ_beta, situ_linear_beta = MegaMoEDeepGemm._resolve_activation_config( + model_config, + activation=None, + situ_beta=None, + situ_linear_beta=None, + ) + + assert activation == "swiglu" + assert situ_beta is None + assert situ_linear_beta is None + + def test_megamoe_init_rejects_uneven_num_slots_with_value_error(): routing_method = RenormalizeMoeRoutingMethod(top_k=1) model_config = ModelConfig(