From ce9408707214a95fa128fcaa505eb5cd68a65cf4 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Sat, 15 Aug 2026 19:24:53 -0700 Subject: [PATCH 1/9] [None][fix] Add SiTu to the Python gated-activation list is_gated_activation is documented to stay aligned with isGatedActivation in cpp/tensorrt_llm/kernels/cutlass_kernels/include/ moe_gemm_kernels.h. The C++ side lists Swiglu, Geglu, SwigluBias and SiTu; the Python side was never updated when SiTu was added. The Python list is what feeds is_gated_activation -> intermediate_size_expand_ratio, so a MoE configured with ActivationType.SiTu sized w3_w1_weight for a non-gated activation - half the rows a gated FC1 needs. Nothing complains on the way there. The Cutlass NVFP4 loader concatenates the two halves and fits the result to the destination with torch.nn.functional.pad, and a negative pad truncates rather than raising, so a wrongly-sized destination reads as a successful load. The first and only complaint was a shape check inside the kernel at warmup: RuntimeError: fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size. The test parses isGatedActivation out of the header and asserts the two sets are equal, so the alignment the comment asks for is now enforced rather than remembered. Found while bringing up Kimi K3 NVFP4, whose CUTLASS MoE path is currently the only caller passing ActivationType.SiTu - so this cannot move any other model's geometry. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- tensorrt_llm/_torch/utils.py | 3 +- .../_torch/test_gated_activation_parity.py | 37 +++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) create mode 100644 tests/unittest/_torch/test_gated_activation_parity.py diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 00bcd5c6d96d..7cf9e33b77f0 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -88,7 +88,8 @@ class ActType_TrtllmGen(IntEnum): # And make sure it aligned with cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h::isGatedActivation function. def is_gated_activation(activation_type: ActivationType) -> bool: return activation_type in [ - ActivationType.Swiglu, ActivationType.SwigluBias, ActivationType.Geglu + ActivationType.Swiglu, ActivationType.SwigluBias, ActivationType.Geglu, + ActivationType.SiTu ] diff --git a/tests/unittest/_torch/test_gated_activation_parity.py b/tests/unittest/_torch/test_gated_activation_parity.py new file mode 100644 index 000000000000..b56dee801d1f --- /dev/null +++ b/tests/unittest/_torch/test_gated_activation_parity.py @@ -0,0 +1,37 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""The Python and C++ gated-activation lists must agree.""" + +import re +from pathlib import Path + +from tensorrt_llm._torch.utils import ActivationType, is_gated_activation + +_HEADER = ( + Path(__file__).parents[3] + / "cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h" +) + + +def test_is_gated_activation_matches_the_cutlass_header(): + """``is_gated_activation`` decides ``intermediate_size_expand_ratio`` and + therefore the whole FC1 weight geometry, while the kernel validates the + shapes it receives against the C++ list. When the two disagree, a MoE + allocates w3_w1 for the wrong geometry and the loader pads/truncates into + it without complaint -- ``torch.nn.functional.pad`` with a negative pad + truncates rather than raising -- so the first sign is a shape check deep + inside the kernel, or nothing at all. + + The comment above the Python function already asks for this alignment. + This asserts it. + """ + body = re.search( + r"constexpr bool isGatedActivation\(ActivationType activation_type\)\s*\{(.*?)\}", + _HEADER.read_text(), + re.DOTALL, + ) + assert body is not None, f"isGatedActivation not found in {_HEADER}" + cpp_gated = set(re.findall(r"ActivationType::(\w+)", body.group(1))) + + python_gated = {a.name for a in ActivationType if is_gated_activation(a)} + assert python_gated == cpp_gated From dabdb41b6822e11535637fc00256b3906629318d Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Sat, 15 Aug 2026 19:25:50 -0700 Subject: [PATCH 2/9] [None][fix] Record the checkpoint directory on lazily-loaded weights _load_lazy_safetensors returns PySafeSlice objects so a model can stream a huge checkpoint and read only its rank-local shard. A slice does not carry the file it came from, so a model that wants to re-open shards itself - to keep each safetensors handle short-lived instead of holding the whole mapping open for the duration of the load - has no way back to the directory. The available fallback no longer works: transformers no longer populates PretrainedConfig._name_or_path, so a model guarding its per-shard path on finding an index under that directory silently takes the other branch forever. Kimi K3 does exactly that, and the result was not a wrong answer but an OOM - every shard stayed mapped, which is the failure the per-shard path exists to prevent. It went unnoticed because a smaller topology maps less per node and survived it. The loader that opened the directory is the authoritative source, so it now records it on the dict it returns. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../models/checkpoints/hf/weight_loader.py | 9 +++++- .../_torch/models/modeling_kimi_k25.py | 4 +++ .../_torch/modeling/test_modeling_kimi_k25.py | 21 ++++++++++++++ .../checkpoints/hf/test_weight_loader.py | 28 +++++++++++++++++++ 4 files changed, 61 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py index fb05f96a7ce9..204428c33ec7 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py @@ -293,7 +293,14 @@ def _load_lazy_safetensors( self._lazy_handles = handles logger.info(f"Lazily opened {len(weight_files)} safetensors files " f"({len(weights)} tensors) from {checkpoint_dir}") - return ConsumableWeightsDict(weights) + lazy_weights = ConsumableWeightsDict(weights) + # A lazy slice does not carry the file it came from, and a model that + # wants to re-open shards itself (Kimi K3 streams rank-local experts + # per shard file, precisely to avoid holding this mapping open) has no + # other reliable source: transformers no longer sets + # ``PretrainedConfig._name_or_path``. + lazy_weights.checkpoint_dir = checkpoint_dir + return lazy_weights def load_weights(self, checkpoint_dir: str, diff --git a/tensorrt_llm/_torch/models/modeling_kimi_k25.py b/tensorrt_llm/_torch/models/modeling_kimi_k25.py index 22484c533f2b..d5ad4564a43f 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_k25.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_k25.py @@ -1733,6 +1733,10 @@ def load_weights(self, weights) -> None: if any(k.startswith(self._LANG_PREFIX) for k in weights): lm_weights = filter_weights("language_model", weights) lm_weights = ConsumableWeightsDict(lm_weights) + checkpoint_dir = getattr(weights, "checkpoint_dir", None) + if checkpoint_dir is not None: + lm_weights.checkpoint_dir = checkpoint_dir + lm_weights.checkpoint_prefix = self._LANG_PREFIX else: lm_weights = weights self.llm.load_weights(lm_weights) diff --git a/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py b/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py index 232893dfca3e..b1c045dc083c 100644 --- a/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py +++ b/tests/unittest/_torch/modeling/test_modeling_kimi_k25.py @@ -39,7 +39,9 @@ import tempfile import unittest from copy import deepcopy +from types import SimpleNamespace from typing import Optional, Tuple +from unittest import mock import numpy as np import pytest @@ -49,6 +51,7 @@ from utils.util import skip_pre_blackwell_unittest from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict from tensorrt_llm._torch.models.modeling_kimi_k25 import ( KimiK25ForConditionalGeneration, KimiK25InputProcessor, @@ -257,6 +260,24 @@ def test_model_instantiation(self): # Media placeholder token ID self.assertEqual(model._media_placeholder_token_id, 163605) + def test_language_weights_preserve_checkpoint_dir(self): + """The text backbone retains metadata needed for shard streaming.""" + weights = ConsumableWeightsDict({"language_model.foo": object()}) + weights.checkpoint_dir = "/checkpoint" + mm_encoder = torch.nn.Module() + mm_encoder.load_weights = mock.Mock() + model = SimpleNamespace( + _LANG_PREFIX="language_model.", mm_encoder=mm_encoder, llm=mock.Mock() + ) + + KimiK25ForConditionalGeneration.load_weights(model, weights) + + lm_weights = model.llm.load_weights.call_args.args[0] + self.assertIsInstance(lm_weights, ConsumableWeightsDict) + self.assertEqual(lm_weights.checkpoint_dir, "/checkpoint") + self.assertEqual(lm_weights.checkpoint_prefix, "language_model.") + self.assertIn("foo", lm_weights) + def test_vision_encoder_uses_trtllm_modules(self): """Vision blocks reuse TRT-LLM normalization, attention, linear, and MLP modules.""" config = _build_hf_config() diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py index 9fc73d04050e..27defe3720be 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py @@ -491,3 +491,31 @@ def test_prefetch_files_emits_progress_heartbeat(tmp_path, monkeypatch): # Every chunk logs when the interval is zero: 3 files x 4 KB at a 1 KB # chunk size means at least 12 heartbeats (short reads only add more). assert len(progress_logs) >= 12 + + +def test_kimi_k3_lazy_load_records_the_checkpoint_dir(tmp_path): + """A model that re-opens shards itself needs the directory back. + + Kimi K3 streams its rank-local experts per shard file to avoid holding + the whole mapping open. A lazy slice does not carry its file, and + transformers 5.x no longer sets ``PretrainedConfig._name_or_path``, so + without this the model silently fell back to the shared mapping and the + step was OOM-killed. + """ + import json + + import safetensors.torch + import torch + + (tmp_path / "config.json").write_text(json.dumps({"model_type": "kimi_k3"})) + safetensors.torch.save_file( + {"w": torch.zeros(2, 2)}, tmp_path / "model-00001-of-00001.safetensors" + ) + + loader = HfWeightLoader() + try: + weights = loader.load_weights(str(tmp_path), Mapping()) + assert isinstance(weights, ConsumableWeightsDict) + assert weights.checkpoint_dir == str(tmp_path) + finally: + loader.cleanup() From 57e11fc5ebaead464bb01eed18190f6dd40290aa Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Sun, 16 Aug 2026 02:52:07 -0700 Subject: [PATCH 3/9] [None][feat] Kimi K3: serve the nvidia/Kimi-K3-NVFP4 checkpoint Brings up nvidia/Kimi-K3-NVFP4 on the PyTorch backend and validates it at DEP16 against the MXFP4 baseline: GSM8K 96.40 vs 96.47, GPQA-Diamond 91.92 +/- 1.94 vs a published 92.77 NVFP4 / 93.21 original. Both are inside their error bars. Four things blocked loading the checkpoint at all. The config parse. 2790 attention entries are spelled FP8_PB_WO in quantized_layers and the alias matched lowercase only, and only at the top level, so each raised "'FP8_PB_WO' is not a valid QuantAlgo". Canonicalization is now shared, case-insensitive, and applied to the top level and to every per-layer entry; unknown values still pass through so QuantAlgo keeps ownership of rejecting bad names. The routed-expert quantization was hardcoded to W4A8_MXFP4_MXFP8. It is now resolved per layer from the checkpoint, falling back to that value when nothing is declared - which is exactly the original moonshotai/Kimi-K3, so its behaviour is unchanged. The expert loader only knew the MXFP4 packed layout, at the model level and independently of the MoE backend, so the first NVFP4 run died with 30912 missing keys. Rather than writing an NVFP4 loader, load_streaming_nvfp4_expert calls exactly the primitives the whole-checkpoint path calls, on the same destinations, leaving the same staging state - so every backend-specific decision stays owned by the backend. That is what makes the [w1 | w3] ordering trap a non-issue: Cutlass concatenates [w3 | w1], and delegating means never encoding a guess about which. And the attention weights, which is what the accuracy collapse turned out to be. This checkpoint stores them FP8 E4M3 plus a 128x128 FP32 block scale where the original stores BF16 - at the SAME shape, so the loader's shape check passed, src.to(param.dtype) converted quantized values as though they were real ones, and weight_scale sat in the checkpoint as a key no parameter asked for. Every projection came out wrong by its block scale and nothing raised; the model loaded clean and answered nonsense. FP8 without a companion scale now raises. Also here: per-expert draining of the Cutlass w3_w1 staging and per-layer lazy preparation, which together bound a footprint that would otherwise grow with the load; DEP16 configs (DEP8 does not fit - see the plan doc); a task-agnostic launcher; and the tests, registered individually in l0_b300.yml because the module carries 19 failures that predate this work. Depends on two generic fixes that ship separately, because neither is K3-specific: - is_gated_activation was missing SiTu, which the C++ side has. The Python list feeds intermediate_size_expand_ratio, so CUTLASS sized w3_w1 for a non-gated activation and F.pad's negative-pad truncation hid it. - The lazy weight loader now records the directory it opened; transformers 5.x no longer sets _name_or_path, so the per-shard streaming path had silently never run. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../kernels/cutlass_kernels/include/common.h | 1 + .../include/moe_gemm_kernels.h | 2 +- .../cutlass_kernels/include/moe_kernels.h | 4 +- .../cutlass_kernels/moe_gemm/moe_kernels.cu | 7 + .../cutlass_kernels/moe_gemm/moe_kernels.cuh | 16 + cpp/tensorrt_llm/thop/moeOp.cpp | 30 +- .../eval_extra_llm_options_nvfp4_dep16.yaml | 48 + ...al_extra_llm_options_nvfp4_dep16_gpqa.yaml | 37 + .../eval_extra_llm_options_nvfp4_dep8.yaml | 30 + examples/kimi_k3/run_eval_kimi_k3.sbatch | 31 +- tensorrt_llm/_torch/auto_deploy/_compat.py | 26 +- tensorrt_llm/_torch/model_config.py | 16 +- .../_torch/models/modeling_kimi_linear.py | 576 +++++++++++- .../_torch/modules/fused_moe/quantization.py | 239 ++++- tensorrt_llm/_torch/utils.py | 9 +- tensorrt_llm/quantization/modelopt_config.py | 50 +- .../test_lists/test-db/l0_b300.yml | 17 + .../modules/moe/test_kimi_k3_situ_moe.py | 853 +++++++++++++++++- tests/unittest/llmapi/test_llm_quant.py | 42 + 19 files changed, 1956 insertions(+), 78 deletions(-) create mode 100644 examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16.yaml create mode 100644 examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_gpqa.yaml create mode 100644 examples/kimi_k3/eval_extra_llm_options_nvfp4_dep8.yaml diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/common.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/common.h index 8a9937c62099..5db1586d1a4a 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/common.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/common.h @@ -36,6 +36,7 @@ enum class ActivationType Geglu = 6, SwigluBias = 7, Relu2 = 8, + SiTu = 9, }; } // namespace kernels::cutlass_kernels diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h index dca53bd071d2..550fde1d9fbb 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h @@ -251,7 +251,7 @@ struct TmaWarpSpecializedGroupedGemmInput constexpr bool isGatedActivation(ActivationType activation_type) { return activation_type == ActivationType::Swiglu || activation_type == ActivationType::Geglu - || activation_type == ActivationType::SwigluBias; + || activation_type == ActivationType::SwigluBias || activation_type == ActivationType::SiTu; } template > : activation_type == ActivationType::SwigluBias ? &doGatedActivationKernel + : activation_type == ActivationType::SiTu + ? &doGatedActivationKernel : nullptr; TLLM_CHECK_WITH_INFO(fn != nullptr, "Invalid activation type"); fn<<>>(output, gemm_result, expert_first_token_offset, inter_size, num_experts_per_node, @@ -2809,6 +2811,9 @@ void doActivation(T* output, GemmOutputType const* gemm_result, float const* fp8 case ActivationType::SwigluBias: return &doActivationKernel; + case ActivationType::SiTu: + return &doActivationKernel; case ActivationType::Relu2: return &doActivationKernel, decltype(block_scaling_type)::value, @@ -2963,6 +2968,8 @@ void doActivationDynamic(T* output, GemmOutputType const* gemm_result, float con case ActivationType::SwigluBias: return &doActivationKernel; + case ActivationType::SiTu: + return &doActivationKernel; case ActivationType::Relu2: return &doActivationKernel, NVFP4_TYPE, kRows, true>; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cuh b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cuh index 36e271228d74..893ab819b75e 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cuh +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cuh @@ -75,6 +75,22 @@ struct SwigluBiasAdaptor } }; +struct SiTuAdaptor +{ + constexpr static bool IS_GLU = true; + float alpha = 1.0f; + float beta = 1.0f; + float limit = std::numeric_limits::infinity(); + + template + __device__ T operator()(T const& gate, T const& linear) const + { + cutlass::epilogue::thread::Sigmoid sigmoid{}; + cutlass::epilogue::thread::Tanh tanhFn{}; + return tanhFn(gate * (1.0f / alpha)) * alpha * sigmoid(gate) * tanhFn(linear * (1.0f / beta)) * beta; + } +}; + } // namespace kernels::cutlass_kernels TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index fbce03c01712..e7f9b1ff97a5 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -543,21 +543,30 @@ class FusedMoeRunner : public torch::CustomClassHolder CHECK_INPUT(swiglu_alpha.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_alpha.value().sizes()[0] == num_experts_on_rank, "swiglu_alpha must have num_experts_on_rank elements."); - base_activation_type = ActivationType::SwigluBias; + if (base_activation_type != ActivationType::SiTu) + { + base_activation_type = ActivationType::SwigluBias; + } } if (swiglu_beta.has_value()) { CHECK_INPUT(swiglu_beta.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_beta.value().sizes()[0] == num_experts_on_rank, "swiglu_beta must have num_experts_on_rank elements."); - base_activation_type = ActivationType::SwigluBias; + if (base_activation_type != ActivationType::SiTu) + { + base_activation_type = ActivationType::SwigluBias; + } } if (swiglu_limit.has_value()) { CHECK_INPUT(swiglu_limit.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_limit.value().sizes()[0] == num_experts_on_rank, "swiglu_limit must have num_experts_on_rank elements."); - base_activation_type = ActivationType::SwigluBias; + if (base_activation_type != ActivationType::SiTu) + { + base_activation_type = ActivationType::SwigluBias; + } } auto activation_params = ActivationParams(base_activation_type, reinterpret_cast(swiglu_alpha.has_value() ? swiglu_alpha.value().const_data_ptr() : nullptr), @@ -797,21 +806,30 @@ class FusedMoeRunner : public torch::CustomClassHolder CHECK_INPUT(swiglu_alpha.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_alpha.value().sizes()[0] == num_experts_on_rank, "swiglu_alpha must have num_experts_on_rank elements."); - base_activation_type = ActivationType::SwigluBias; + if (base_activation_type != ActivationType::SiTu) + { + base_activation_type = ActivationType::SwigluBias; + } } if (swiglu_beta.has_value()) { CHECK_INPUT(swiglu_beta.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_beta.value().sizes()[0] == num_experts_on_rank, "swiglu_beta must have num_experts_on_rank elements."); - base_activation_type = ActivationType::SwigluBias; + if (base_activation_type != ActivationType::SiTu) + { + base_activation_type = ActivationType::SwigluBias; + } } if (swiglu_limit.has_value()) { CHECK_INPUT(swiglu_limit.value(), at::ScalarType::Float); TORCH_CHECK(swiglu_limit.value().sizes()[0] == num_experts_on_rank, "swiglu_limit must have num_experts_on_rank elements."); - base_activation_type = ActivationType::SwigluBias; + if (base_activation_type != ActivationType::SiTu) + { + base_activation_type = ActivationType::SwigluBias; + } } auto activation_params = ActivationParams(base_activation_type, reinterpret_cast(swiglu_alpha.has_value() ? swiglu_alpha.value().const_data_ptr() : nullptr), diff --git a/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16.yaml b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16.yaml new file mode 100644 index 000000000000..71b07d93dd5b --- /dev/null +++ b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16.yaml @@ -0,0 +1,48 @@ +# Kimi K3 NVFP4 on DEP16 (4 nodes x 4 GPU). +# Derived from the DEP8 variant; kept separate from the shared MXFP4 DEP16 +# template so that one stays untouched. +# +# DEP8 does not fit. The parameters alone exceed a 288 GB GB300 at EP8: the +# per-rank footprint is NOT checkpoint_bytes / 8, because enable_attention_dp +# replicates every non-routed weight on every rank, and Kimi K3 converts the +# checkpoint's FP8 attention to BF16 online (KIMI_K3_FP8_WEIGHT_READ is off by +# default), doubling it. Halving the routed-expert share is what buys the +# headroom back. See Phase 3 for making the attention weights read as FP8. +# +# max_batch_size is 8, not the template's 32. The V2 Mamba cache reserves a +# full recurrent-state slot per resident sequence (69 KDA layers, fp32 state, +# ~0.42 GiB/slot), so the manager's minimum live quota scales with it: at 32 it +# demands 14.39 GiB. NVFP4 cannot pay that here. Its in-memory weights are +# ~4.8 GiB LARGER per rank than MXFP4's (200.36 vs 195.58 GiB) even though its +# checkpoint is smaller -- NVFP4 carries an FP8 block scale per 16 elements +# where MXFP4 carries one UE8M0 per 32, i.e. 0.5625 vs 0.53125 bytes/element -- +# and the MXFP4 DEP16 recipe only had 1.06 GiB of slack in the second KV +# sizing pass. Lowering the batch lowers the minimum instead of the headroom. +# It costs eval wall time, not accuracy. +# +# backend: CUTLASS is required, not a preference -- AUTO resolves Kimi K3 to +# TRTLLM, and trtllm-gen ships SiTu cubins for W4A8_MXFP4_MXFP8 only. +# +# moe_config.max_num_tokens stays at the inherited value: for CUTLASS it is a +# per-call chunking bound. Do NOT carry it over to MEGAMOE_* backends, where +# it is the SymmBuffer capacity and this value over-provisions it 32x. +tensor_parallel_size: 16 +enable_attention_dp: true +moe_expert_parallel_size: 16 +max_batch_size: 8 +max_num_tokens: 8192 +max_seq_len: 8192 +trust_remote_code: true +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + enable_padding: true + max_batch_size: 8 +moe_config: + backend: CUTLASS + max_num_tokens: 131072 + use_low_precision_moe_combine: true +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.25 + tokens_per_block: 64 diff --git a/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_gpqa.yaml b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_gpqa.yaml new file mode 100644 index 000000000000..b0cf9f60715d --- /dev/null +++ b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_gpqa.yaml @@ -0,0 +1,37 @@ +# Kimi K3 NVFP4 on DEP16 for GPQA-Diamond (4 nodes x 4 GPU). +# +# Differs from the GSM8K DEP16 config in exactly one thing that matters: +# max_seq_len. GSM8K answers are a few hundred tokens, so 8192 covered prompt +# plus generation. GPQA-Diamond is a reasoning benchmark and the published +# Kimi-K3 numbers were measured with a 65536-token generation budget, so the +# sequence budget has to hold 4096 of prompt plus all of that. +# +# max_batch_size is 8, revised up from 4 after the 8-question smoke (job +# 475469) measured what this model actually generates here: 3431 and 6000 +# tokens, not the 65536 the budget allows. The budget still has to cover the +# worst case, but sizing CONCURRENCY for it was wrong -- at ~6k tokens and +# ~63 KiB/token a sequence wants ~0.4 GiB, not ~4 GiB, so the pool funds +# several. The V2 Mamba minimum also stays comfortable: ~0.42 GiB per resident +# slot means ~4.2 GiB at batch 8, well under the ~15 GiB quota, and the GSM8K +# DEP16 run already ran at 8. If several questions do run long the scheduler +# simply admits fewer of them. +tensor_parallel_size: 16 +enable_attention_dp: true +moe_expert_parallel_size: 16 +max_batch_size: 8 +max_num_tokens: 8192 +max_seq_len: 69632 +trust_remote_code: true +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + enable_padding: true + max_batch_size: 8 +moe_config: + backend: CUTLASS + max_num_tokens: 131072 + use_low_precision_moe_combine: true +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.35 + tokens_per_block: 64 diff --git a/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep8.yaml b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep8.yaml new file mode 100644 index 000000000000..9bfea16b0cb7 --- /dev/null +++ b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep8.yaml @@ -0,0 +1,30 @@ +# Kimi K3 NVFP4 on DEP8 (2 nodes x 4 GPU). +# Derived from eval_extra_llm_options.yaml; kept separate so the shared +# DEP16 template stays untouched. +# +# backend: CUTLASS is required, not a preference -- AUTO resolves Kimi K3 to +# TRTLLM, and trtllm-gen ships SiTu cubins for W4A8_MXFP4_MXFP8 only. +# +# moe_config.max_num_tokens stays at the inherited value: for CUTLASS it is a +# per-call chunking bound. Do NOT carry it over to MEGAMOE_* backends, where +# it is the SymmBuffer capacity and this value over-provisions it 32x. +tensor_parallel_size: 8 +enable_attention_dp: true +moe_expert_parallel_size: 8 +max_batch_size: 32 +max_num_tokens: 8192 +max_seq_len: 8192 +trust_remote_code: true +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + enable_padding: true + max_batch_size: 32 +moe_config: + backend: CUTLASS + max_num_tokens: 131072 + use_low_precision_moe_combine: true +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.25 + tokens_per_block: 64 diff --git a/examples/kimi_k3/run_eval_kimi_k3.sbatch b/examples/kimi_k3/run_eval_kimi_k3.sbatch index 1c269fbf93a9..7bc009781435 100644 --- a/examples/kimi_k3/run_eval_kimi_k3.sbatch +++ b/examples/kimi_k3/run_eval_kimi_k3.sbatch @@ -190,9 +190,10 @@ VENV=${TRTLLM_VENV:-$REPO/.venv-3.12} # footprint and KV budget: drafter weights, capture buffer, and DFlash # context-KV slots leave less headroom than the SA config assumes # (0.25/8192 OOM'd in warmup on GB300). -MAX_BATCH_SIZE=32 -MAX_NUM_TOKENS=8192 -KV_FRAC=0.25 +# These defaults remain overridable for NVFP4 topology-specific runs. +MAX_BATCH_SIZE=${MAX_BATCH_SIZE:-32} +MAX_NUM_TOKENS=${MAX_NUM_TOKENS:-8192} +KV_FRAC=${KV_FRAC:-0.25} KV_REUSE_ARG="--disable_kv_cache_reuse" case "$MODE" in sa) @@ -219,23 +220,29 @@ case "$MODE" in SPEC_STATS_DEFAULT=1 ;; *) - EVAL_CONFIG=$REPO/examples/kimi_k3/eval_extra_llm_options.yaml + # Honour an externally exported EVAL_CONFIG so alternative topologies + # (e.g. the NVFP4 DEP8 config) can reuse this script unchanged. + EVAL_CONFIG=${EVAL_CONFIG:-$REPO/examples/kimi_k3/eval_extra_llm_options.yaml} ;; esac # Per-task eval settings (see the header for the mmmu rationale). The task # subcommand string is expanded into the container command below, so keep it # on a single line. Per-sample MMMU outputs land next to the job log. -MAX_SEQ_LEN=8192 case "$TASK" in gsm8k) + MAX_SEQ_LEN=${MAX_SEQ_LEN:-8192} TASK_CMD="gsm8k" ;; mmmu) - MAX_SEQ_LEN=24576 + MAX_SEQ_LEN=${MAX_SEQ_LEN:-24576} TASK_CMD="mmmu --post_process_fn kimi_k3_mmmu --max_input_length 8192 --max_output_length 16384 --preserve_caller_max_tokens --output_path '$REPO/kimi-k3-mmmu-results-$SLURM_JOB_ID'" ;; esac +# Advanced evals such as GPQA can override the subcommand and its flags while +# reusing the same server launch and topology wrapper. +EVAL_TASK=${EVAL_TASK:-$TASK_CMD} +EVAL_TASK_ARGS=${EVAL_TASK_ARGS:-} # Per-parallel-layout adjustments (see the header). The TEP rewrite happens on # a per-job copy so the checked-in YAMLs stay canonical for DEP16. In the @@ -327,6 +334,14 @@ srun --mpi=pmix \ # sa/dflash modes, off otherwise (explicit env still wins either way). export TLLM_EVAL_SPEC_STATS=\"\${TLLM_EVAL_SPEC_STATS:-$SPEC_STATS_DEFAULT}\" + # Kimi K3 weight/precision switches. Forwarded explicitly rather than + # relying on environment propagation into the container, so that a run + # either has the setting it was submitted with or fails visibly - a + # silently-unset precision flag would look like a clean A/B while + # actually re-running the same configuration twice. + export KIMI_K3_FP8_WEIGHT_READ=\"${KIMI_K3_FP8_WEIGHT_READ:-0}\" + echo \"[k3] KIMI_K3_FP8_WEIGHT_READ=\$KIMI_K3_FP8_WEIGHT_READ\" + # Import tensorrt_llm from $REPO, not from wherever the venv's # in-place install points (for git-worktree submits those differ; # without this the job silently tests the main checkout's code). @@ -336,7 +351,7 @@ srun --mpi=pmix \ '$VENV/bin/trtllm-eval' \ --model \"$MODEL\" \ --backend pytorch \ - --tp_size 16 \ + --tp_size ${TP_SIZE:-16} \ --max_batch_size $MAX_BATCH_SIZE \ --max_seq_len $MAX_SEQ_LEN \ --max_num_tokens $MAX_NUM_TOKENS \ @@ -344,5 +359,5 @@ srun --mpi=pmix \ $KV_REUSE_ARG \ --trust_remote_code \ --config '$EVAL_CONFIG' \ - $TASK_CMD + $EVAL_TASK $EVAL_TASK_ARGS " diff --git a/tensorrt_llm/_torch/auto_deploy/_compat.py b/tensorrt_llm/_torch/auto_deploy/_compat.py index 6e79568df948..aa980387cf86 100644 --- a/tensorrt_llm/_torch/auto_deploy/_compat.py +++ b/tensorrt_llm/_torch/auto_deploy/_compat.py @@ -56,6 +56,7 @@ # --------------------------------------------------------------------------- if TRTLLM_AVAILABLE: from tensorrt_llm.quantization.modelopt_config import ( + canonicalize_quant_algo, is_modelopt_quant_config, read_modelopt_quant_config, ) @@ -76,6 +77,15 @@ def is_modelopt_quant_config(raw: Any) -> bool: } _KV_SCHEME_STRING_ALGOS = {"FP8", "NVFP4", "INT8"} + # Keep in sync with ``tensorrt_llm.quantization.modelopt_config``. + _QUANT_ALGO_ALIASES = {"fp8_pb_wo": "FP8_BLOCK_SCALES"} + + def canonicalize_quant_algo(value: Any) -> Any: + """Map a ModelOpt ``quant_algo`` spelling onto its ``QuantAlgo`` name.""" + if not isinstance(value, str): + return value + return _QUANT_ALGO_ALIASES.get(value.lower(), value) + def _kv_cache_scheme_to_algo(scheme: Any) -> Optional[str]: if scheme is None: return None @@ -112,8 +122,20 @@ def read_modelopt_quant_config(raw: Dict[str, Any]) -> Dict[str, Any]: f"Not a modelopt quant config (producer={raw.get('producer')!r}, " f"quant_method={raw.get('quant_method')!r})" ) - if result.get("quant_algo") == "fp8_pb_wo": - result["quant_algo"] = "FP8_BLOCK_SCALES" + # Canonicalize both the top-level algo and the per-layer entries: + # MIXED_PRECISION checkpoints carry the real names in quantized_layers. + if "quant_algo" in result: + result["quant_algo"] = canonicalize_quant_algo(result["quant_algo"]) + layers = result.get("quantized_layers") + if isinstance(layers, dict): + result["quantized_layers"] = { + name: ( + {**cfg, "quant_algo": canonicalize_quant_algo(cfg["quant_algo"])} + if isinstance(cfg, dict) and "quant_algo" in cfg + else cfg + ) + for name, cfg in layers.items() + } return result diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 57ef2259ebef..e246f9f88b6c 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -47,8 +47,8 @@ update_quant_config_from_compressed_tensors from tensorrt_llm.quantization.mode import QuantAlgo from tensorrt_llm.quantization.modelopt_config import ( - is_modelopt_quant_config, read_modelopt_quant_config, - warn_if_inline_diverges) + canonicalize_quant_algo, is_modelopt_quant_config, + read_modelopt_quant_config, warn_if_inline_diverges) if TYPE_CHECKING: from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp @@ -506,9 +506,12 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, quant_config = QuantConfig() layer_quant_config = None - quant_config.quant_algo = (QuantAlgo(json_quant_configs['quant_algo']) - if json_quant_configs.get('quant_algo') - is not None else None) + # ``canonicalize_quant_algo`` is applied again here (and per layer + # below) because the ``quant_cfg.json`` overlay merged in for + # MIXED_PRECISION bypasses ``read_modelopt_quant_config``. + quant_config.quant_algo = ( + QuantAlgo(canonicalize_quant_algo(json_quant_configs['quant_algo'])) + if json_quant_configs.get('quant_algo') is not None else None) quant_config.kv_cache_quant_algo = ( QuantAlgo(json_quant_configs['kv_cache_quant_algo']) if json_quant_configs.get('kv_cache_quant_algo') is not None else None) @@ -560,7 +563,8 @@ def _build_modelopt_quant_config(json_quant_configs, checkpoint_dir, layer_cfg = mixed_quant_configs[layer] config = QuantConfig() config.kv_cache_quant_algo = kv_cache_quant_algo - config.quant_algo = QuantAlgo(layer_cfg['quant_algo']) + config.quant_algo = QuantAlgo( + canonicalize_quant_algo(layer_cfg['quant_algo'])) config.group_size = layer_cfg.get('group_size', None) # AWQ-specific extras emitted by modelopt per-layer. if 'has_zero_point' in layer_cfg: diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 7f15e37fe0c7..6b138bf40695 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -85,8 +85,20 @@ import json import math import os +import threading from contextlib import ExitStack -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Set, Tuple +from typing import ( + TYPE_CHECKING, + Any, + Callable, + Dict, + List, + Literal, + NamedTuple, + Optional, + Set, + Tuple, +) import torch from safetensors import safe_open @@ -106,7 +118,7 @@ from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm from ..modules.situ import SituAndMul -from ..utils import ActType_TrtllmGen +from ..utils import ActivationType, ActType_TrtllmGen from .modeling_speculative import SpecDecOneEngineForCausalLM from .modeling_utils import DecoderModel, register_auto_model, run_concurrently @@ -500,12 +512,155 @@ def quantize_weight(weight: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: ) return weight_fp8, weight_scale + @staticmethod + def prepare_checkpoint_scale( + weight_fp8: torch.Tensor, weight_scale: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Checkpoint FP8 + FP32 128x128 block scale -> deep_gemm-ready pair. + + This is ``quantize_weight`` without its first step. A checkpoint that + already stores FP8_PB_WO weights (``nvidia/Kimi-K3-NVFP4``) supplies + exactly what ``per_block_cast_to_fp8`` would have produced, so the only + work left is the scale bridge: resmooth to UE8M0 and pack into + deep_gemm's TMA-aligned MN-major layout, because ``fp8_swap_ab_gemm`` + runs with ``disable_ue8m0_cast=True`` and would misread a plain FP32 + block scale. + + Taking this route instead of dequantizing to BF16 and re-quantizing + avoids a lossy round trip: the re-quantization recomputes each block's + scale from the dequantized values, which need not reproduce the + scale the checkpoint was written with. + """ + from ...quantization.utils.fp8_utils import ( + resmooth_to_fp8_e8m0, + transform_sf_into_required_layout, + ) + + weight_fp8, weight_scale = resmooth_to_fp8_e8m0( + weight_fp8.contiguous(), weight_scale.contiguous().float() + ) + weight_scale = transform_sf_into_required_layout( + weight_scale, + mn=weight_fp8.shape[0], + k=weight_fp8.shape[1], + recipe=(1, 128, 128), + is_sfa=False, + ) + return weight_fp8, weight_scale + @classmethod def from_linear(cls, linear: nn.Linear | TrtllmLinear) -> "_Fp8BlockScaleWeightReadLinear": assert linear.bias is None, "FP8 weight read expects a bias-free Linear" + # If the checkpoint stored this projection as FP8_PB_WO, the loader + # kept the original pair on the parameter. Using it skips a lossy + # round trip: dequantizing to BF16 and re-quantizing recomputes each + # block's scale from the dequantized values, which need not reproduce + # the scale the checkpoint was written with. + ckpt_pair = getattr(linear.weight, _K3_CKPT_FP8_ATTR, None) + if ckpt_pair is not None: + return cls.from_checkpoint_fp8(ckpt_pair[0], ckpt_pair[1], linear.out_features) weight_fp8, weight_scale = cls.quantize_weight(linear.weight.data) return cls(weight_fp8, weight_scale, linear.out_features) + @classmethod + def empty_placeholder( + cls, out_features: int, in_features: int + ) -> "_Fp8BlockScaleWeightReadLinear": + """A module that occupies (almost) nothing until the loader fills it. + + This is the contract construction-time FP8 dispatch has to satisfy: + build one of these instead of an ``nn.Linear`` and the BF16 weight is + never allocated at all. That matters because the DEP8 failure was an + OOM during module CONSTRUCTION, before a single weight had been read, + so any scheme that allocates BF16 and converts afterwards cannot fix + it however cheap the steady state ends up. + + Same shape of trick MegaMoE uses for its raw NVFP4 params -- and with + the same hazard, learned there: anything that indexes these before the + loader has filled them sees an empty tensor. ``load_checkpoint_pair`` + is the only way to populate them, and ``forward`` refuses to run while + they are still empty rather than producing garbage. + """ + # Shaped [0, in_features] rather than a flat empty: __init__ reads + # in_features off the weight, and numel() == 0 still marks it unfilled. + return cls( + torch.empty(0, in_features, dtype=torch.float8_e4m3fn), + torch.empty(0, 0, dtype=torch.int32), + out_features, + ) + + @property + def is_placeholder(self) -> bool: + return self.weight.numel() == 0 + + def load_checkpoint_pair(self, pairs) -> None: + """Fill a placeholder from one or more checkpoint FP8_PB_WO pairs. + + Several pairs fuse along ``out`` (KDA's q/k/v/g); one is the plain + case. Rejects a second fill so a double-load is loud rather than + silently keeping whichever arrived last. + """ + if not self.is_placeholder: + raise RuntimeError( + "Kimi K3 FP8 placeholder was filled twice; the second load would silently win." + ) + pairs = list(pairs) + filled = ( + self.fuse_checkpoint_fp8(pairs) + if len(pairs) > 1 + else self.from_checkpoint_fp8(pairs[0][0], pairs[0][1], self.out_features) + ) + if filled.out_features != self.out_features: + raise ValueError( + f"Kimi K3 FP8 placeholder expects out_features=" + f"{self.out_features}, checkpoint gives {filled.out_features}." + ) + self.register_buffer("weight", filled.weight, persistent=False) + self.register_buffer("weight_scale", filled.weight_scale, persistent=False) + self.in_features = filled.in_features + + @classmethod + def fuse_checkpoint_fp8(cls, pairs) -> "_Fp8BlockScaleWeightReadLinear": + """Fuse per-projection checkpoint FP8_PB_WO pairs along ``out``. + + ``pairs`` is an ordered sequence of ``(weight_fp8, fp32 block scale)`` + sharing one input dim -- e.g. KDA's q/k/v/g. + + This exists so a fused projection can be built WITHOUT ever + materializing the BF16 concatenation the current path goes through, + which is what makes constructing attention as FP8 possible at all. + It is sound for the reason ``quantize_weight`` documents: with every + ``out`` a multiple of 128 no 128x128 block crosses a boundary, so the + fused quantization is per-block identical to the individual ones -- + hence the fused FP8 rows ARE the individual FP8 rows, and likewise + for the block-scale rows. + + The layout transform is applied ONCE to the assembled scale rather + than per part: ``transform_sf_into_required_layout`` packs against the + full ``mn``, so transforming the pieces and stacking afterwards would + not produce the same bytes. + """ + weights = [w for w, _ in pairs] + scales = [s for _, s in pairs] + if any(w.shape[0] % _FP8_BLOCK for w in weights): + raise ValueError( + "Kimi K3 fused FP8 read requires every out dim to be a " + f"multiple of {_FP8_BLOCK}; got " + f"{[tuple(w.shape) for w in weights]}." + ) + fused_weight = torch.cat(weights, dim=0) + fused_scale = torch.cat(scales, dim=0) + prepared, scale = cls.prepare_checkpoint_scale(fused_weight, fused_scale) + return cls(prepared, scale, fused_weight.shape[0]) + + @classmethod + def from_checkpoint_fp8( + cls, weight_fp8: torch.Tensor, weight_scale: torch.Tensor, out_features: int + ) -> "_Fp8BlockScaleWeightReadLinear": + """Build directly from a checkpoint's FP8_PB_WO pair, no BF16 detour.""" + prepared, scale = cls.prepare_checkpoint_scale(weight_fp8, weight_scale) + return cls(prepared, scale, out_features) + def forward( self, x: torch.Tensor, @@ -514,6 +669,15 @@ def forward( lora_params: Optional[dict] = None, layer_idx: Optional[int] = None, ) -> torch.Tensor: + if self.is_placeholder: + # An unfilled placeholder would otherwise reach fp8_swap_ab_gemm + # with empty operands. Fail here instead: this whole path has a + # habit of turning wrong weights into plausible-looking numbers + # rather than errors. + raise RuntimeError( + "Kimi K3 FP8 placeholder was never filled by the loader; " + "load_checkpoint_pair must run before forward." + ) if lora_params: raise NotImplementedError("Kimi K3 FP8 weight read does not support LoRA.") out_shape = x.shape[:-1] + (self.out_features,) @@ -715,6 +879,97 @@ def _convert_mla_projections_to_fp8_weight_read(model: nn.Module) -> int: # Latent MoE block using the unified ConfigurableMoE stack. # --------------------------------------------------------------------------- +# Routed-expert key spellings that ModelOpt emits for Kimi K3. The NVFP4 +# checkpoint (``nvidia/Kimi-K3-NVFP4``) lists every prefix x module-name +# combination in ``quantized_layers``, so a lookup over this product finds it +# without needing the MiniMax-M3-style prefix normalization in ``ModelConfig``. +_K3_ROUTED_EXPERT_KEY_PREFIXES = ("language_model.model.", "model.", "") +_K3_ROUTED_EXPERT_KEY_SUFFIXES = ("block_sparse_moe.experts", "mlp.experts") + +# Routed-expert quantization used when the checkpoint declares nothing per +# layer. The original ``moonshotai/Kimi-K3`` ships a compressed-tensors +# ``mxfp4-pack-quantized`` config with no ModelOpt per-layer entries, and that +# checkpoint is what this default has always served. +_K3_DEFAULT_ROUTED_QUANT_ALGO = QuantAlgo.W4A8_MXFP4_MXFP8 + + +# --------------------------------------------------------------------------- +# Routed-expert checkpoint layouts. +# +# K3 streams routed experts one at a time (see ``load_weights``), so the tensor +# names, the per-expert loader and the finalization it needs are all decided by +# the checkpoint's routed-expert quantization rather than by the MoE backend. +# Three call sites need that decision — the expected-key plan, the loader, and +# the file grouping — so it lives in one spec instead of three conditionals. +# --------------------------------------------------------------------------- + + +def _load_packed_mxfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: + backend.quant_method.load_packed_mxfp4_expert( + backend, + global_expert_id=expert_idx, + local_slot_id=local_slot_id, + w1_weight=get_tensor(f"{base}.{expert_idx}.w1.weight_packed"), + w1_weight_scale=get_tensor(f"{base}.{expert_idx}.w1.weight_scale"), + w2_weight=get_tensor(f"{base}.{expert_idx}.w2.weight_packed"), + w2_weight_scale=get_tensor(f"{base}.{expert_idx}.w2.weight_scale"), + w3_weight=get_tensor(f"{base}.{expert_idx}.w3.weight_packed"), + w3_weight_scale=get_tensor(f"{base}.{expert_idx}.w3.weight_scale"), + ) + + +def _load_nvfp4_expert(backend, base, expert_idx, local_slot_id, get_tensor) -> None: + backend.quant_method.load_streaming_nvfp4_expert( + backend, + global_expert_id=expert_idx, + local_slot_id=local_slot_id, + **{ + f"{w}_{kind}": get_tensor(f"{base}.{expert_idx}.{w}.{kind}") + for w in ("w1", "w2", "w3") + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") + }, + ) + + +class _K3ExpertCkptSpec(NamedTuple): + """How one routed-expert quantization is spelled and loaded.""" + + # Per-``w{1,2,3}`` checkpoint tensor suffixes this layout stores. + kinds: Tuple[str, ...] + loader: Callable[..., None] + # Set of filled slots the loader maintains, checked after the load. + loaded_slots_attr: str + # NVFP4 defers cat/pad/interleave and the alpha computation to + # ``process_weights_after_loading``; the MXFP4 loaders write through. + needs_layer_finalize: bool + + +_K3_EXPERT_CKPT_SPECS = { + QuantAlgo.W4A8_MXFP4_MXFP8: _K3ExpertCkptSpec( + kinds=("weight_packed", "weight_scale"), + loader=_load_packed_mxfp4_expert, + loaded_slots_attr="_packed_mxfp4_loaded_slots", + needs_layer_finalize=False, + ), + QuantAlgo.NVFP4: _K3ExpertCkptSpec( + kinds=("weight", "weight_scale", "weight_scale_2", "input_scale"), + loader=_load_nvfp4_expert, + loaded_slots_attr="_streamed_expert_slots", + needs_layer_finalize=True, + ), +} + + +def _k3_expert_ckpt_spec(quant_algo: Optional[QuantAlgo]) -> _K3ExpertCkptSpec: + spec = _K3_EXPERT_CKPT_SPECS.get(quant_algo) + if spec is None: + raise NotImplementedError( + f"Kimi K3 routed experts are quantized as {quant_algo}, for which " + "no per-expert checkpoint layout is known. Supported: " + f"{sorted(a.name for a in _K3_EXPERT_CKPT_SPECS)}." + ) + return spec + class KimiK3MoERuntime(nn.Module): """Kimi K3 latent MoE block backed by ConfigurableMoE.""" @@ -760,7 +1015,10 @@ def __init__( self.gate = KimiK3MoEGate(cfg, logits_gemm_dtype=torch.bfloat16 if _router_bf16 else None) routed_moe_model_config = self._routed_moe_model_config(model_config) - routed_quant_config = QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8) + routed_quant_config = self._resolve_routed_quant_config(model_config, layer_idx) + # Resolved here so ``load_weights`` reads the checkpoint layout off the + # module instead of re-deriving it at each of its three call sites. + self.expert_ckpt_spec = _k3_expert_ckpt_spec(routed_quant_config.quant_algo) routed_moe_kwargs = dict( routing_method=self.gate.routing_method, num_experts=self.num_experts, @@ -776,7 +1034,44 @@ def __init__( # Let CommunicationFactory select the best available strategy. communication_method=None, ) - if routed_moe_model_config.moe_backend == "TRTLLM": + # trtllm-gen ships SiTu cubins for exactly one dtype combination + # (``Bmm_MxE4m3_MxE2m1MxE4m3`` = MXFP8 act x MXFP4 weight) and has no + # standalone SiTu activation kernel to fall back on, so a non-MXFP4 + # routed-expert format cannot be served here. Fail with the fix rather + # than with a cubin lookup error deep inside the runner. Checked against + # the resolved backend, not the K3 architecture branch, because the + # generic FP8_BLOCK_SCALES fallback in ``resolve_moe_backend`` can also + # land on TRTLLM. + if ( + routed_moe_model_config.moe_backend == "TRTLLM" + and routed_quant_config.quant_algo != QuantAlgo.W4A8_MXFP4_MXFP8 + ): + raise ValueError( + f"Kimi K3 routed experts are quantized as " + f"{routed_quant_config.quant_algo}, which the TRTLLM " + "(trtllm-gen) MoE backend cannot serve: its SiTu activation " + "exists only for W4A8_MXFP4_MXFP8. Set moe_config.backend to " + "CUTLASS or MEGAMOE_CUTEDSL." + ) + + if routed_moe_model_config.moe_backend == "CUTLASS": + local_num_experts = self.num_experts // routed_moe_model_config.mapping.moe_ep_size + device = torch.device("cuda", torch.cuda.current_device()) + self.routed_situ_alpha = torch.full( + (local_num_experts,), float(situ_beta), dtype=torch.float32, device=device + ) + self.routed_situ_beta = torch.full( + (local_num_experts,), + float(situ_linear_beta if situ_linear_beta is not None else 1.0), + dtype=torch.float32, + device=device, + ) + routed_moe_kwargs.update( + activation_type=ActivationType.SiTu, + swiglu_alpha=self.routed_situ_alpha, + swiglu_beta=self.routed_situ_beta, + ) + elif 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 @@ -913,6 +1208,40 @@ def _select_moe_tp_ep(mapping: Mapping) -> Tuple[int, int]: return mapping.moe_tp_size, mapping.moe_ep_size return 1, tp_size + @staticmethod + def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> QuantConfig: + """Routed-expert quantization for ``layer_idx``, taken from the checkpoint. + + ``nvidia/Kimi-K3-NVFP4`` declares the routed experts per layer as + ``NVFP4`` with ``group_size=16``; the original ``moonshotai/Kimi-K3`` + declares nothing per layer and keeps the historical + ``W4A8_MXFP4_MXFP8`` default. Reading the checkpoint instead of + hardcoding is what lets one code path serve both. + """ + per_layer = getattr(model_config, "quant_config_dict", None) + if per_layer: + for prefix in _K3_ROUTED_EXPERT_KEY_PREFIXES: + for suffix in _K3_ROUTED_EXPERT_KEY_SUFFIXES: + cfg = per_layer.get(f"{prefix}layers.{layer_idx}.{suffix}") + if cfg is not None and cfg.quant_algo is not None: + # Logged once per layer: the routed-expert format decides + # which MoE backends can serve this checkpoint at all. + logger.debug( + "Kimi K3 layer %d routed experts: %s (group_size=%s) " + "from the checkpoint", + layer_idx, + cfg.quant_algo, + cfg.group_size, + ) + return cfg + logger.debug( + "Kimi K3 layer %d routed experts: no per-layer quant config in the " + "checkpoint, defaulting to %s", + layer_idx, + _K3_DEFAULT_ROUTED_QUANT_ALGO, + ) + return QuantConfig(quant_algo=_K3_DEFAULT_ROUTED_QUANT_ALGO) + @staticmethod def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: """Build a private routed-expert mapping without mutating the shared @@ -2330,10 +2659,86 @@ def forward( # --------------------------------------------------------------------------- +_FP8_BLOCK_SCALE_SUFFIX = "_scale" +_FP8_BLOCK = 128 +# Set on a Parameter by the loader when the checkpoint stored that projection +# as FP8_PB_WO, so the FP8 weight-read conversion can reuse the original pair +# instead of re-deriving it from the dequantized BF16. +_K3_CKPT_FP8_ATTR = "_k3_ckpt_fp8_pair" + + +def _fp8_block_scale_key(weight_key: str) -> str: + return weight_key + _FP8_BLOCK_SCALE_SUFFIX + + +def _checkpoint_fp8_pair(ckpt_key: str, src: torch.Tensor, weights): + """``(fp8 weight, fp32 128x128 block scale)`` if this tensor is FP8_PB_WO. + + ``None`` for an ordinary (BF16) checkpoint tensor. Raises when an FP8 + tensor has no companion scale rather than letting the caller fall through + to a cast that would silently drop it. + """ + if src.dtype != torch.float8_e4m3fn: + return None + scale_key = _fp8_block_scale_key(ckpt_key) + if scale_key not in weights: + raise KeyError( + f"Kimi K3: {ckpt_key} is FP8 E4M3 but has no {scale_key}; refusing " + "to load it as if it were unquantized." + ) + scale = _materialize(weights[scale_key]).float() + # The checkpoint stores the block scale 4-D as + # [ceil(N/128), 1, ceil(K/128), 1]. Normalize here, once, so every + # consumer sees the plain 2-D [n_blocks_m, n_blocks_k] that both the + # dequantization below and deep_gemm's transform_sf_into_required_layout + # expect -- the latter asserts on rank and gave an unhelpful + # ``assert sf.dim() == ...`` when handed the raw 4-D tensor. + if scale.dim() == 4: + scale = scale.reshape(scale.shape[0], scale.shape[2]) + return src, scale + + +def _dequantize_fp8_block_scaled(ckpt_key: str, src: torch.Tensor, weights) -> torch.Tensor: + """Undo FP8_PB_WO block quantization, returning a BF16-compatible tensor. + + ``nvidia/Kimi-K3-NVFP4`` stores the attention projections as FP8 E4M3 + weights plus a per-128x128-block FP32 scale, where ``moonshotai/Kimi-K3`` + stores plain BF16. **The two are the same shape**, so nothing downstream + can tell them apart: the shape check passes and ``src.to(param.dtype)`` + happily reinterprets quantized values as if they were the real ones, with + the scale left behind as an unreferenced checkpoint key. The model then + loads clean and generates nonsense. + + Dequantizing here keeps the runtime on its existing BF16 attention path + (the same tensors the MXFP4 checkpoint would have supplied) instead of + requiring the FP8 weight-read path to be finished first. + """ + pair = _checkpoint_fp8_pair(ckpt_key, src, weights) + if pair is None: + return src + src, scale = pair + out = src.to(torch.float32) + if scale.numel() == 1: + return (out * scale.reshape(())).to(torch.bfloat16) + expanded = scale.repeat_interleave(_FP8_BLOCK, dim=0).repeat_interleave(_FP8_BLOCK, dim=1) + if expanded.shape[0] < out.shape[0] or expanded.shape[1] < out.shape[1]: + raise ValueError( + f"Kimi K3: {_fp8_block_scale_key(ckpt_key)} covers " + f"{tuple(expanded.shape)} but {ckpt_key} is {tuple(out.shape)}." + ) + return (out * expanded[: out.shape[0], : out.shape[1]]).to(torch.bfloat16) + + def _materialize(value) -> torch.Tensor: """Materialize a (possibly lazy safetensors slice) weight value.""" if isinstance(value, torch.Tensor): return value + # ``[:]`` is how a lazy slice is realized, but it is invalid on a 0-dim + # entry — and the NVFP4 checkpoint stores weight_scale_2 / input_scale as + # scalars. Same fix as ``_ReopenSafeTensorSlice._tensor``. + get_shape = getattr(value, "get_shape", None) + if get_shape is not None and len(get_shape()) == 0: + return value[()] return value[:] @@ -2476,10 +2881,11 @@ def checkpoint_name_plan( continue moe = layer.block_sparse_moe base = f"{prefix}model.layers.{layer_idx}.block_sparse_moe.experts" + kinds = moe.expert_ckpt_spec.kinds for expert_idx in moe.local_expert_ids: for w in ("w1", "w2", "w3"): - expected_keys.add(f"{base}.{expert_idx}.{w}.weight_packed") - expected_keys.add(f"{base}.{expert_idx}.{w}.weight_scale") + for kind in kinds: + expected_keys.add(f"{base}.{expert_idx}.{w}.{kind}") expert_jobs.append((layer_idx, moe, base)) return name_map, expected_keys, expert_jobs @@ -2558,6 +2964,9 @@ def _load_trunk_params( # divide model TP uses a smaller repeated TP subgroup, so its local # shard rank is model tp_rank modulo the parameter's shard count. model_tp_rank = self.model_config.mapping.tp_rank + # Keep each FP8_PB_WO checkpoint pair alongside the BF16 + # parameter only when the later weight-read conversion consumes it. + stash_ckpt_fp8 = os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_ENV, "0") != "0" and is_sm_100f() # KDA head-shard (attention-DP off): rank r loads head rows/cols # [r*local : (r+1)*local] of every head-major KDA tensor. kda_tp_size, kda_tp_rank = 1, 0 @@ -2574,8 +2983,12 @@ def load_param(name: str, param: torch.nn.Parameter): # Row-concat the checkpoint's separate gate_proj / up_proj # tensors into the fused [gate | up] parameter. gate_key, up_key = _gate_up_ckpt_keys(name_map[name]) - gate = _materialize(weights[gate_key]) - up = _materialize(weights[up_key]) + # This branch materializes its own sources, so it needs the + # same FP8 handling as the single-tensor path below. + gate = _dequantize_fp8_block_scaled( + gate_key, _materialize(weights[gate_key]), weights + ) + up = _dequantize_fp8_block_scaled(up_key, _materialize(weights[up_key]), weights) inter = param.shape[0] // 2 if gate.shape[0] != inter and gate.shape[0] % inter == 0: # TP-sharded fused MLP (shared experts on the direct @@ -2596,6 +3009,10 @@ def load_param(name: str, param: torch.nn.Parameter): param.data[inter:].copy_(up.to(param.dtype)) return src = _materialize(weights[name_map[name]]) + ckpt_fp8_pair = ( + _checkpoint_fp8_pair(name_map[name], src, weights) if stash_ckpt_fp8 else None + ) + src = _dequantize_fp8_block_scaled(name_map[name], src, weights) if name == "lm_head.weight": # LMHead is vocab-sharded (TP column) + gathered; its # load_weights shards the full checkpoint tensor. @@ -2708,6 +3125,13 @@ def load_param(name: str, param: torch.nn.Parameter): f"{tuple(param.shape)}" ) param.data.copy_(src.to(param.dtype)) + # Keep the checkpoint's FP8 pair for the weight-read conversion, + # but ONLY on this path -- the branches above shard or pad, and the + # conversion consumes linear.weight in its module-local shape, so a + # full-size pair would not match. Under enable_attention_dp the + # attention projections are replicated, which is exactly this path. + if ckpt_fp8_pair is not None and ckpt_fp8_pair[0].shape == param.shape: + setattr(param, _K3_CKPT_FP8_ATTR, ckpt_fp8_pair) param_jobs = [(name, params[name]) for name in name_map] run_concurrently(load_param, param_jobs, num_workers=8) @@ -2726,23 +3150,81 @@ def _load_expert_slices( the backend expert slots, then verify every slot was filled.""" device = next(self.parameters()).device + # Layouts whose per-expert loader only stages its input must have the + # staging containers created before any thread runs, and must be + # finalized once the layer's last slot lands: nothing else on this path + # calls process_weights_after_loading. Which thread completes a layer is + # a race, so the completion test and the "already finalized" bookkeeping + # are one critical section. + finalize_lock = threading.Lock() + finalized_backends = set() + prepared_backends = set() + + def ensure_prepared(moe: KimiK3MoERuntime): + """Prepare a layer's streaming state on its FIRST expert, not up front. + + Preparing every layer before the load starts is what OOM-ed the + MegaMoE CuteDSL backend: it keeps its raw NVFP4 source params as + 0-element placeholders and rematerializes them at full shape here, + so preparing all 92 layers held 92 layers of raw weights at once + instead of the handful actually being filled. Cutlass did not care + because its parameters are allocated either way. + + Bounded lazily instead. The paired shrink already happens per layer + in process_weights_after_loading, so the live set is whatever is + genuinely in flight -- measured at 1 layer per shard file for this + checkpoint (its rank-local experts are 1:1 with files), so ~4 with + 4 loader threads. Correctness does not depend on that layout + though; a checkpoint that split a layer across files would only + raise the peak, not break this. + + The membership add happens AFTER preparing, so the lock-free fast + path can only ever be stale in the safe direction. + """ + spec = moe.expert_ckpt_spec + if not spec.needs_layer_finalize: + return + backend = moe.routed_experts.backend + if id(backend) in prepared_backends: + return + with finalize_lock: + if id(backend) in prepared_backends: + return + backend.quant_method.prepare_streaming_expert_load(backend) + prepared_backends.add(id(backend)) + + def maybe_finalize_layer(moe: KimiK3MoERuntime): + spec = moe.expert_ckpt_spec + if not spec.needs_layer_finalize: + return + backend = moe.routed_experts.backend + with finalize_lock: + loaded = len(getattr(backend, spec.loaded_slots_attr, ())) + if loaded != backend.expert_size_per_partition: + return + if id(backend) in finalized_backends: + return + finalized_backends.add(id(backend)) + # Computes the alphas, interleaves the w2 scales, and (MegaMoE) + # packs into the mega buffers and shrinks the raw source params + # back to placeholders. Paired with ensure_prepared above, that + # pairing is what bounds the per-layer footprint. The per-expert + # drain inside load_streaming_nvfp4_expert separately bounds the + # Cutlass w3_w1 staging, which is per expert rather than per layer. + backend.process_weights_after_loading() + def load_expert( moe: KimiK3MoERuntime, base: str, local_slot_id: int, expert_idx: int, get_tensor ): if device.type == "cuda": torch.cuda.set_device(device) backend = moe.routed_experts.backend - backend.quant_method.load_packed_mxfp4_expert( - backend, - global_expert_id=expert_idx, - local_slot_id=local_slot_id, - w1_weight=get_tensor(f"{base}.{expert_idx}.w1.weight_packed"), - w1_weight_scale=get_tensor(f"{base}.{expert_idx}.w1.weight_scale"), - w2_weight=get_tensor(f"{base}.{expert_idx}.w2.weight_packed"), - w2_weight_scale=get_tensor(f"{base}.{expert_idx}.w2.weight_scale"), - w3_weight=get_tensor(f"{base}.{expert_idx}.w3.weight_packed"), - w3_weight_scale=get_tensor(f"{base}.{expert_idx}.w3.weight_scale"), - ) + # Every route into an expert goes through here (file-grouped, + # split-file and the shared-dict fallback alike), so this is the + # one place preparation has to be hooked. + ensure_prepared(moe) + moe.expert_ckpt_spec.loader(backend, base, expert_idx, local_slot_id, get_tensor) + maybe_finalize_layer(moe) def load_experts_from_weights(layer_idx: int, moe: KimiK3MoERuntime, base: str): del layer_idx @@ -2764,8 +3246,19 @@ def load_experts_from_weights(layer_idx: int, moe: KimiK3MoERuntime, base: str): # trays). Instead, group the rank-local expert tensors by shard file # and stream each file through a short-lived handle: # open -> copy -> close (unmap) -> fadvise(DONTNEED). - ckpt_dir = getattr(self.model_config.pretrained_config, "_name_or_path", None) + # The lazy loader records the directory it opened; prefer it over + # ``_name_or_path``, which transformers no longer populates (it is + # empty on transformers 5.x, which silently sent the whole load down + # the fallback below and OOM-killed the step). + ckpt_dir = getattr(weights, "checkpoint_dir", None) or getattr( + self.model_config.pretrained_config, "_name_or_path", None + ) index_path = os.path.join(ckpt_dir or "", "model.safetensors.index.json") + checkpoint_prefix = getattr(weights, "checkpoint_prefix", "") + + def checkpoint_key(key: str) -> str: + return f"{checkpoint_prefix}{key}" + if expert_jobs and ckpt_dir and os.path.isfile(index_path): with open(index_path) as f: weight_map = json.load(f)["weight_map"] @@ -2777,9 +3270,9 @@ def load_experts_from_weights(layer_idx: int, moe: KimiK3MoERuntime, base: str): keys = [ f"{base}.{expert_idx}.{w}.{kind}" for w in ("w1", "w2", "w3") - for kind in ("weight_packed", "weight_scale") + for kind in moe.expert_ckpt_spec.kinds ] - files = {weight_map[key] for key in keys} + files = {weight_map[checkpoint_key(key)] for key in keys} job = (moe, base, local_slot_id, expert_idx) if len(files) == 1: per_file.setdefault(files.pop(), []).append(job) @@ -2803,7 +3296,13 @@ def load_expert_file(file_name: str, jobs: list): path = os.path.join(ckpt_dir, file_name) with safe_open(path, framework="pt", device="cpu") as fh: for moe, base, local_slot_id, expert_idx in jobs: - load_expert(moe, base, local_slot_id, expert_idx, fh.get_tensor) + load_expert( + moe, + base, + local_slot_id, + expert_idx, + lambda key: fh.get_tensor(checkpoint_key(key)), + ) # Handle closed -> pages unmapped -> the drop takes effect. drop_file_pages(file_name) @@ -2821,7 +3320,8 @@ def load_split_file_expert(job, files): } def get_tensor(key): - return handles[weight_map[key]].get_tensor(key) + source_key = checkpoint_key(key) + return handles[weight_map[source_key]].get_tensor(source_key) load_expert(*job, get_tensor) for file_name in files: @@ -2830,18 +3330,42 @@ def get_tensor(key): run_concurrently(load_expert_file, sorted(per_file.items()), num_workers=4) run_concurrently(load_split_file_expert, split_file_jobs, num_workers=4) else: + # Falling back is a silent loss of the whole point of the block + # above: the shared lazy dict keeps every shard mapped, which is + # the OOM this streaming path exists to avoid. Say so. + if expert_jobs: + logger.warning( + f"Kimi K3: no safetensors index at '{index_path}', so routed " + "experts are loaded from the shared lazy weight dict instead " + "of being streamed per shard file. Every shard stays mapped " + "for the whole load, which OOM-kills the step at DEP8 scale." + ) run_concurrently(load_experts_from_weights, expert_jobs, num_workers=4) for _, moe, _ in expert_jobs: + spec = moe.expert_ckpt_spec backend = moe.routed_experts.backend - loaded_slots = getattr(backend, "_packed_mxfp4_loaded_slots", set()) + loaded_slots = getattr(backend, spec.loaded_slots_attr, set()) expected_slots = set(range(backend.expert_size_per_partition)) if loaded_slots != expected_slots: missing_slots = sorted(expected_slots - loaded_slots) raise RuntimeError( - "Kimi K3 packed expert loading did not fill all backend " + "Kimi K3 streaming expert loading did not fill all backend " f"slots; missing {missing_slots[:10]}." ) + if ( + spec.needs_layer_finalize + and expected_slots + and id(backend) not in finalized_backends + ): + # Unreachable via load_expert (the last slot finalizes), so + # reaching it means the two bookkeeping paths disagree. Guarded + # on expected_slots because a layer that owns no local slots is + # never prepared and so is legitimately never finalized. + raise RuntimeError( + "Kimi K3 streaming expert loading filled every slot but " + "never finalized the layer." + ) backend._weights_transformed = False def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 90828ccd6f5d..15a9383c44a6 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -2345,6 +2345,141 @@ def _get_fc2_alpha_input_scale( del local_slot_id, expert_id return module.fc2_input_scale.data + def prepare_streaming_expert_load(self, module: torch.nn.Module) -> None: + """Pre-create every staging container ``load_streaming_nvfp4_expert`` + mutates, before any loader thread starts. + + The whole-checkpoint path creates these lazily with + ``if not hasattr(module, ...)`` / ``getattr(module, ..., {})``, which is + a read-modify-write and therefore loses entries when a model streams + experts from several threads. Creating them up front keeps the + per-expert mutations to ``dict.__setitem__`` / ``set.add``, which are + atomic. + """ + module.tmp_raw_input_scales = {} + module.tmp_weight_scale_2 = {} + module._streamed_expert_slots = set() + + def finalize_streamed_expert(self, module: torch.nn.Module, + local_slot_id: int) -> None: + """Resolve any per-expert staging left by ``load_streaming_nvfp4_expert``. + + No-op by default: a backend that writes straight through to its + destination has nothing to drain. Backends that stage (the Cutlass + child, and MegaMoE CuteDSL) override this so a streaming loader does + not accumulate a second copy of the routed-expert weights. + """ + + def load_streaming_nvfp4_expert( + self, + module: torch.nn.Module, + *, + global_expert_id: int, + local_slot_id: int, + w1_weight: torch.Tensor, + w1_weight_scale: torch.Tensor, + w1_weight_scale_2: torch.Tensor, + w1_input_scale: torch.Tensor, + w2_weight: torch.Tensor, + w2_weight_scale: torch.Tensor, + w2_weight_scale_2: torch.Tensor, + w2_input_scale: torch.Tensor, + w3_weight: torch.Tensor, + w3_weight_scale: torch.Tensor, + w3_weight_scale_2: torch.Tensor, + w3_input_scale: torch.Tensor, + ) -> None: + """Load one NVFP4 checkpoint expert into a local slot. + + Per-expert counterpart of the ``load_weights`` / + ``load_quant_scales`` pair, for models whose checkpoint is too large to + keep mapped while the whole layer loads. It calls exactly the same + primitives on exactly the same destinations and leaves the same staging + state behind, so ``process_weights_after_loading`` finalizes a streamed + layer and a whole-checkpoint layer identically. + + ``prepare_streaming_expert_load`` must have run first, and + ``module.process_weights_after_loading()`` must run once the layer's + last slot is in — nothing else calls it on this path. + + One deliberate difference from ``load_quant_scales``: it reduces the + activation ``input_scale`` over every expert in the checkpoint, whereas + a streaming caller only ever holds its own rank's experts, so the + reduction is over the rank-local slice. The two agree only while + ``input_scale`` is uniform across experts — which is what a + static-activation-scale checkpoint (e.g. ``nvidia/Kimi-K3-NVFP4``, all + 1.0) gives. A checkpoint with genuinely per-expert activation scales + would make expert-parallel ranks disagree, so such a checkpoint needs a + cross-rank reduction added here first. + """ + 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}.") + if not hasattr(module, "_streamed_expert_slots"): + raise RuntimeError( + "prepare_streaming_expert_load() must run before " + "load_streaming_nvfp4_expert().") + if local_slot_id in module._streamed_expert_slots: + raise ValueError( + f"NVFP4 local slot {local_slot_id} was loaded twice.") + + for name, value in (("w1_weight", w1_weight), ("w2_weight", w2_weight), + ("w3_weight", w3_weight)): + if value.dtype != torch.uint8: + raise TypeError(f"{name} must contain packed NVFP4 uint8 data, " + f"got {value.dtype}.") + + # Cutlass stages the two halves of w3_w1 per expert and needs the slot + # id as the staging key; other backends write straight through. + w3_w1_weight_kargs = {} + if "expert_idx" in inspect.getfullargspec( + self.load_expert_w3_w1_weight).args: + w3_w1_weight_kargs["expert_idx"] = local_slot_id + w3_w1_scale_kargs = {} + if "expert_idx" in inspect.getfullargspec( + self.load_expert_w3_w1_weight_scale_nvfp4).args: + w3_w1_scale_kargs["expert_idx"] = local_slot_id + + self.load_expert_w3_w1_weight(module, w1_weight, w3_weight, + module.w3_w1_weight.data[local_slot_id], + **w3_w1_weight_kargs) + self.load_expert_w2_weight(module, w2_weight, + module.w2_weight.data[local_slot_id]) + self.load_expert_w3_w1_weight_scale_nvfp4( + module, w1_weight_scale, w3_weight_scale, + module.w3_w1_weight_scale.data[local_slot_id], **w3_w1_scale_kargs) + self.load_expert_w2_weight_scale_nvfp4( + module, w2_weight_scale, module.w2_weight_scale.data[local_slot_id]) + + # ``_reconcile_and_compute_alphas`` keys weight_scale_2 by local slot + # and ``process_weights_after_loading`` keys the input scales by global + # expert id; both are only reduced, never indexed positionally. + module.tmp_weight_scale_2[local_slot_id] = { + 'w1': w1_weight_scale_2, + 'w3': w3_weight_scale_2, + 'w2': w2_weight_scale_2, + } + module.tmp_raw_input_scales[global_expert_id] = { + 'w1': w1_input_scale[...].reshape([]), + 'w3': w3_input_scale[...].reshape([]), + 'w2': w2_input_scale[...].reshape([]), + } + + # Backends that only stage their input here resolve this slot right + # away, so the staged footprint stays proportional to the experts in + # flight rather than growing with the load. + self.finalize_streamed_expert(module, local_slot_id) + + # MegaMoE backends assert this in transform_weights; it is normally set + # by the load_weights this path replaces. + module._weights_loaded = True + module._streamed_expert_slots.add(local_slot_id) + def load_fp4_weight_block_scales( self, module: torch.nn.Module, @@ -2772,6 +2907,14 @@ class NVFP4CutlassFusedMoEMethod(NVFP4FusedMoEMethod): NVFP4_ROW_ALIGNMENT = 128 NVFP4_COL_ALIGNMENT = 4 + def prepare_streaming_expert_load(self, module: torch.nn.Module) -> None: + # This backend defers the cat + pad + interleave of w3_w1 to + # process_weights_after_loading, so it stages both halves per expert in + # two more dicts that a streaming loader writes from several threads. + super().prepare_streaming_expert_load(module) + module.tmp_cutlass_w3_w1_weights = {} + module.tmp_cutlass_w3_w1_weight_scales = {} + def get_weights_shapes(self, module: torch.nn.Module, weight_vec_size: int, block_scales_vec_size: int): """Override the base method to get aligned weights shapes for Cutlass nvfp4 alignment.""" @@ -2996,30 +3139,83 @@ def _maybe_padding_shape(self, source_tensor, dst_tensor): "constant", 0).contiguous() return source_tensor + def _resolve_staged_w3_w1_weight(self, entry: Dict) -> None: + """Cat + pad one staged expert's w3_w1 halves into its destination.""" + w3 = entry.get('w3') + w1 = entry.get('w1') + dst = entry['dst'] + if w3 is not None and w1 is not None: + cat_weight = torch.cat([w3, w1], dim=0) + cat_weight = self._maybe_padding_shape(cat_weight, dst) + dst.copy_(cat_weight, non_blocking=True) + elif w1 is not None: + # Non-gated MoE (e.g. Relu2): dst holds only w1; w3 source is empty. + w1 = self._maybe_padding_shape(w1, dst) + dst.copy_(w1, non_blocking=True) + + def _resolve_staged_w3_w1_weight_scale(self, entry: Dict) -> None: + """Cat + pad + interleave one staged expert's w3_w1 block scales.""" + w3_scale = entry.get('w3') + w1_scale = entry.get('w1') + dst = entry['dst'] + if w3_scale is not None and w1_scale is not None: + cat_scale = torch.cat([w3_scale, w1_scale], dim=0) + cat_scale = self._maybe_padding_shape(cat_scale, dst) + dst.copy_(cat_scale) + self._interleave_w3_w1_weight_scale(dst) + elif w1_scale is not None: + # Non-gated MoE (e.g. Relu2): dst holds only w1; w3 source is empty. + w1_scale = self._maybe_padding_shape(w1_scale, dst) + dst.copy_(w1_scale) + self._interleave_w3_w1_weight_scale(dst) + + def finalize_streamed_expert(self, module: torch.nn.Module, + local_slot_id: int) -> None: + """Resolve just this slot's staged halves, as soon as it is loaded. + + Staging is per-expert, so it can be drained per-expert. Doing so bounds + the staged footprint to the few experts in flight instead of letting it + grow with the load: the halves are a second copy of the routed-expert + weights, and a streaming loader that groups its work by shard FILE + (Kimi K3 does) finishes a given layer only when the last of its slots + happens to land, which can be arbitrarily late. + + ``dict.pop`` is atomic, so concurrent loader threads draining different + slots need no lock. + """ + # The key must be built exactly as the staging site builds it. The two + # sites do not agree on the accessor -- the weight one uses + # untyped_storage(), the scale one the deprecated storage() -- so + # mirror each rather than assume they return the same address. + for attr, dst_base, resolve in ( + ('tmp_cutlass_w3_w1_weights', lambda: module.w3_w1_weight.data[ + local_slot_id].untyped_storage().data_ptr(), + self._resolve_staged_w3_w1_weight), + ('tmp_cutlass_w3_w1_weight_scales', + lambda: module.w3_w1_weight_scale.data[local_slot_id].storage( + ).data_ptr(), self._resolve_staged_w3_w1_weight_scale), + ): + staged = getattr(module, attr, None) + if not staged: + continue + entry = staged.pop((dst_base(), local_slot_id), None) + if entry is not None: + resolve(entry) + def process_weights_after_loading(self, module: torch.nn.Module): - # Finalize w3_w1 weights: cat + pad + # Finalize w3_w1 weights: cat + pad. Streamed loads drain these + # per-expert in finalize_streamed_expert, so this handles whatever is + # left -- everything, for a whole-checkpoint load; nothing, for a fully + # streamed one. if hasattr(module, 'tmp_cutlass_w3_w1_weights'): for entry in module.tmp_cutlass_w3_w1_weights.values(): - w3 = entry.get('w3') - w1 = entry.get('w1') - dst = entry['dst'] - if w3 is not None and w1 is not None: - cat_weight = torch.cat([w3, w1], dim=0) - cat_weight = self._maybe_padding_shape(cat_weight, dst) - dst.copy_(cat_weight, non_blocking=True) + self._resolve_staged_w3_w1_weight(entry) delattr(module, 'tmp_cutlass_w3_w1_weights') # Finalize w3_w1 weight scales: cat + pad + interleave if hasattr(module, 'tmp_cutlass_w3_w1_weight_scales'): for entry in module.tmp_cutlass_w3_w1_weight_scales.values(): - w3_scale = entry.get('w3') - w1_scale = entry.get('w1') - dst = entry['dst'] - if w3_scale is not None and w1_scale is not None: - cat_scale = torch.cat([w3_scale, w1_scale], dim=0) - cat_scale = self._maybe_padding_shape(cat_scale, dst) - dst.copy_(cat_scale) - self._interleave_w3_w1_weight_scale(dst) + self._resolve_staged_w3_w1_weight_scale(entry) delattr(module, 'tmp_cutlass_w3_w1_weight_scales') # Finalize w2 weight scales: interleave (regular experts) @@ -3584,6 +3780,17 @@ class NVFP4MegaMoECuteDslMethod(NVFP4FusedMoEMethod): weight_dtype = FUSED_MOE_NVFP4_WEIGHT_DTYPE block_scales_dtype = FUSED_MOE_NVFP4_WEIGHT_BLOCK_SCALE_DTYPE + def prepare_streaming_expert_load(self, module: torch.nn.Module) -> None: + # Like the Cutlass child this backend stages the w3_w1 halves, and it + # additionally tracks which w2 rows a partial load covered. Every one of + # those containers is created lazily on the load path, which a + # multi-threaded streaming loader cannot do safely. + super().prepare_streaming_expert_load(module) + module.tmp_cutlass_w3_w1_weights = {} + module.tmp_cutlass_w3_w1_weight_scales = {} + module._streamed_w2_covered = set() + module._streamed_w2_scale_covered = set() + def _get_fc2_alpha_input_scale( self, module: torch.nn.Module, diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 7cf9e33b77f0..e9b6444d1095 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -61,14 +61,13 @@ class ActivationType(IntEnum): Geglu = 6 SwigluBias = 7 Relu2 = 8 + SiTu = 9 # TRTLLM-Gen-local activation encoding, kept separate from the shared -# ActivationType above ON PURPOSE: ActivationType mirrors the cutlass enum in -# common.h and drives cutlass MoE kernels, whereas SiTu exists only in the -# trtllm-gen batched-GEMM kernels. Adding SiTu to the shared ActivationType -# would force a matching cutlass enum member that no cutlass kernel implements. -# So SiTu stays here (TRTLLM-15177 item 1.2(a): decided keep-backend-local). +# ActivationType above: ActivationType mirrors the CUTLASS enum in common.h, +# while ActType_TrtllmGen mirrors the independent batched-GEMM encoding below. +# SiTu is supported by both backends, but its numeric value is backend-local. # Keep this in sync with the ActType enum in # cpp/tensorrt_llm/kernels/trtllmGenKernels/batchedGemm/KernelRunner.h class ActType_TrtllmGen(IntEnum): diff --git a/tensorrt_llm/quantization/modelopt_config.py b/tensorrt_llm/quantization/modelopt_config.py index afd482f0ea95..407710b3b718 100644 --- a/tensorrt_llm/quantization/modelopt_config.py +++ b/tensorrt_llm/quantization/modelopt_config.py @@ -46,6 +46,51 @@ def is_modelopt_quant_config(raw: Any) -> bool: } _KV_SCHEME_STRING_ALGOS = {"FP8", "NVFP4", "INT8"} +# ModelOpt ships algo spellings that predate the ``QuantAlgo`` enum names. +# Lookup is case-insensitive: the same alias appears lowercase as a top-level +# ``quant_algo`` (modelopt 0.x) and uppercase inside ``quantized_layers`` +# (modelopt >= 0.45, e.g. ``nvidia/Kimi-K3-NVFP4``). +_QUANT_ALGO_ALIASES = { + # Per-block FP8, weight-only. Same on-disk layout as the DeepSeek block + # recipe (E4M3 weights + one FP32 scale per 128x128 block), so it maps onto + # FP8_BLOCK_SCALES; ``FP8BlockScalesLinearMethod.load_weights_vanilla`` + # squeezes ModelOpt's two extra singleton scale dimensions. + "fp8_pb_wo": "FP8_BLOCK_SCALES", +} + + +def canonicalize_quant_algo(value: Any) -> Any: + """Map a ModelOpt ``quant_algo`` spelling onto its ``QuantAlgo`` name. + + Case-insensitive. Unknown values pass through untouched so that ``QuantAlgo`` + keeps ownership of rejecting genuinely invalid names. + """ + if not isinstance(value, str): + return value + return _QUANT_ALGO_ALIASES.get(value.lower(), value) + + +def _canonicalize_algos(result: Dict[str, Any]) -> Dict[str, Any]: + """Canonicalize ``quant_algo`` at the top level and inside ``quantized_layers``. + + MIXED_PRECISION checkpoints carry the real algo names per layer, so the + top-level key alone is not enough. Nested dicts are rebuilt rather than + mutated in place: ``result`` is only a shallow copy of the caller's config. + """ + if "quant_algo" in result: + result["quant_algo"] = canonicalize_quant_algo(result["quant_algo"]) + layers = result.get("quantized_layers") + if isinstance(layers, dict): + result["quantized_layers"] = { + name: ( + {**cfg, "quant_algo": canonicalize_quant_algo(cfg["quant_algo"])} + if isinstance(cfg, dict) and "quant_algo" in cfg + else cfg + ) + for name, cfg in layers.items() + } + return result + def _kv_cache_scheme_to_algo(scheme: Any) -> Optional[str]: """Translate modelopt 1.x ``kv_cache_scheme`` to a legacy algo name. @@ -97,10 +142,7 @@ def read_modelopt_quant_config(raw: Dict[str, Any]) -> Dict[str, Any]: f"Not a modelopt quant config (producer={raw.get('producer')!r}, " f"quant_method={raw.get('quant_method')!r})" ) - # Canonicalize the fp8_pb_wo legacy alias. - if result.get("quant_algo") == "fp8_pb_wo": - result["quant_algo"] = "FP8_BLOCK_SCALES" - return result + return _canonicalize_algos(result) def warn_if_inline_diverges( diff --git a/tests/integration/test_lists/test-db/l0_b300.yml b/tests/integration/test_lists/test-db/l0_b300.yml index 747679751e5b..a351e591b04e 100644 --- a/tests/integration/test_lists/test-db/l0_b300.yml +++ b/tests/integration/test_lists/test-db/l0_b300.yml @@ -38,6 +38,23 @@ l0_b300: - unittest/_torch/modules/test_moe_routing.py - unittest/_torch/modules/fused_moe/test_deepgemm_fused_gather_finalize.py - unittest/_torch/modules/fused_moe/test_deepgemm_fused_expand_quant.py + # ------------- Kimi K3 NVFP4: expert loading + SiTU --------------- + # Listed test-by-test on purpose. The rest of this file has 19 failures that + # predate this work (verified by running the base revision's copy against the + # same source), so registering the whole module would import them into CI as + # if they were ours. Those are reported separately; these are green. + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_kimi_k3_expert_ckpt_spec_selection + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_fp8_block_scaled_dequantization + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_materialize_handles_scalar_lazy_safetensors + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_matches_whole_checkpoint[threads1] + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_matches_whole_checkpoint[threads4] + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_w1_w3_order_is_load_bearing + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_expert_load_rejects_duplicate_slot + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streaming_drains_staging_per_expert + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streamed_experts_forward_runs + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_cutlass_situ_bf16_matches_reference + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_kernel_actually_applies_situ + - unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py::test_nvfp4_streamed_experts_match_situ_reference[static_1.0] # ------------- MoE: test_moe_backend (by backend) --------------- - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_bf16_unquantized_moe - unittest/_torch/modules/moe/test_moe_backend.py::test_trtllm_fp8_block_scales_fused_shared_experts 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 11ef6d557e50..65c6eaa83187 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 @@ -48,8 +48,9 @@ _MEGA_MOE_SYMM_BUFFER_CACHE, ) from tensorrt_llm._torch.utils import ActType_TrtllmGen -from tensorrt_llm._utils import get_free_port +from tensorrt_llm._utils import get_free_port, get_sm_version from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantAlgo situ_supported = pytest.mark.skipif( not is_native_situ_supported(), @@ -623,6 +624,7 @@ def _make_routed_moe( gate, num_experts=_TP_EXPERTS, moe_backend="TRTLLM", + routed_quant_config=None, ): """Mirror KimiK3MoERuntime's create_moe call on a single-rank mapping.""" from transformers.configuration_utils import PretrainedConfig @@ -650,7 +652,11 @@ def _make_routed_moe( dtype=torch.bfloat16, reduce_results=True, model_config=model_config, - override_quant_config=QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8), + override_quant_config=( + routed_quant_config + if routed_quant_config is not None + else QuantConfig(quant_algo=QuantAlgo.W4A8_MXFP4_MXFP8) + ), layer_idx=0, communication_method=None, ) @@ -660,6 +666,17 @@ def _make_routed_moe( trtllm_gen_activation_alpha=4.0, trtllm_gen_activation_beta=25.0, ) + elif moe_backend == "CUTLASS": + # Mirror KimiK3MoERuntime exactly: CUTLASS is the one backend that + # takes SiTU as an ActivationType (the others carry it out of band), + # and that choice decides the FC1 weight geometry. + from tensorrt_llm._torch.utils import ActivationType + + moe_kwargs.update( + activation_type=ActivationType.SiTu, + swiglu_alpha=torch.full((num_experts,), 4.0, dtype=torch.float32, device="cuda"), + swiglu_beta=torch.full((num_experts,), 25.0, dtype=torch.float32, device="cuda"), + ) else: moe_kwargs.update( activation="situ", @@ -866,3 +883,835 @@ def test_megamoe_deepgemm_situ_matches_trtllm_gen( ) assert cosine > 0.998 assert relative_l2 < 0.06 + + +# --------------------------------------------------------------------------- +# NVFP4 streaming expert loading (nvidia/Kimi-K3-NVFP4). +# +# K3 never calls the MoE backend's whole-checkpoint ``load_weights``: the +# checkpoint is 1.5 TB, so experts are streamed one at a time through a +# per-expert adapter. For NVFP4 that adapter is +# ``NVFP4FusedMoEMethod.load_streaming_nvfp4_expert``, and the property that +# matters is that it lands the model in exactly the state ``load_weights`` +# would have. Nothing about a wrong [w1 | w3] order or a lost slot raises — +# shapes still match and only accuracy moves — so it is asserted here. +# --------------------------------------------------------------------------- + +_NVFP4_GROUP_SIZE = 16 + +nvfp4_moe_supported = pytest.mark.skipif( + not torch.cuda.is_available() or get_sm_version() < 100, + reason="NVFP4 Cutlass MoE requires Blackwell (SM100+)", +) + + +def _make_nvfp4_expert_bank(num_experts, intermediate, hidden, seed=907): + """Random NVFP4 tensors in ``nvidia/Kimi-K3-NVFP4`` checkpoint layout.""" + gen = torch.Generator().manual_seed(seed) + + def nibbles(*shape): + return torch.randint(0, 256, shape, generator=gen, dtype=torch.uint8) + + def block_scales(*shape): + # Modest exponents keep the dequantized weights well-conditioned. + return ( + torch.randint(120, 132, shape, generator=gen, dtype=torch.int32) + .to(torch.float32) + .div(126.0) + .to(torch.float8_e4m3fn) + ) + + bank = [] + for _ in range(num_experts): + expert = { + "w1.weight": nibbles(intermediate, hidden // 2), + "w1.weight_scale": block_scales(intermediate, hidden // _NVFP4_GROUP_SIZE), + "w3.weight": nibbles(intermediate, hidden // 2), + "w3.weight_scale": block_scales(intermediate, hidden // _NVFP4_GROUP_SIZE), + "w2.weight": nibbles(hidden, intermediate // 2), + "w2.weight_scale": block_scales(hidden, intermediate // _NVFP4_GROUP_SIZE), + } + for w in ("w1", "w2", "w3"): + # The real checkpoint stores one global weight scale per tensor and + # a static activation scale of 1.0. + expert[f"{w}.weight_scale_2"] = torch.tensor(0.00012207, dtype=torch.float32) + expert[f"{w}.input_scale"] = torch.tensor(1.0, dtype=torch.float32) + bank.append(expert) + return bank + + +def _make_nvfp4_moe(gate, num_experts=_TP_EXPERTS): + from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig + + return _make_routed_moe( + _TP_INTERMEDIATE, + gate, + num_experts=num_experts, + moe_backend="CUTLASS", + routed_quant_config=QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=_NVFP4_GROUP_SIZE), + ) + + +def _stream_nvfp4_bank(moe, bank, swap_w1_w3=False, num_threads=1): + """Load a bank through the per-expert streaming adapter K3 uses.""" + backend = moe.backend + quant_method = backend.quant_method + quant_method.prepare_streaming_expert_load(backend) + + def load_one(expert_id): + tensors = dict(bank[expert_id]) + if swap_w1_w3: + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale"): + tensors[f"w1.{kind}"], tensors[f"w3.{kind}"] = ( + tensors[f"w3.{kind}"], + tensors[f"w1.{kind}"], + ) + quant_method.load_streaming_nvfp4_expert( + backend, + global_expert_id=expert_id, + local_slot_id=expert_id, + **{ + f"{w}_{kind}": tensors[f"{w}.{kind}"] + for w in ("w1", "w2", "w3") + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") + }, + ) + + if num_threads > 1: + from concurrent.futures import ThreadPoolExecutor + + with ThreadPoolExecutor(max_workers=num_threads) as pool: + list(pool.map(load_one, range(len(bank)))) + else: + for expert_id in range(len(bank)): + load_one(expert_id) + + backend.process_weights_after_loading() + return backend + + +def _load_nvfp4_bank_whole(moe, bank): + """Load the same bank through the stock whole-checkpoint path.""" + weights = {} + for expert_id, tensors in enumerate(bank): + for key, value in tensors.items(): + weights[f"{expert_id}.{key}"] = value + moe.backend.load_weights([weights]) + return moe.backend + + +_NVFP4_LOADED_STATE = ( + "w3_w1_weight", + "w2_weight", + "w3_w1_weight_scale", + "w2_weight_scale", + "fc31_alpha", + "fc2_alpha", + "fc31_input_scale", + "fc2_input_scale", +) + + +def _bitwise_equal(actual, expected): + """Exact equality, including for the packed / float8 buffers. + + ``torch.equal`` refuses float8 operands, so those are compared through a + uint8 reinterpretation — which in turn needs a non-0-dim tensor, hence the + ``reshape(-1)`` for the scalar input scales. + """ + if actual.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + actual = actual.reshape(-1).view(torch.uint8) + expected = expected.reshape(-1).view(torch.uint8) + return torch.equal(actual, expected) + + +@nvfp4_moe_supported +@pytest.mark.parametrize("num_threads", [1, 4], ids=lambda n: f"threads{n}") +def test_nvfp4_streaming_expert_load_matches_whole_checkpoint(num_threads): + """Streaming per-expert loading == stock whole-checkpoint loading. + + ``num_threads=4`` is the concurrency K3's file-grouped loader actually + uses; the staging containers are pre-created for exactly this reason, and + a lost slot would show up here as a mismatching buffer rather than as an + error. + """ + bank = _make_nvfp4_expert_bank(_TP_EXPERTS, _TP_INTERMEDIATE, _TP_HIDDEN) + gate = _make_test_gate() + + reference = _load_nvfp4_bank_whole(_make_nvfp4_moe(gate), bank) + streamed = _stream_nvfp4_bank(_make_nvfp4_moe(gate), bank, num_threads=num_threads) + + assert streamed._streamed_expert_slots == set(range(_TP_EXPERTS)) + for name in _NVFP4_LOADED_STATE: + expected = getattr(reference, name).data + actual = getattr(streamed, name).data + assert actual.shape == expected.shape, name + assert _bitwise_equal(actual, expected), name + + +@nvfp4_moe_supported +def test_nvfp4_streaming_expert_load_w1_w3_order_is_load_bearing(): + """Mutation test for the trap that does not raise. + + ``w3_w1_weight`` holds the two halves in a fixed order. Feeding w1 and w3 + the wrong way round keeps every shape valid and every load silent, so + without this assertion the only detector would be an accuracy run. + """ + bank = _make_nvfp4_expert_bank(_TP_EXPERTS, _TP_INTERMEDIATE, _TP_HIDDEN) + gate = _make_test_gate() + + correct = _stream_nvfp4_bank(_make_nvfp4_moe(gate), bank) + swapped = _stream_nvfp4_bank(_make_nvfp4_moe(gate), bank, swap_w1_w3=True) + + assert not _bitwise_equal(swapped.w3_w1_weight.data, correct.w3_w1_weight.data) + + +@nvfp4_moe_supported +def test_nvfp4_streaming_expert_load_rejects_duplicate_slot(): + bank = _make_nvfp4_expert_bank(2, _TP_INTERMEDIATE, _TP_HIDDEN) + gate = _make_test_gate() + moe = _make_nvfp4_moe(gate) + backend = moe.backend + backend.quant_method.prepare_streaming_expert_load(backend) + + kwargs = dict( + global_expert_id=0, + local_slot_id=0, + **{ + f"{w}_{kind}": bank[0][f"{w}.{kind}"] + for w in ("w1", "w2", "w3") + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") + }, + ) + backend.quant_method.load_streaming_nvfp4_expert(backend, **kwargs) + with pytest.raises(ValueError, match="loaded twice"): + backend.quant_method.load_streaming_nvfp4_expert(backend, **kwargs) + + +def test_kimi_k3_expert_ckpt_spec_selection(): + """The three loader call sites share one layout decision.""" + from tensorrt_llm._torch.models.modeling_kimi_linear import _k3_expert_ckpt_spec + + mxfp4 = _k3_expert_ckpt_spec(QuantAlgo.W4A8_MXFP4_MXFP8) + assert mxfp4.kinds == ("weight_packed", "weight_scale") + assert not mxfp4.needs_layer_finalize + + nvfp4 = _k3_expert_ckpt_spec(QuantAlgo.NVFP4) + assert nvfp4.kinds == ("weight", "weight_scale", "weight_scale_2", "input_scale") + assert nvfp4.needs_layer_finalize + + with pytest.raises(NotImplementedError, match="no per-expert checkpoint layout"): + _k3_expert_ckpt_spec(QuantAlgo.FP8) + + +def test_materialize_handles_scalar_lazy_safetensors(tmp_path): + """The NVFP4 checkpoint stores weight_scale_2 / input_scale as 0-dim. + + ``_materialize`` realizes a lazy slice with ``[:]``, which a 0-dim entry + rejects — so the fallback (non-file-grouped) expert load died on the very + first NVFP4 scalar it touched. + """ + import safetensors.torch + from safetensors import safe_open + + from tensorrt_llm._torch.models.modeling_kimi_linear import _materialize + + scalar = torch.tensor(0.00012207, dtype=torch.float32) + matrix = torch.arange(6, dtype=torch.float32).reshape(2, 3) + path = tmp_path / "model.safetensors" + safetensors.torch.save_file({"scalar": scalar, "matrix": matrix}, path) + + with safe_open(str(path), framework="pt", device="cpu") as handle: + assert torch.equal(_materialize(handle.get_slice("scalar")), scalar) + assert torch.equal(_materialize(handle.get_slice("matrix")), matrix) + assert torch.equal(_materialize(scalar), scalar) + + +@nvfp4_moe_supported +def test_nvfp4_streamed_experts_forward_runs(): + """End-to-end guard for the geometry the loader writes into. + + The buffer-equality tests compare two loaders against each other, so they + stay green even when both write into a wrongly-shaped destination. Only + running the kernel catches that. + """ + bank = _make_nvfp4_expert_bank(_TP_EXPERTS, _TP_INTERMEDIATE, _TP_HIDDEN) + gate = _make_test_gate() + moe = _make_nvfp4_moe(gate) + + backend = _stream_nvfp4_bank(moe, bank) + assert backend.w3_w1_weight.shape[1] == 2 * _TP_INTERMEDIATE + backend._weights_transformed = False + moe.post_load_weights() + + torch.manual_seed(5) + x = torch.randn(16, _TP_HIDDEN, dtype=torch.bfloat16, device="cuda") * 0.5 + out = moe.forward(x, gate.compute_logits(x), all_rank_num_tokens=None) + assert out.shape == x.shape + assert torch.isfinite(out).all() + + +def _situ_reference_moe(x, router_logits, routing_method, w1, w2, w3, beta, linear_beta): + """Golden K3 routed-MoE, straight from the SituAndMul definition. + + ``out = [beta*tanh(g/beta)*sigmoid(g)] * [linear_beta*tanh(u/linear_beta)]`` + with g = w1(x) (gate_proj) and u = w3(x) (up_proj). + """ + ids, weights = routing_method.apply(router_logits) + out = torch.zeros_like(x, dtype=torch.float32) + xf = x.float() + for token in range(x.shape[0]): + for slot in range(ids.shape[1]): + e = int(ids[token, slot]) + g = xf[token] @ w1[e].float().t() + u = xf[token] @ w3[e].float().t() + situ = beta * torch.tanh(g / beta) * torch.sigmoid(g) + u = linear_beta * torch.tanh(u / linear_beta) + out[token] += float(weights[token, slot]) * ((situ * u) @ w2[e].float().t()) + return out + + +@nvfp4_moe_supported +def test_cutlass_situ_bf16_matches_reference(): + """CUTLASS + SiTU in BF16, against the golden activation. No quantization. + + This is the unambiguous form of the question the NVFP4 GSM8K collapse + raised. With unquantized weights there is no 4-bit noise to hide behind, + so a wrong FC1 half assignment or a swapped alpha/beta cannot pass, and + nothing about the NVFP4 loader is involved. + """ + from tensorrt_llm.models.modeling_utils import QuantConfig + + num_experts, hidden, inter = _TP_EXPERTS, _TP_HIDDEN, 512 + gate = _make_test_gate(num_experts=num_experts) + moe = _make_routed_moe( + inter, + gate, + num_experts=num_experts, + moe_backend="CUTLASS", + routed_quant_config=QuantConfig(), + ) + + torch.manual_seed(77) + w1 = [ + torch.randn(inter, hidden, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + w3 = [ + torch.randn(inter, hidden, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + w2 = [ + torch.randn(hidden, inter, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + weights = {} + for e in range(num_experts): + weights[f"{e}.w1.weight"] = w1[e] + weights[f"{e}.w2.weight"] = w2[e] + weights[f"{e}.w3.weight"] = w3[e] + moe.backend.load_weights([weights]) + moe.backend._weights_transformed = False + moe.post_load_weights() + + x = torch.randn(8, hidden, dtype=torch.bfloat16, device="cuda") * 0.5 + router_logits = gate.compute_logits(x) + actual = moe.forward(x, router_logits, all_rank_num_tokens=None).float() + expected = _situ_reference_moe( + x, router_logits, gate.routing_method, w1, w2, w3, beta=4.0, linear_beta=25.0 + ) + + cosine = torch.nn.functional.cosine_similarity(actual.flatten(), expected.flatten(), dim=0) + rel_l2 = torch.linalg.vector_norm(actual - expected) / torch.linalg.vector_norm(expected) + print(f"cutlass-situ-bf16 vs reference: cosine={cosine.item():.6f} rel_l2={rel_l2.item():.6f}") + assert cosine > 0.999, f"cosine={cosine.item()}, rel_l2={rel_l2.item()}" + + +def _quantize_expert_to_nvfp4(w1, w2, w3, input_scale): + """One expert's checkpoint tensors, in nvidia/Kimi-K3-NVFP4's layout. + + w1 and w3 must share a global scale: the loader asserts their + weight_scale_2 match when it folds them into a single fc31 alpha. + """ + sv = _NVFP4_GROUP_SIZE + w13_global = torch.min((448 * 6) / w1.abs().max().float(), (448 * 6) / w3.abs().max().float()) + w2_global = (448 * 6) / w2.abs().max().float() + + def q(w, g): + packed, sf = torch.ops.trtllm.fp4_quantize(w, g, sv, False) + # The checkpoint stores plain per-block scales; the swizzle is applied + # by the backend at load time. + return packed.cpu(), torch.ops.trtllm.block_scale_interleave_reverse( + sf.cpu().view(w.shape[0], -1) + ) + + out = {} + for name, w, g in (("w1", w1, w13_global), ("w3", w3, w13_global), ("w2", w2, w2_global)): + packed, sf = q(w, g) + out[f"{name}.weight"] = packed + out[f"{name}.weight_scale"] = sf + out[f"{name}.weight_scale_2"] = (1.0 / g).cpu().float() + out[f"{name}.input_scale"] = torch.tensor(input_scale, dtype=torch.float32) + return out + + +@nvfp4_moe_supported +@pytest.mark.parametrize( + "input_scale", + [ + 1.0, + pytest.param( + None, + marks=pytest.mark.xfail( + strict=True, + reason="OPEN: the conventional NVFP4 activation global scale is the WORSE " + "of the two. Measured cosine/rel_l2 vs the golden activation: " + "static 1.0 -> 0.9727/0.234, derived (448*6)/amax -> 0.8757/0.538. " + "Correct NVFP4 scaling maps amax onto the FP4xE4M3 range and should " + "reduce error, not double it, so this points at the activation scale " + "being applied more than once somewhere on the CUTLASS SiTU path -- " + "which a checkpoint shipping input_scale=1.0 cannot expose, because " + "applying 1.0 twice is still 1.0. nvidia/Kimi-K3-NVFP4 ships exactly " + "that, so this does not affect the current bring-up (GSM8K 96.40 " + "confirms the shipped configuration). strict=True so that fixing it " + "reports here instead of passing silently.", + ), + ), + ], + ids=["static_1.0", "derived_act_scale"], +) +def test_nvfp4_streamed_experts_match_situ_reference(input_scale): + """Streamed NVFP4 experts against the golden activation. + + The streamed-vs-whole-checkpoint tests compare two loaders that share all + the scale handling, so they agree whether or not that handling is right. + This compares against the definition instead. + + ``static_1.0`` is what nvidia/Kimi-K3-NVFP4 actually ships for + ``input_scale``; ``derived_act_scale`` is the conventional value computed + from the activations. Splitting them separates "the loader is wrong" from + "this checkpoint's static activation scale is being interpreted wrongly". + """ + num_experts, hidden, inter = _TP_EXPERTS, _TP_HIDDEN, _TP_INTERMEDIATE + gate = _make_test_gate(num_experts=num_experts) + + torch.manual_seed(91) + x = torch.randn(8, hidden, dtype=torch.bfloat16, device="cuda") * 0.5 + act_scale = float(x.abs().max().float() / (448 * 6)) if input_scale is None else input_scale + + w1 = [ + torch.randn(inter, hidden, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + w3 = [ + torch.randn(inter, hidden, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + w2 = [ + torch.randn(hidden, inter, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + bank = [_quantize_expert_to_nvfp4(w1[e], w2[e], w3[e], act_scale) for e in range(num_experts)] + + moe = _make_nvfp4_moe(gate, num_experts=num_experts) + _stream_nvfp4_bank(moe, bank) + moe.backend._weights_transformed = False + moe.post_load_weights() + + router_logits = gate.compute_logits(x) + actual = moe.forward(x, router_logits, all_rank_num_tokens=None).float() + expected = _situ_reference_moe( + x, router_logits, gate.routing_method, w1, w2, w3, beta=4.0, linear_beta=25.0 + ) + + cosine = torch.nn.functional.cosine_similarity(actual.flatten(), expected.flatten(), dim=0) + rel_l2 = torch.linalg.vector_norm(actual - expected) / torch.linalg.vector_norm(expected) + print( + f"nvfp4-streamed[{input_scale}] vs reference: cosine={cosine.item():.6f} rel_l2={rel_l2.item():.6f}" + ) + # The shipped configuration measures 0.9727 / 0.234. That residual is + # activation quantization: this path takes the ACTIVATIONS to FP4 as well + # as the weights, on random Gaussian data whose per-block dynamic range is + # a worst case real activations do not have. 0.95 leaves room for that + # while still catching a gross scale error, which lands far below it. + # + # An earlier version of this comment justified the bound by claiming both + # parameterizations landed identically, "the signature of a noise floor". + # That observation was an artifact of the harness truncating pytest output + # to the last failure; the two in fact differ a lot (see the xfail above), + # and the claim was never measured. The structural guards do not rest on + # this number at all: they are the BF16 comparison against the same golden + # reference and the SiTU-vs-SwiGLU discriminator, both tolerance-free. The + # accuracy gate for the real checkpoint is GSM8K. + assert cosine > 0.95, f"cosine={cosine.item()}, rel_l2={rel_l2.item()}" + + +def _swiglu_reference_moe(x, router_logits, routing_method, w1, w2, w3, alpha, beta): + """Same routing/geometry as the SiTU reference, but CUTLASS's SwigluBias. + + ``gate*sigmoid(gate*alpha)*(linear+beta)`` -- what the FC1 epilogue would + compute if the activation enum did not resolve to SiTu. + """ + ids, weights = routing_method.apply(router_logits) + out = torch.zeros_like(x, dtype=torch.float32) + xf = x.float() + for token in range(x.shape[0]): + for slot in range(ids.shape[1]): + e = int(ids[token, slot]) + g = xf[token] @ w1[e].float().t() + u = xf[token] @ w3[e].float().t() + h = g * torch.sigmoid(g * alpha) * (u + beta) + out[token] += float(weights[token, slot]) * (h @ w2[e].float().t()) + return out + + +@nvfp4_moe_supported +def test_nvfp4_kernel_actually_applies_situ(): + """Which activation does the QUANTIZED kernel actually run? + + BF16 + SiTU already matches the golden reference, so the activation, the + FC1 half assignment and the alpha/beta mapping are all correct in the + unquantized path. If the NVFP4 path silently resolved to a different + gated activation -- e.g. because SiTu is not instantiated for the FP4 + epilogue and something falls back -- the output would be structurally + wrong in exactly the way the GSM8K collapse showed, while every + shape-and-buffer check stayed green. + + Reported rather than merely asserted: which reference the kernel is + closer to is the diagnosis. + """ + num_experts, hidden, inter = _TP_EXPERTS, _TP_HIDDEN, _TP_INTERMEDIATE + gate = _make_test_gate(num_experts=num_experts) + + torch.manual_seed(91) + x = torch.randn(8, hidden, dtype=torch.bfloat16, device="cuda") * 0.5 + act_scale = float(x.abs().max().float() / (448 * 6)) + w1 = [ + torch.randn(inter, hidden, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + w3 = [ + torch.randn(inter, hidden, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + w2 = [ + torch.randn(hidden, inter, dtype=torch.bfloat16, device="cuda") * 0.05 + for _ in range(num_experts) + ] + bank = [_quantize_expert_to_nvfp4(w1[e], w2[e], w3[e], act_scale) for e in range(num_experts)] + + moe = _make_nvfp4_moe(gate, num_experts=num_experts) + _stream_nvfp4_bank(moe, bank) + moe.backend._weights_transformed = False + moe.post_load_weights() + + router_logits = gate.compute_logits(x) + actual = moe.forward(x, router_logits, all_rank_num_tokens=None).float() + + situ = _situ_reference_moe( + x, router_logits, gate.routing_method, w1, w2, w3, beta=4.0, linear_beta=25.0 + ) + swiglu = _swiglu_reference_moe( + x, router_logits, gate.routing_method, w1, w2, w3, alpha=4.0, beta=25.0 + ) + + def score(ref): + cos = torch.nn.functional.cosine_similarity(actual.flatten(), ref.flatten(), dim=0) + l2 = torch.linalg.vector_norm(actual - ref) / torch.linalg.vector_norm(ref) + return cos.item(), l2.item() + + situ_cos, situ_l2 = score(situ) + swiglu_cos, swiglu_l2 = score(swiglu) + print( + f"NVFP4 kernel vs SiTU ref: cosine={situ_cos:.6f} rel_l2={situ_l2:.6f}\n" + f"NVFP4 kernel vs SwiGLU ref: cosine={swiglu_cos:.6f} rel_l2={swiglu_l2:.6f}" + ) + assert situ_cos > swiglu_cos, ( + f"the NVFP4 kernel matches a SwiGLU reference better than the SiTU one " + f"(situ={situ_cos:.6f}, swiglu={swiglu_cos:.6f}): the quantized path is " + f"not applying SiTU" + ) + + +def test_fp8_block_scaled_dequantization(): + """FP8_PB_WO attention weights must be dequantized, never reinterpreted. + + nvidia/Kimi-K3-NVFP4 stores the attention projections as FP8 E4M3 plus a + per-128x128-block FP32 scale; moonshotai/Kimi-K3 stores plain BF16 at the + SAME shape. Nothing downstream can tell them apart, so before this the + loader's shape check passed and src.to(param.dtype) reinterpreted the + quantized values as real ones, silently dropping the scale. + """ + from tensorrt_llm._torch.models.modeling_kimi_linear import _dequantize_fp8_block_scaled + + n, k = 256, 384 + torch.manual_seed(3) + quantized = (torch.randn(n, k) * 2).to(torch.float8_e4m3fn) + scale = torch.rand(n // 128, 1, k // 128, 1) * 3 + 0.5 + weights = {"w.weight": quantized, "w.weight_scale": scale} + + got = _dequantize_fp8_block_scaled("w.weight", quantized, weights) + assert got.dtype == torch.bfloat16 + expected = quantized.float() * scale.reshape(n // 128, k // 128).repeat_interleave( + 128, 0 + ).repeat_interleave(128, 1) + torch.testing.assert_close(got.float(), expected.to(torch.bfloat16).float()) + + # A BF16 checkpoint (the MXFP4 original) must pass through untouched. + bf16 = torch.randn(8, 8, dtype=torch.bfloat16) + assert _dequantize_fp8_block_scaled("w.weight", bf16, {}) is bf16 + + # FP8 without a scale is the silent-corruption case; it must raise. + with pytest.raises(KeyError, match="refusing"): + _dequantize_fp8_block_scaled("w.weight", quantized, {"w.weight": quantized}) + + +@nvfp4_moe_supported +def test_nvfp4_streaming_drains_staging_per_expert(): + """The staged halves must not accumulate across the load. + + Cutlass defers the cat+pad of w3_w1 to process_weights_after_loading by + staging both halves per expert, which is a second copy of the + routed-expert weights. Draining per LAYER does not bound that for Kimi + K3, whose loader groups work by shard FILE: a layer completes only when + the last of its slots happens to land, so many layers stage at once. + Draining per EXPERT does bound it. + """ + bank = _make_nvfp4_expert_bank(_TP_EXPERTS, _TP_INTERMEDIATE, _TP_HIDDEN) + gate = _make_test_gate() + moe = _make_nvfp4_moe(gate) + backend = moe.backend + quant_method = backend.quant_method + quant_method.prepare_streaming_expert_load(backend) + + for expert_id in range(_TP_EXPERTS): + quant_method.load_streaming_nvfp4_expert( + backend, + global_expert_id=expert_id, + local_slot_id=expert_id, + **{ + f"{w}_{kind}": bank[expert_id][f"{w}.{kind}"] + for w in ("w1", "w2", "w3") + for kind in ("weight", "weight_scale", "weight_scale_2", "input_scale") + }, + ) + # Nothing may still be staged for a slot that has been loaded. + assert not backend.tmp_cutlass_w3_w1_weights, ( + f"{len(backend.tmp_cutlass_w3_w1_weights)} staged weight entries " + f"still held after loading slot {expert_id}" + ) + assert not backend.tmp_cutlass_w3_w1_weight_scales, ( + f"{len(backend.tmp_cutlass_w3_w1_weight_scales)} staged scale entries " + f"still held after loading slot {expert_id}" + ) + + backend.process_weights_after_loading() + # And the result must still equal the whole-checkpoint path. + reference = _load_nvfp4_bank_whole(_make_nvfp4_moe(gate), bank) + for name in _NVFP4_LOADED_STATE: + assert _bitwise_equal(getattr(backend, name).data, getattr(reference, name).data), name + + +def _checkpoint_scale_2d(scale: "torch.Tensor") -> "torch.Tensor": + """Mirror the loader's normalization of the checkpoint's 4-D block scale.""" + return scale.reshape(scale.shape[0], scale.shape[2]) if scale.dim() == 4 else scale + + +@nvfp4_moe_supported +def test_fp8_checkpoint_scale_bridge_matches_the_roundtrip_path(): + """FP8_PB_WO straight from the checkpoint == the BF16 round trip. + + The shipping FP8 weight-read path starts from a BF16 weight and runs + per_block_cast_to_fp8 -> resmooth_to_fp8_e8m0 -> deep_gemm layout. A + checkpoint that already stores FP8_PB_WO has done the first step, so + reading it resident should be the same two remaining steps. + + That equality is the whole premise of keeping attention FP8 on device + rather than expanding it to BF16, and it is not obvious: fp8_swap_ab_gemm + runs with disable_ue8m0_cast=True, so it consumes a pre-formatted UE8M0 + scale and would misread the plain FP32 block scale the checkpoint ships. + This asserts the bridge, on the same values, before any loader is rewired + to depend on it. + """ + from tensorrt_llm._torch.models.modeling_kimi_linear import ( + _Fp8BlockScaleWeightReadLinear as FP8Linear, + ) + from tensorrt_llm.deep_gemm.utils.math import per_block_cast_to_fp8 + + out_features, in_features = 256, 512 + torch.manual_seed(4) + weight = torch.randn(out_features, in_features, dtype=torch.bfloat16, device="cuda") * 0.05 + + # Path A: what production does today, from a BF16 weight. + w_a, s_a = FP8Linear.quantize_weight(weight) + + # Path B: what a checkpoint hands us -- FP8 + FP32 128x128 block scale -- + # fed straight into the scale bridge. + ckpt_fp8, ckpt_scale = per_block_cast_to_fp8(weight, use_ue8m0=False) + # nvidia/Kimi-K3-NVFP4 stores the block scale 4-D as + # [ceil(N/128), 1, ceil(K/128), 1]. Feed that exact shape, not the 2-D one + # per_block_cast_to_fp8 happens to return: an earlier version of this test + # used the 2-D form, passed, and the real checkpoint then tripped + # transform_sf_into_required_layout's rank assert on a four-node run. + nb_m, nb_k = ckpt_scale.shape + ckpt_scale_4d = ckpt_scale.float().reshape(nb_m, 1, nb_k, 1) + w_b, s_b = FP8Linear.prepare_checkpoint_scale(ckpt_fp8, _checkpoint_scale_2d(ckpt_scale_4d)) + + assert w_a.shape == w_b.shape and s_a.shape == s_b.shape + # _bitwise_equal reinterprets only float8 (which torch.equal refuses); the + # prepared scale is int32 in deep_gemm's packed layout, whose stride makes a + # uint8 view invalid anyway. + assert _bitwise_equal(w_a, w_b), "FP8 weights differ" + assert _bitwise_equal(s_a, s_b), "prepared scales differ" + + # And the GEMM they drive agrees. + x = torch.randn(8, in_features, dtype=torch.bfloat16, device="cuda") * 0.5 + lin_a = FP8Linear(w_a, s_a, out_features) + lin_b = FP8Linear(w_b, s_b, out_features) + torch.testing.assert_close(lin_a.forward(x), lin_b.forward(x)) + + +@nvfp4_moe_supported +def test_fp8_weight_read_prefers_the_checkpoint_pair(): + """``from_linear`` must use a stashed FP8_PB_WO pair when one is present. + + Three call sites convert attention projections to the FP8 weight-read + Linear, all through ``from_linear``, so that is the one place the + checkpoint-direct route has to be honoured. Without this the loader would + stash the pair and nothing would read it -- silently falling back to the + BF16 round trip, which is exactly what this is meant to avoid. + """ + import torch.nn as nn + + from tensorrt_llm._torch.models.modeling_kimi_linear import _K3_CKPT_FP8_ATTR + from tensorrt_llm._torch.models.modeling_kimi_linear import ( + _Fp8BlockScaleWeightReadLinear as FP8Linear, + ) + from tensorrt_llm.deep_gemm.utils.math import per_block_cast_to_fp8 + + out_features, in_features = 256, 512 + torch.manual_seed(6) + weight = torch.randn(out_features, in_features, dtype=torch.bfloat16, device="cuda") * 0.05 + ckpt_fp8, ckpt_scale = per_block_cast_to_fp8(weight, use_ue8m0=False) + + linear = nn.Linear(in_features, out_features, bias=False, dtype=torch.bfloat16, device="cuda") + with torch.no_grad(): + # Deliberately NOT the weight the pair came from: if from_linear + # re-quantized linear.weight instead of using the pair, the result + # would follow this garbage and the assert below would catch it. + linear.weight.copy_(torch.zeros_like(weight)) + setattr(linear.weight, _K3_CKPT_FP8_ATTR, (ckpt_fp8, ckpt_scale.float())) + + converted = FP8Linear.from_linear(linear) + expected_w, expected_s = FP8Linear.prepare_checkpoint_scale(ckpt_fp8, ckpt_scale.float()) + assert _bitwise_equal(converted.weight, expected_w) + assert _bitwise_equal(converted.weight_scale, expected_s) + # float8 supports almost no arithmetic (no abs_cuda), so probe the bytes. + assert converted.weight.view(torch.uint8).any().item(), ( + "fell back to re-quantizing the zeroed BF16 weight" + ) + + +@nvfp4_moe_supported +def test_fused_fp8_from_checkpoint_slices_matches_the_bf16_concat(): + """Fusing checkpoint FP8 slices == quantizing the BF16 concatenation. + + The KDA conversion fuses q/k/v/g into one qkvg_proj by concatenating their + BF16 weights and quantizing the result. Constructing attention as FP8 + requires building that fused weight from the per-projection checkpoint + FP8 instead, with no BF16 anywhere -- which is only valid because every + out dim is a multiple of 128, so no 128x128 block straddles a boundary. + + That is the load-bearing assumption behind 3a, so it is checked directly + rather than inferred from the docstring that states it. + """ + from tensorrt_llm._torch.models.modeling_kimi_linear import ( + _Fp8BlockScaleWeightReadLinear as FP8Linear, + ) + from tensorrt_llm.deep_gemm.utils.math import per_block_cast_to_fp8 + + in_features = 512 + outs = [256, 256, 128, 384] # all multiples of 128, mixed sizes like q/k/v/g + torch.manual_seed(9) + parts = [torch.randn(o, in_features, dtype=torch.bfloat16, device="cuda") * 0.05 for o in outs] + + # What production does today: concatenate in BF16, then quantize. + fused_ref = FP8Linear.from_linear( + type( + "L", + (), + { + "weight": type("W", (), {"data": torch.cat(parts, dim=0)})(), + "bias": None, + "out_features": sum(outs), + }, + )() + ) + + # 3a's route: each projection arrives already FP8 from the checkpoint. + pairs = [] + for part in parts: + w, sc = per_block_cast_to_fp8(part, use_ue8m0=False) + pairs.append((w, sc.float())) + fused_new = FP8Linear.fuse_checkpoint_fp8(pairs) + + assert _bitwise_equal(fused_new.weight, fused_ref.weight), "fused FP8 weights differ" + assert _bitwise_equal(fused_new.weight_scale, fused_ref.weight_scale), "fused scales differ" + + x = torch.randn(8, in_features, dtype=torch.bfloat16, device="cuda") * 0.5 + torch.testing.assert_close(fused_new.forward(x), fused_ref.forward(x)) + + # And a non-128 out dim must be rejected, not silently mis-fused. + bad = torch.randn(96, in_features, dtype=torch.bfloat16, device="cuda") + w, sc = per_block_cast_to_fp8(bad, use_ue8m0=False) + with pytest.raises(ValueError, match="multiple of 128"): + FP8Linear.fuse_checkpoint_fp8([(w, sc.float())]) + + +@nvfp4_moe_supported +def test_fp8_placeholder_fill_contract(): + """The contract construction-time FP8 dispatch has to satisfy. + + A placeholder must cost nothing until filled, must produce exactly what + the eager constructors produce once filled, must refuse to run while + empty, and must refuse a second fill. The first property is the whole + point - DEP8's OOM was during module construction - and the last three + are what stop that saving from turning into silent garbage, which is how + every other mistake on this path has presented. + """ + from tensorrt_llm._torch.models.modeling_kimi_linear import ( + _Fp8BlockScaleWeightReadLinear as FP8Linear, + ) + from tensorrt_llm.deep_gemm.utils.math import per_block_cast_to_fp8 + + in_features, outs = 512, [256, 128] + torch.manual_seed(11) + parts = [torch.randn(o, in_features, dtype=torch.bfloat16, device="cuda") * 0.05 for o in outs] + pairs = [] + for part in parts: + w, sc = per_block_cast_to_fp8(part, use_ue8m0=False) + pairs.append((w, sc.float())) + + ph = FP8Linear.empty_placeholder(sum(outs), in_features) + assert ph.is_placeholder and ph.weight.numel() == 0 + + x = torch.randn(4, in_features, dtype=torch.bfloat16, device="cuda") * 0.5 + with pytest.raises(RuntimeError, match="never filled"): + ph.forward(x) + + ph.load_checkpoint_pair(pairs) + assert not ph.is_placeholder + + expected = FP8Linear.fuse_checkpoint_fp8(pairs) + assert _bitwise_equal(ph.weight, expected.weight) + assert _bitwise_equal(ph.weight_scale, expected.weight_scale) + torch.testing.assert_close(ph.forward(x), expected.forward(x)) + + with pytest.raises(RuntimeError, match="filled twice"): + ph.load_checkpoint_pair(pairs) + + # The single-projection case takes the same route. + single = FP8Linear.empty_placeholder(outs[0], in_features) + single.load_checkpoint_pair([pairs[0]]) + ref = FP8Linear.from_checkpoint_fp8(pairs[0][0], pairs[0][1], outs[0]) + assert _bitwise_equal(single.weight, ref.weight) diff --git a/tests/unittest/llmapi/test_llm_quant.py b/tests/unittest/llmapi/test_llm_quant.py index 141ea98ab644..316643827912 100644 --- a/tests/unittest/llmapi/test_llm_quant.py +++ b/tests/unittest/llmapi/test_llm_quant.py @@ -273,6 +273,48 @@ def test_quant_cfg_fp8_pb_wo_alias_canonicalized(): assert quant_config.group_size == 128 +@pytest.mark.cpu_only +def test_quant_cfg_fp8_pb_wo_alias_uppercase_per_layer(): + """Uppercase ``FP8_PB_WO`` inside ``quantized_layers`` is canonicalized. + + Shape taken from ``nvidia/Kimi-K3-NVFP4`` (modelopt 0.45): the top level is + MIXED_PRECISION and the real algos live per layer, spelled in upper case. + Before canonicalization reached the per-layer entries this raised + ``ValueError: 'FP8_PB_WO' is not a valid QuantAlgo``. + """ + attn = "language_model.model.layers.0.self_attn.q_proj" + experts = "language_model.model.layers.1.block_sparse_moe.experts" + with tempfile.TemporaryDirectory() as tmp_dir: + model_dir = Path(tmp_dir) + hf_quant_config_file = _write_hf_quant_config( + model_dir, { + "producer": { + "name": "modelopt", + "version": "0.45.0" + }, + "quantization": { + "quant_algo": "MIXED_PRECISION", + "kv_cache_quant_algo": None, + "quantized_layers": { + attn: { + "quant_algo": "FP8_PB_WO" + }, + experts: { + "quant_algo": "NVFP4", + "group_size": 16 + }, + }, + }, + }) + quant_config, layer_quant_config = ModelConfig.load_modelopt_quant_config( + hf_quant_config_file, model_dir, None) + assert quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + assert layer_quant_config[attn].quant_algo == QuantAlgo.FP8_BLOCK_SCALES + # Unaliased algos and their per-layer extras survive untouched. + assert layer_quant_config[experts].quant_algo == QuantAlgo.NVFP4 + assert layer_quant_config[experts].group_size == 16 + + @pytest.mark.cpu_only def test_quant_cfg_fp8_block_scales_trtllm_default_excludes(): """TRTLLM moe_backend + FP8_BLOCK_SCALES + no excludes → defaults applied.""" From 698f8fe370daf1d896aaf23ff46109403e88a8e1 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Sun, 16 Aug 2026 02:00:09 -0700 Subject: [PATCH 4/9] [None][feat] Kimi K3: MegaMoE CuteDSL SiTU support for the NVFP4 checkpoint Adds the MegaMoE CuteDSL path for Kimi K3 NVFP4 and validates it against the CUTLASS path at DEP16: GSM8K 96.32 vs 96.40, GPQA-Diamond 93.94 +/- 1.70 vs 91.92 +/- 1.94. Both differences are inside their error bars, and the configs differ from the CUTLASS ones in moe_config alone so the comparison is attributable to the backend. SiTU in the kernel. The activation is computed inside the existing alpha_swiglu_clamp loop with no tanh call at all: beta*tanh(x/beta) is rewritten as 2*beta*sigmoid(2x/beta) - beta, the same identity the codebase already uses for gelu_tanh, which keeps the whole core on the packed f32x2 path instead of dropping to scalar for a tanh that has no packed form. Validated against the golden SituAndMul in float64 at K3's beta=4.0 / linear_beta=25.0: max relative error ~1e-9. situ_beta and situ_linear_beta are baked into the generated kernel, so they join unique_id() - without that a SwiGLU-compiled kernel would be silently reused for a SiTU launch - and name() in both kernel files. The workspace probe takes the same parameters as the runner, so a parameter added to one and not the other would size a buffer for a kernel that is not the one launched. Five things had to be fixed before it ran, all found by running it: - An edit that inserted a method directly above another took its @staticmethod. No test constructed the class, so nothing was red. - K3 routes top-16 while the wrapper rejected anything above 13. That bound is EXPERIMENTALLY widened here and is NOT kernel-confirmed; see TODO-C in KIMI_K3_NVFP4_BRINGUP.md. Two benchmarks passing is evidence at two shapes, not coverage. - The backend keeps its raw NVFP4 params as 0-element placeholders and rematerializes them inside its own load_weights, which the streaming loader bypasses - so the per-expert writes indexed empty tensors. - Its aux-scale coverage check compares against num_experts, which is right for a whole-checkpoint load and wrong for a streaming EP one, where the rank's own 56 of 896 IS complete. - create_moe falls back silently when a backend declines a config, and MegaMoE declines easily (EP-only, its own token and top-k limits), so an explicit MEGAMOE_CUTEDSL request that quietly became CUTLASS would be benchmarked as if it were MegaMoE. Now guarded, mirroring MEGAMOE_DEEPGEMM. Three of those are one mistake: the streaming API sits on the shared NVFP4 base class, so it was assumed to cover this backend. The API was shared; the backend's assumptions were not. Known gap: finalize_streamed_expert is deliberately not implemented for this backend. Its staging dicts are not only staging - _initial_slot_ coverage() counts their entries - so draining per expert would report zero coverage. Bounding that footprint means moving the accounting onto _streamed_expert_slots first. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py --- ...extra_llm_options_nvfp4_dep16_megamoe.yaml | 37 ++++++++ ..._llm_options_nvfp4_dep16_megamoe_gpqa.yaml | 52 +++++++++++ .../custom_ops/cute_dsl_megamoe_custom_op.py | 20 ++++ .../mega_moe_nvfp4/epilogue_refactor.py | 93 ++++++++++++++++--- .../mega_moe_nvfp4/kernel_fc12.py | 11 ++- .../mega_moe_nvfp4/megamoe_kernel.py | 7 +- .../_torch/models/modeling_kimi_linear.py | 29 +++++- .../fused_moe/mega_moe/mega_moe_cute_dsl.py | 91 +++++++++++++++++- .../_torch/modules/fused_moe/quantization.py | 36 ++++++- .../modules/moe/test_kimi_k3_situ_moe.py | 5 +- 10 files changed, 354 insertions(+), 27 deletions(-) create mode 100644 examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_megamoe.yaml create mode 100644 examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_megamoe_gpqa.yaml diff --git a/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_megamoe.yaml b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_megamoe.yaml new file mode 100644 index 000000000000..4fee54900506 --- /dev/null +++ b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_megamoe.yaml @@ -0,0 +1,37 @@ +# Kimi K3 NVFP4 on DEP16 with the MegaMoE CuteDSL backend (4 nodes x 4 GPU). +# +# Same topology and sequence budget as the CUTLASS DEP16 GSM8K config, so a +# score difference is attributable to the MoE backend and nothing else. The +# CUTLASS run on this config scored 96.40 (job 475308) against a 96.47 MXFP4 +# baseline, which is what this is measured against. +# +# moe_config.max_num_tokens is deliberately ABSENT. For CUTLASS it is a +# per-call chunking bound and 131072 is harmless; for MEGAMOE_* the same key is +# the SymmBuffer capacity, further divided by ep_size under attention-DP. +# Carrying 131072 over would over-provision the generation SymmBuffer ~32x and +# eat the KV cache. ModelConfig's default (max_num_tokens x dp_size) lands +# correctly, so leave it unset -- see the MegaMoE disagg run record. +# +# MegaMoE is EP-only, which DEP16 already is (moe_ep=16, moe_tp=1). It also +# uses MNNVL symmetric memory, so all 4 nodes must sit in one NVL72 domain: +# pin them with sbatch -w to a single nvl72dNNN rack or the rendezvous fails +# with "invalid resource handle". +tensor_parallel_size: 16 +enable_attention_dp: true +moe_expert_parallel_size: 16 +max_batch_size: 8 +max_num_tokens: 8192 +max_seq_len: 8192 +trust_remote_code: true +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + enable_padding: true + max_batch_size: 8 +moe_config: + backend: MEGAMOE_CUTEDSL + use_low_precision_moe_combine: true +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.25 + tokens_per_block: 64 diff --git a/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_megamoe_gpqa.yaml b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_megamoe_gpqa.yaml new file mode 100644 index 000000000000..4409e5be2b49 --- /dev/null +++ b/examples/kimi_k3/eval_extra_llm_options_nvfp4_dep16_megamoe_gpqa.yaml @@ -0,0 +1,52 @@ +# Kimi K3 NVFP4 on DEP16 for GPQA-Diamond, MegaMoE CuteDSL (4 nodes x 4 GPU). +# +# Identical to the CUTLASS GPQA config except for moe_config, so a difference +# is attributable to the MoE backend alone. Fills the one empty cell in the +# validation matrix: CUTLASS has both GSM8K (96.40) and GPQA-D (91.92); +# MegaMoE had only GSM8K (96.32), and GSM8K is the cheap gate while GPQA is +# the sensitive one -- parity on the former does not imply parity on the +# latter, least of all across backends with different kernels and scale +# handling. +# +# NOTE this run still carries the EXPERIMENTAL experts_per_token 13 -> 16 +# widening (see TODO-C). A good score here is evidence at one more shape, not +# kernel-side coverage. +# +# Differs from the GSM8K DEP16 config in exactly one thing that matters: +# max_seq_len. GSM8K answers are a few hundred tokens, so 8192 covered prompt +# plus generation. GPQA-Diamond is a reasoning benchmark and the published +# Kimi-K3 numbers were measured with a 65536-token generation budget, so the +# sequence budget has to hold 4096 of prompt plus all of that. +# +# max_batch_size is 8, revised up from 4 after the 8-question smoke (job +# 475469) measured what this model actually generates here: 3431 and 6000 +# tokens, not the 65536 the budget allows. The budget still has to cover the +# worst case, but sizing CONCURRENCY for it was wrong -- at ~6k tokens and +# ~63 KiB/token a sequence wants ~0.4 GiB, not ~4 GiB, so the pool funds +# several. The V2 Mamba minimum also stays comfortable: ~0.42 GiB per resident +# slot means ~4.2 GiB at batch 8, well under the ~15 GiB quota, and the GSM8K +# DEP16 run already ran at 8. If several questions do run long the scheduler +# simply admits fewer of them. +tensor_parallel_size: 16 +enable_attention_dp: true +moe_expert_parallel_size: 16 +max_batch_size: 8 +max_num_tokens: 8192 +max_seq_len: 69632 +trust_remote_code: true +disable_overlap_scheduler: false +enable_chunked_prefill: true +cuda_graph_config: + enable_padding: true + max_batch_size: 8 +moe_config: + backend: MEGAMOE_CUTEDSL + # max_num_tokens deliberately absent: for MEGAMOE_* it is the SymmBuffer + # capacity (divided again by ep_size under attention-DP), not the per-call + # chunking bound it is for CUTLASS. Carrying 131072 over would over-provision + # the generation buffer ~32x and eat the KV cache. + use_low_precision_moe_combine: true +kv_cache_config: + enable_block_reuse: false + free_gpu_memory_fraction: 0.35 + tokens_per_block: 64 diff --git a/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py b/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py index 60a2b18e0397..d9a661ed0f61 100644 --- a/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py +++ b/tensorrt_llm/_torch/custom_ops/cute_dsl_megamoe_custom_op.py @@ -1012,6 +1012,8 @@ def query_megamoe_shared_workspace_bytes( tactic: Optional[Tuple] = None, apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, in_kernel_fc2_reduce: bool = False, combine_format: str = "bf16", ) -> int: @@ -1069,6 +1071,8 @@ def query_megamoe_shared_workspace_bytes( epi_flag_batch=tuple(epi_flag_batch), apply_topk_in_fc1=bool(apply_topk_in_fc1), gate_up_clamp=(None if gate_up_clamp is None else float(gate_up_clamp)), + situ_beta=(None if situ_beta is None else float(situ_beta)), + situ_linear_beta=(None if situ_linear_beta is None else float(situ_linear_beta)), **_LOCKED_KERNEL_KWARGS, ) # The probe MUST build the SAME kernel that runs (same combine_format): @@ -1115,6 +1119,8 @@ def __init__( output_dtype: torch.dtype, apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, in_kernel_fc2_reduce: bool = False, combine_format: str = "bf16", tactic_autotune: bool = False, @@ -1151,6 +1157,10 @@ def __init__( # per-call runtime kwargs. self.apply_topk_in_fc1 = bool(apply_topk_in_fc1) self.gate_up_clamp = None if gate_up_clamp is None else float(gate_up_clamp) + # SiTU constants are baked into the generated kernel exactly like + # gate_up_clamp, so they are part of unique_id() below. + self.situ_beta = None if situ_beta is None else float(situ_beta) + self.situ_linear_beta = None if situ_linear_beta is None else float(situ_linear_beta) # Symmetric profiling scratch, set by the op around ``choose_one`` # so the pre-hook routes cross-rank inputs through symmetric # memory; None outside tuning and for single-rank. @@ -1183,6 +1193,8 @@ def unique_id(self): str(self.output_dtype), self.apply_topk_in_fc1, self.gate_up_clamp, + self.situ_beta, + self.situ_linear_beta, self.in_kernel_fc2_reduce, self.combine_format, ) @@ -1375,6 +1387,8 @@ def _build_kernel(self, tactic: Tuple): epi_flag_batch=tuple(epi_flag_batch), apply_topk_in_fc1=self.apply_topk_in_fc1, gate_up_clamp=self.gate_up_clamp, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, **_LOCKED_KERNEL_KWARGS, ) kernel_cls, CombineFormat = _import_megamoe_kernel() @@ -1686,6 +1700,8 @@ def cute_dsl_megamoe_nvfp4_blackwell( peer_offsets: List[int], apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, in_kernel_fc2_reduce: bool = False, combine_format: str = "bf16", tactic_autotune: bool = False, @@ -1740,6 +1756,8 @@ def cute_dsl_megamoe_nvfp4_blackwell( output_dtype=combine_output.dtype, apply_topk_in_fc1=apply_topk_in_fc1, gate_up_clamp=gate_up_clamp, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, in_kernel_fc2_reduce=in_kernel_fc2_reduce, combine_format=combine_format, tactic_autotune=tactic_autotune, @@ -1852,6 +1870,8 @@ def _( peer_offsets: List[int], apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, in_kernel_fc2_reduce: bool = False, combine_format: str = "bf16", tactic_autotune: bool = False, diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py index 290d7b861f43..28d960eff0b7 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/epilogue_refactor.py @@ -1016,6 +1016,14 @@ def __init__( static_expert_shape: Optional[Tuple[ int, int, int]] = None, # [expert, intermediate, hidden] gate_up_clamp: Optional[float] = None, # Swiglu style only + # SiTU (Kimi K3). Both None -> SwiGLU. Both set -> SiTU, and + # ``gate_up_clamp`` must be None (SiTU has no clamp, matching + # DeepGEMM's ``DG_HOST_ASSERT(not use_situ or not + # activation_clamp_opt.has_value())``). Baked in as codegen-time + # constants exactly like ``gate_up_clamp``, so they must also be + # part of the compiled-kernel cache key. + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, epi_flag_batch: Optional[Tuple[int, int]] = ( 1, 1), # (fc1, fc2) done-counter publish batch ) -> None: @@ -1042,6 +1050,22 @@ def __init__( self.sf_vec_size = sf_vec_size # Swiglu gate/up clamp limit; None disables clamping. self.gate_up_clamp = gate_up_clamp + # SiTU activation constants; None/None selects SwiGLU. + if (situ_beta is None) != (situ_linear_beta is None): + raise ValueError( + "situ_beta and situ_linear_beta must be set together, got " + f"{situ_beta} and {situ_linear_beta}.") + if situ_beta is not None: + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError("SiTU beta parameters must be positive, got " + f"{situ_beta} and {situ_linear_beta}.") + if gate_up_clamp is not None: + raise ValueError( + "SiTU does not support gate_up_clamp (matches DeepGEMM, " + "which rejects activation_clamp together with SiTU).") + self.situ_beta = None if situ_beta is None else float(situ_beta) + self.situ_linear_beta = (None if situ_linear_beta is None else + float(situ_linear_beta)) # Done-counter publish batch granularity _fc1_eb, _fc2_eb = (1, 1) if epi_flag_batch is None else epi_flag_batch self.fc1_epi_flag_batch = max(1, min(32, int(_fc1_eb))) @@ -1732,20 +1756,61 @@ def alpha_swiglu_clamp( asm_dialect=llvm.AsmDialect.AD_ATT, )) - # 3) swiglu on the dequanted (and clamped) real values: - # out = up * gate * sigmoid(gate) - ug = cute.arch.mul_packed_f32x2((u0, u1), (g0, g1)) - neg_g_log2e = cute.arch.mul_packed_f32x2((g0, g1), neg_log2e_pair) - exp_pair = ( - cute.math.exp2(neg_g_log2e[0], fastmath=True), - cute.math.exp2(neg_g_log2e[1], fastmath=True), - ) - one_plus_exp = cute.arch.add_packed_f32x2(exp_pair, one_pair) - sigmoid_pair = ( - cute.arch.rcp_approx(one_plus_exp[0]), - cute.arch.rcp_approx(one_plus_exp[1]), - ) - out_pair = cute.arch.mul_packed_f32x2(ug, sigmoid_pair) + # 3) gated activation on the dequanted (and clamped) real values. + # sigmoid(x) = rcp(1 + exp2(-x * log2e)) -- shared by both cores. + def _sigmoid(p0, p1): + neg = cute.arch.mul_packed_f32x2((p0, p1), neg_log2e_pair) + e = ( + cute.math.exp2(neg[0], fastmath=True), + cute.math.exp2(neg[1], fastmath=True), + ) + d = cute.arch.add_packed_f32x2(e, one_pair) + return (cute.arch.rcp_approx(d[0]), cute.arch.rcp_approx(d[1])) + + sigmoid_pair = _sigmoid(g0, g1) + + if cutlass.const_expr(self.situ_beta is None): + # SwiGLU: out = up * gate * sigmoid(gate) + ug = cute.arch.mul_packed_f32x2((u0, u1), (g0, g1)) + out_pair = cute.arch.mul_packed_f32x2(ug, sigmoid_pair) + else: + # SiTU (Kimi K3), matching ``kimi_k3_moe/_mlp.py::SituAndMul`` + # (itself byte-identical to HF ``modeling_kimi.py``): + # situ_gate = beta * tanh(gate / beta) * sigmoid(gate) + # situ_up = linear_beta * tanh(up / linear_beta) + # out = situ_gate * situ_up + # + # ``tanh(z) = 2 * sigmoid(2z) - 1`` (same identity + # ``blackwell/utils.py::gelu_tanh_f32`` uses) keeps the whole + # core on the packed f32x2 path -- there is no packed tanh, so + # calling one would force this loop back to scalar. + # + # beta * tanh(x/beta) = beta * (2*sigmoid(2x/beta) - 1) + # = 2*beta*sigmoid(2x/beta) - beta + # so the reciprocals and the 2*beta factors fold at trace time. + inv_2beta = cutlass.Float32(2.0 / self.situ_beta) + two_beta = cutlass.Float32(2.0 * self.situ_beta) + neg_beta = cutlass.Float32(-self.situ_beta) + inv_2lbeta = cutlass.Float32(2.0 / self.situ_linear_beta) + two_lbeta = cutlass.Float32(2.0 * self.situ_linear_beta) + neg_lbeta = cutlass.Float32(-self.situ_linear_beta) + + gs = _sigmoid( + *cute.arch.mul_packed_f32x2((g0, g1), (inv_2beta, + inv_2beta))) + tanh_g = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(gs, (two_beta, two_beta)), + (neg_beta, neg_beta)) + + us = _sigmoid( + *cute.arch.mul_packed_f32x2((u0, u1), (inv_2lbeta, + inv_2lbeta))) + tanh_u = cute.arch.add_packed_f32x2( + cute.arch.mul_packed_f32x2(us, (two_lbeta, two_lbeta)), + (neg_lbeta, neg_lbeta)) + + situ_gate = cute.arch.mul_packed_f32x2(tanh_g, sigmoid_pair) + out_pair = cute.arch.mul_packed_f32x2(situ_gate, tanh_u) out[i] = out_pair[0] out[i + 1] = out_pair[1] diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py index ef06de0d1f03..8dbe12b49504 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/kernel_fc12.py @@ -83,6 +83,8 @@ def __init__( token_back_by_dispatch: bool = False, apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), ) -> None: if not force_static_sched: @@ -127,6 +129,8 @@ def __init__( self.token_back_by_dispatch = token_back_by_dispatch self.apply_topk_in_fc1 = apply_topk_in_fc1 self.gate_up_clamp = gate_up_clamp + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta self.epi_flag_batch = epi_flag_batch self._validate_mma_tiler_and_cluster_shape() @@ -198,8 +202,9 @@ def name(self) -> str: f"_padding_{self.token_padding_block}x{self.sf_padding_block}" f"_{fc2store}_{inkred}_{apply_topk}" f"_fc2out{self.fc2_output_dtype.__name__}_sfvec{self.sf_vec_size}" - f"_acc{self.acc_dtype.__name__}_clamp{self.gate_up_clamp}_epiflag{epiflag}" - ) + f"_acc{self.acc_dtype.__name__}_clamp{self.gate_up_clamp}" + # situ constants are baked into the kernel, so they belong in the name + f"_situ{self.situ_beta}x{self.situ_linear_beta}_epiflag{epiflag}") def _validate_mma_tiler_and_cluster_shape(self) -> None: """Validate user-provided geometry against v1 fused-fc12 constraints. @@ -368,6 +373,8 @@ def _setup_attributes(self) -> None: allow_overlap_acc=True, static_expert_shape=self.static_expert_shape, gate_up_clamp=self.gate_up_clamp, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, ) if self.num_sched_stages is None: diff --git a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py index 0ea3f48ae9ac..b30188050ac8 100644 --- a/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py +++ b/tensorrt_llm/_torch/cute_dsl_kernels/mega_moe_nvfp4/megamoe_kernel.py @@ -183,6 +183,8 @@ def __init__( "reuse_dispatch_warps"] = "epi_warps", apply_topk_in_fc1: bool = True, gate_up_clamp: Optional[float] = None, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, epi_flag_batch: Optional[Tuple[int, int]] = (1, 1), flag_batch: int = 1, ) -> None: @@ -255,6 +257,8 @@ def __init__( token_back_by_dispatch=token_back_by_dispatch, apply_topk_in_fc1=apply_topk_in_fc1, gate_up_clamp=gate_up_clamp, + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, epi_flag_batch=epi_flag_batch, ) @@ -889,7 +893,8 @@ def name(self) -> str: f"_padding_{self.token_padding_block}x{self.sf_padding_block}" f"_{fc2store}_{inkred}_token_back_by_{token_back}_{apply_topk}" f"_fc2out{self.fc2_output_dtype.__name__}_combine{self.combine_format}_sfvec{self.sf_vec_size}" - f"_acc{self.acc_dtype.__name__}_clamp{self.gate_up_clamp}_epiflag{epiflag}" + f"_acc{self.acc_dtype.__name__}_clamp{self.gate_up_clamp}" + f"_situ{self.situ_beta}x{self.situ_linear_beta}_epiflag{epiflag}" # MegaMoE-specific constexpr: f"_ep_{self.world_size}_topk_{self.num_topk}_maxtoken_{self.max_tokens_per_rank}" f"_flagbatch_{self.flag_batch}") diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 6b138bf40695..d827957e0e6d 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -1100,6 +1100,20 @@ def __init__( "Kimi K3 explicitly requested MEGAMOE_DEEPGEMM, but the " f"MoE factory selected {type(self.routed_experts.backend).__name__}." ) + if routed_moe_model_config.moe_backend == "MEGAMOE_CUTEDSL": + from ..modules.fused_moe.mega_moe import MegaMoECuteDsl + + # Same guard as MEGAMOE_DEEPGEMM above, and for the same reason: + # create_moe silently falls back when a backend declines the + # config, and for MegaMoE the decline is easy to trigger (it is + # EP-only and has its own token/top-k limits), so an explicit + # request that quietly became CUTLASS would be measured as if it + # were MegaMoE. + if not isinstance(self.routed_experts.backend, MegaMoECuteDsl): + raise RuntimeError( + "Kimi K3 explicitly requested MEGAMOE_CUTEDSL, 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 " @@ -1246,11 +1260,15 @@ def _resolve_routed_quant_config(model_config: ModelConfig, layer_idx: int) -> Q 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"} + supported_backends = { + "TRTLLM", + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + } if model_config.moe_backend not in supported_backends: raise ValueError( - "Kimi K3 SiTU routed experts only support the TRTLLM and " - "MEGAMOE_DEEPGEMM backends; " + "Kimi K3 SiTU routed experts only support the TRTLLM, " + "MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " f"got {model_config.moe_backend!r}." ) if model_config.moe_load_balancer is not None: @@ -1295,7 +1313,10 @@ def _routed_moe_model_config(model_config: ModelConfig) -> ModelConfig: # 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": + if routed_model_config.moe_backend in { + "MEGAMOE_DEEPGEMM", + "MEGAMOE_CUTEDSL", + }: default_moe_max_num_tokens = routed_model_config.max_num_tokens * routed_mapping.dp_size configured_moe_max_num_tokens = int(routed_model_config.moe_max_num_tokens or 0) if configured_moe_max_num_tokens < default_moe_max_num_tokens: diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py index 6a65635867e9..7f1ffe2d4ec7 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -467,6 +467,11 @@ def __init__( init_load_balancer: bool = True, activation_type: ActivationType = ActivationType.Swiglu, swiglu_limit: Optional[torch.Tensor] = None, + # ``activation=None`` infers Kimi K3 SiTU from the pretrained config and + # otherwise defaults to SwiGLU. Mirrors MegaMoEDeepGemm. + activation: Optional[str] = None, + situ_beta: Optional[float] = None, + situ_linear_beta: Optional[float] = None, **kwargs, ) -> None: # ``aux_stream_dict`` is accepted for ``create_moe_backend`` signature @@ -496,6 +501,22 @@ def __init__( "MegaMoECuteDsl does not support apply_router_weight_on_input; " "the fused kernel applies routing weights on the MoE output." ) + # ``ActivationType.Swiglu`` describes the gated FC1 tensor geometry shared + # by SwiGLU and SiTU; the elementwise function is selected by + # ``activation`` below. Same reasoning as MegaMoEDeepGemm. + if activation_type != ActivationType.Swiglu: + raise ValueError( + f"MegaMoECuteDsl 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, + ) + self.activation = activation + self.situ_beta = situ_beta + self.situ_linear_beta = situ_linear_beta self.apply_router_weight_on_input = apply_router_weight_on_input # topk-score application point. v2 default is the deepgemm graph @@ -529,6 +550,11 @@ def __init__( # used directly (NO trtllm-gen-style div_(fc31_alpha) normalization). # Reject non-uniform / per-expert clamp: the kernel bakes one constant. self.gate_up_clamp = self._resolve_gate_up_clamp(swiglu_limit) + if self.activation == "situ" and self.gate_up_clamp is not None: + raise ValueError( + "MegaMoECuteDsl SiTU does not support a gate/up clamp; " + "drop swiglu_limit for SiTU checkpoints." + ) # Buffer sizing. MoE layers execute serially per forward; one pool # sized to the worst-case per-rank tokens covers every layer. The @@ -608,6 +634,52 @@ def __init__( # ------------------------------------------------------------------ # Topology # ------------------------------------------------------------------ + @staticmethod + def _resolve_activation_config( + model_config, + *, + activation: Optional[str], + situ_beta: Optional[float], + situ_linear_beta: Optional[float], + ): + """Resolve the elementwise activation and its SiTU constants. + + ``activation=None`` infers SiTU from ``activation_situ_beta`` in the + pretrained config (Kimi K3 sets it) and otherwise selects SwiGLU. Kept + byte-for-byte equivalent to ``MegaMoEDeepGemm._resolve_activation_config`` + so the two MegaMoE backends cannot disagree about the same checkpoint. + """ + pretrained_config = getattr(model_config, "pretrained_config", None) + text_config = getattr(pretrained_config, "text_config", None) + cfg_beta = getattr(pretrained_config, "activation_situ_beta", None) + cfg_lbeta = getattr(pretrained_config, "activation_situ_linear_beta", None) + if cfg_beta is None: + cfg_beta = getattr(text_config, "activation_situ_beta", None) + if cfg_lbeta is None: + cfg_lbeta = getattr(text_config, "activation_situ_linear_beta", None) + if activation is None: + activation = "situ" if cfg_beta is not None else "swiglu" + activation = activation.lower() + if activation not in ("swiglu", "situ"): + raise ValueError( + f"MegaMoECuteDsl 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 = cfg_beta if situ_beta is None else situ_beta + situ_linear_beta = cfg_lbeta if situ_linear_beta is None else situ_linear_beta + if situ_beta is None or situ_linear_beta is None: + raise ValueError( + "MegaMoECuteDsl 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("MegaMoECuteDsl SiTU beta parameters must be positive.") + return activation, float(situ_beta), float(situ_linear_beta) + @staticmethod def _resolve_gate_up_clamp( swiglu_limit: Optional[torch.Tensor], @@ -719,10 +791,21 @@ def validate_configurable_moe(self, moe) -> None: f"of the fused kernel); got moe.comm={type(moe.comm).__name__}." ) top_k = moe.routing_method.experts_per_token - if top_k > 13: + # EXPERIMENTAL, pending kernel-side confirmation: raised 13 -> 16 to + # admit Kimi K3, whose routing is top-16. The original bound was + # described as matching *external coverage*, i.e. what had been + # validated, not a kernel capability -- and nothing in the CuTe DSL + # kernel hardcodes it: num_topk is a sizing parameter there + # (max_tokens_per_rank * num_topk, (max_tokens, num_topk, hidden)). + # That makes the widening plausible, NOT proven: shared-memory and + # register budgets grow with num_topk, and topk_reduce.py may carry + # fixed unrolls. Accuracy at one shape is evidence, not coverage. + # See TODO-C in KIMI_K3_NVFP4_BRINGUP.md before relying on this. + _MEGAMOE_CUTEDSL_MAX_TOP_K = 16 + if top_k > _MEGAMOE_CUTEDSL_MAX_TOP_K: raise ValueError( - f"MegaMoECuteDsl supports experts_per_token <= 13 " - f"(matches external coverage); got {top_k}." + f"MegaMoECuteDsl supports experts_per_token <= " + f"{_MEGAMOE_CUTEDSL_MAX_TOP_K}; got {top_k}." ) if moe.moe_max_num_tokens <= 0: raise ValueError( @@ -1442,6 +1525,8 @@ def _deferred_scratch_factory(_self=self, _get=get_megamoe_profiling_scratch): peer_offsets=peer_offsets, apply_topk_in_fc1=bool(self.apply_topk_in_fc1), gate_up_clamp=self.gate_up_clamp, + situ_beta=self.situ_beta, + situ_linear_beta=self.situ_linear_beta, # Keep the deterministic standalone TopkReduce until form-B # has dedicated GPU correctness and performance coverage. in_kernel_fc2_reduce=False, diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 15a9383c44a6..ecbdd5c59031 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -3786,11 +3786,28 @@ def prepare_streaming_expert_load(self, module: torch.nn.Module) -> None: # those containers is created lazily on the load path, which a # multi-threaded streaming loader cannot do safely. super().prepare_streaming_expert_load(module) + # This backend keeps its raw NVFP4 source params as 0-element + # placeholders and only rematerializes them inside its own + # load_weights, which a streaming loader bypasses -- so without this + # the first per-expert write indexes an empty tensor + # ("index 0 is out of bounds for dimension 0 with size 0"). + # process_weights_after_loading shrinks them again after packing. + self._materialize_source_params(module) module.tmp_cutlass_w3_w1_weights = {} module.tmp_cutlass_w3_w1_weight_scales = {} module._streamed_w2_covered = set() module._streamed_w2_scale_covered = set() + # NOTE: this backend deliberately does NOT override + # finalize_streamed_expert, unlike its Cutlass sibling. Its staged + # dicts are not only staging: _initial_slot_coverage() COUNTS their + # entries to decide which slots a partial load actually populated, so + # draining them per expert would report zero coverage and make a + # complete load look empty. Bounding the staged footprint here means + # moving that accounting off the dicts first (e.g. onto + # _streamed_expert_slots, which already tracks exactly this) rather + # than copying the Cutlass drain over. + def _get_fc2_alpha_input_scale( self, module: torch.nn.Module, @@ -4220,7 +4237,24 @@ def _check_initial_aux_scale_coverage(self, module: torch.nn.Module) -> None: """Reject partially populated NVFP4 auxiliary-scale families.""" n_slots = module.expert_size_per_partition - n_experts = module.num_experts + # A whole-checkpoint load is handed every expert's input_scale, because + # the weights dict holds the entire checkpoint; a streaming EP load + # only ever reads its own rank's experts, so its complete answer is + # expert_size_per_partition, not num_experts. Expecting the latter + # turns a complete streamed load into a false "partial" report. + # + # This is the one place the documented divergence in + # load_streaming_nvfp4_expert becomes visible: with fewer entries the + # global input-scale reduction in process_weights_after_loading is over + # the rank-local slice rather than all experts. The two agree exactly + # while input_scale is uniform across experts -- which is what a + # static-activation-scale checkpoint gives. A checkpoint with genuinely + # per-expert activation scales needs a cross-rank reduction added + # there; this check is not the place to catch that, because a + # single-rank load of such a checkpoint would be equally wrong and + # equally "complete". + streamed = bool(getattr(module, '_streamed_expert_slots', None)) + n_experts = n_slots if streamed else module.num_experts weight_scale_2 = getattr(module, 'tmp_weight_scale_2', None) or {} raw_input_scales = getattr(module, 'tmp_raw_input_scales', None) or {} 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 65c6eaa83187..f1c184ca3bff 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 @@ -469,7 +469,7 @@ def test_kimi_k3_moe_split_selection(monkeypatch): assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (4, 2) -@pytest.mark.parametrize("backend", ["TRTLLM", "MEGAMOE_DEEPGEMM"]) +@pytest.mark.parametrize("backend", ["TRTLLM", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"]) def test_kimi_k3_routed_config_preserves_explicit_backend(backend): model_config = ModelConfig( mapping=Mapping(world_size=1, rank=0, tp_size=1), @@ -486,7 +486,8 @@ def test_kimi_k3_routed_config_preserves_explicit_backend(backend): "backend,expected_moe_max_num_tokens", [ pytest.param("TRTLLM", 33024, id="trtllm"), - pytest.param("MEGAMOE_DEEPGEMM", 131072, id="megamoe"), + pytest.param("MEGAMOE_DEEPGEMM", 131072, id="megamoe-deepgemm"), + pytest.param("MEGAMOE_CUTEDSL", 131072, id="megamoe-cutedsl"), ], ) def test_kimi_k3_routed_config_scopes_megamoe_capacity(backend, expected_moe_max_num_tokens): From d622b63cb85a2fa320921056d7a51a2ea4de5360 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:37:49 -0700 Subject: [PATCH 5/9] [None][fix] Enable Kimi K3 SiTU on CUTLASS Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_kimi_linear.py | 3 ++- tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py | 6 +++--- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index d827957e0e6d..9f5604545bd5 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -1261,13 +1261,14 @@ 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 = { + "CUTLASS", "TRTLLM", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL", } if model_config.moe_backend not in supported_backends: raise ValueError( - "Kimi K3 SiTU routed experts only support the TRTLLM, " + "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " "MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " f"got {model_config.moe_backend!r}." ) 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 f1c184ca3bff..3370ae3f2269 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 @@ -469,7 +469,7 @@ def test_kimi_k3_moe_split_selection(monkeypatch): assert KimiK3MoERuntime._select_moe_tp_ep(auto) == (4, 2) -@pytest.mark.parametrize("backend", ["TRTLLM", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"]) +@pytest.mark.parametrize("backend", ["CUTLASS", "TRTLLM", "MEGAMOE_DEEPGEMM", "MEGAMOE_CUTEDSL"]) def test_kimi_k3_routed_config_preserves_explicit_backend(backend): model_config = ModelConfig( mapping=Mapping(world_size=1, rank=0, tp_size=1), @@ -539,13 +539,13 @@ def test_kimi_k3_routed_config_logs_megamoe_capacity_override(monkeypatch): 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", + moe_backend="TRITON", ) with pytest.raises(ValueError, match="SiTU routed experts only support"): KimiK3MoERuntime._routed_moe_model_config(model_config) - assert model_config.moe_backend == "CUTLASS" + assert model_config.moe_backend == "TRITON" @pytest.mark.parametrize( From e9d5e7c466f643c90c9a257f30d62906c153635c Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:12:30 -0700 Subject: [PATCH 6/9] Remove local doc reference from MegaMoE comment Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py index 7f1ffe2d4ec7..352b44788d2c 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py +++ b/tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_cute_dsl.py @@ -800,7 +800,6 @@ def validate_configurable_moe(self, moe) -> None: # That makes the widening plausible, NOT proven: shared-memory and # register budgets grow with num_topk, and topk_reduce.py may carry # fixed unrolls. Accuracy at one shape is evidence, not coverage. - # See TODO-C in KIMI_K3_NVFP4_BRINGUP.md before relying on this. _MEGAMOE_CUTEDSL_MAX_TOP_K = 16 if top_k > _MEGAMOE_CUTEDSL_MAX_TOP_K: raise ValueError( From 6b09ef052e2fb17d2838c566eff992a561c09ea1 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:22:45 -0700 Subject: [PATCH 7/9] Align Kimi FP8 checkpoint stash gate Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- tensorrt_llm/_torch/models/modeling_kimi_linear.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 9f5604545bd5..10c57dc8b684 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -2988,7 +2988,7 @@ def _load_trunk_params( model_tp_rank = self.model_config.mapping.tp_rank # Keep each FP8_PB_WO checkpoint pair alongside the BF16 # parameter only when the later weight-read conversion consumes it. - stash_ckpt_fp8 = os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_ENV, "0") != "0" and is_sm_100f() + stash_ckpt_fp8 = _resolve_fp8_weight_read_gates()[0] # KDA head-shard (attention-DP off): rank r loads head rows/cols # [r*local : (r+1)*local] of every head-major KDA tensor. kda_tp_size, kda_tp_rank = 1, 0 From 08abae0b50f94997242cace01a1b60d2a05bab48 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Tue, 18 Aug 2026 07:43:26 -0700 Subject: [PATCH 8/9] Address Kimi SiTu review feedback Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../cutlass_kernels/include/moe_kernels.h | 4 ++ .../cutlass_kernels/moe_gemm/moe_kernels.cu | 46 ++++++++++++--- cpp/tensorrt_llm/thop/moeOp.cpp | 10 ++++ .../_torch/models/modeling_kimi_linear.py | 56 +++++++++++++++---- .../modules/moe/test_kimi_k3_situ_moe.py | 30 ++++++++++ 5 files changed, 126 insertions(+), 20 deletions(-) diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h index 5053d7c571c2..f2356cc21d93 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -148,6 +148,10 @@ struct ActivationParams , swiglu_beta(swiglu_beta) , swiglu_limit(swiglu_limit) { + TLLM_CHECK_WITH_INFO(activation_type != ActivationType::SiTu || (swiglu_alpha && swiglu_beta), + "SiTu requires both alpha and beta activation parameters"); + TLLM_CHECK_WITH_INFO( + activation_type != ActivationType::SiTu || !swiglu_limit, "SiTu does not support a clamp limit"); } // TODO Port everything properly and get rid of these implicit conversions diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu index 09072b546f91..e32b7d461175 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu @@ -5037,6 +5037,18 @@ __global__ void populateRandomBufferKernel(void* buffer_void, size_t size) buffer[tid * elem_per_thread + i] = curand4(&state); } +__global__ void populateProfilerSiTuParamsKernel(float* alpha, float* beta, int const numExperts) +{ + int const tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= numExperts) + { + return; + } + + alpha[tid] = 1.0F; + beta[tid] = 1.0F; +} + template __global__ void prepareMinLatencyBuffer(int* num_active_experts_per_node, int* active_expert_global_ids, int64_t* expert_first_token_offset, int const num_tokens, int const num_experts_per_token, @@ -5299,10 +5311,12 @@ std::map> GemmProfilerBackend::getProfile = mMinLatencyMode ? sizeof(int) * NUM_ROUTING_SAMPLES : 0; // smaller than or equal to num_experts_per_node size_t active_expert_global_ids_size = mMinLatencyMode ? mNumExpertsPerNode * sizeof(int) * NUM_ROUTING_SAMPLES : 0; - bool is_swiglu_bias = mActivationType == ActivationType::SwigluBias && mGemmToProfile == GemmToProfile::GEMM_1; - size_t swiglu_alpha_size = is_swiglu_bias ? num_experts_per_node * sizeof(float) : 0; - size_t swiglu_beta_size = is_swiglu_bias ? num_experts_per_node * sizeof(float) : 0; - size_t swiglu_limit_size = is_swiglu_bias ? num_experts_per_node * sizeof(float) : 0; + bool const profilesGemm1 = mGemmToProfile == GemmToProfile::GEMM_1; + bool const isSwigluBias = mActivationType == ActivationType::SwigluBias && profilesGemm1; + bool const isSitu = mActivationType == ActivationType::SiTu && profilesGemm1; + size_t swiglu_alpha_size = (isSwigluBias || isSitu) ? num_experts_per_node * sizeof(float) : 0; + size_t swiglu_beta_size = (isSwigluBias || isSitu) ? num_experts_per_node * sizeof(float) : 0; + size_t swiglu_limit_size = isSwigluBias ? num_experts_per_node * sizeof(float) : 0; size_t map_offset = 0; std::map> out_map; @@ -5338,14 +5352,14 @@ std::map> GemmProfilerBackend::getProfile ADD(quant_4); ADD(quant_5); ADD(quant_6); - ADD(tma_ws_input_workspace); + ADD(swiglu_alpha); + ADD(swiglu_beta); + ADD(swiglu_limit); ADD(w4a8_alpha); + ADD(tma_ws_input_workspace); ADD(alpha_scale_ptr_array); ADD(fp4_act_scale_flat); ADD(gemm_workspace); - ADD(swiglu_alpha); - ADD(swiglu_beta); - ADD(swiglu_limit); #undef ADD_NAME #undef ADD @@ -5626,6 +5640,22 @@ void GemmProfilerBackend::prepare( auto workspace_size = getWorkspaceSize(num_tokens); populateRandomBuffer(workspace_ptr_char, workspace_size, stream); + if (mActivationType == ActivationType::SiTu && mGemmToProfile == GemmToProfile::GEMM_1) + { + auto const workspaces = getProfilerWorkspaces(num_tokens, mSM >= 90); + auto const& alphaWorkspace = workspaces.at("swiglu_alpha"); + auto const& betaWorkspace = workspaces.at("swiglu_beta"); + size_t const expectedSize = static_cast(mNumExpertsPerNode) * sizeof(float); + TLLM_CHECK_WITH_INFO(alphaWorkspace.first >= expectedSize && betaWorkspace.first >= expectedSize, + "SiTu profiler activation-parameter workspace has the wrong size"); + auto* alpha = reinterpret_cast(workspace_ptr_char + alphaWorkspace.second); + auto* beta = reinterpret_cast(workspace_ptr_char + betaWorkspace.second); + constexpr int kThreadsPerBlock = 128; + populateProfilerSiTuParamsKernel<<>>(alpha, beta, mNumExpertsPerNode); + sync_check_cuda_error(stream); + } + prepareRouting(num_tokens, workspace_ptr_char, stream); prepareQuantParams(num_tokens, workspace_ptr_char, stream); for (auto fusion : {TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::NONE, diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index e7f9b1ff97a5..cd41979bc72b 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -568,6 +568,11 @@ class FusedMoeRunner : public torch::CustomClassHolder base_activation_type = ActivationType::SwigluBias; } } + TORCH_CHECK( + base_activation_type != ActivationType::SiTu || (swiglu_alpha.has_value() && swiglu_beta.has_value()), + "SiTu requires both swiglu_alpha and swiglu_beta."); + TORCH_CHECK(base_activation_type != ActivationType::SiTu || !swiglu_limit.has_value(), + "SiTu does not support swiglu_limit."); auto activation_params = ActivationParams(base_activation_type, reinterpret_cast(swiglu_alpha.has_value() ? swiglu_alpha.value().const_data_ptr() : nullptr), reinterpret_cast(swiglu_beta.has_value() ? swiglu_beta.value().const_data_ptr() : nullptr), @@ -831,6 +836,11 @@ class FusedMoeRunner : public torch::CustomClassHolder base_activation_type = ActivationType::SwigluBias; } } + TORCH_CHECK( + base_activation_type != ActivationType::SiTu || (swiglu_alpha.has_value() && swiglu_beta.has_value()), + "SiTu requires both swiglu_alpha and swiglu_beta."); + TORCH_CHECK(base_activation_type != ActivationType::SiTu || !swiglu_limit.has_value(), + "SiTu does not support swiglu_limit."); auto activation_params = ActivationParams(base_activation_type, reinterpret_cast(swiglu_alpha.has_value() ? swiglu_alpha.value().const_data_ptr() : nullptr), reinterpret_cast(swiglu_beta.has_value() ? swiglu_beta.value().const_data_ptr() : nullptr), diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 10c57dc8b684..44c498b2a74d 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -318,6 +318,23 @@ def _resolve_fp8_weight_read_gates() -> tuple[bool, bool, bool]: return fp8_weight_read, kda_fp8, kda_glue_fp8 +def _resolve_kimi_situ_betas(cfg: Any) -> tuple[float, float]: + """Return the finite SiTu betas required by the routed-expert kernels.""" + config_situ_beta = getattr(cfg, "activation_situ_beta", None) + situ_beta = 1.0 if config_situ_beta is None else config_situ_beta + situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + if situ_linear_beta is None: + raise ValueError( + "Kimi K3 routed SiTu experts require activation_situ_linear_beta; " + "None means an identity linear branch that the fused kernels cannot represent." + ) + if situ_beta <= 0 or situ_linear_beta <= 0: + raise ValueError( + f"Kimi K3 SiTu betas must be positive; got {situ_beta} and {situ_linear_beta}." + ) + return float(situ_beta), float(situ_linear_beta) + + # --------------------------------------------------------------------------- # Config helpers. # --------------------------------------------------------------------------- @@ -558,7 +575,9 @@ def from_linear(cls, linear: nn.Linear | TrtllmLinear) -> "_Fp8BlockScaleWeightR # the scale the checkpoint was written with. ckpt_pair = getattr(linear.weight, _K3_CKPT_FP8_ATTR, None) if ckpt_pair is not None: - return cls.from_checkpoint_fp8(ckpt_pair[0], ckpt_pair[1], linear.out_features) + converted = cls.from_checkpoint_fp8(ckpt_pair[0], ckpt_pair[1], linear.out_features) + delattr(linear.weight, _K3_CKPT_FP8_ATTR) + return converted weight_fp8, weight_scale = cls.quantize_weight(linear.weight.data) return cls(weight_fp8, weight_scale, linear.out_features) @@ -994,8 +1013,7 @@ def __init__( if not getattr(cfg, "latent_moe_use_norm", False): raise ValueError("Kimi K3 runtime expects latent_moe_use_norm=True") - situ_beta = getattr(cfg, "activation_situ_beta", None) or 1.0 - situ_linear_beta = getattr(cfg, "activation_situ_linear_beta", None) + situ_beta, situ_linear_beta = _resolve_kimi_situ_betas(cfg) dtype = torch.bfloat16 # Routing scores stay fp32; with attention-DP off the gate GEMM runs @@ -1062,7 +1080,7 @@ def __init__( ) self.routed_situ_beta = torch.full( (local_num_experts,), - float(situ_linear_beta if situ_linear_beta is not None else 1.0), + situ_linear_beta, dtype=torch.float32, device=device, ) @@ -1076,16 +1094,14 @@ def __init__( 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 - ), + trtllm_gen_activation_alpha=situ_beta, + trtllm_gen_activation_beta=situ_linear_beta, ) 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), + situ_beta=situ_beta, + situ_linear_beta=situ_linear_beta, ) self.routed_experts = create_moe(**routed_moe_kwargs) if not isinstance(self.routed_experts, ConfigurableMoE): @@ -2756,14 +2772,24 @@ def _materialize(value) -> torch.Tensor: if isinstance(value, torch.Tensor): return value # ``[:]`` is how a lazy slice is realized, but it is invalid on a 0-dim - # entry — and the NVFP4 checkpoint stores weight_scale_2 / input_scale as - # scalars. Same fix as ``_ReopenSafeTensorSlice._tensor``. + # entry. The NVFP4 checkpoint stores weight_scale_2 / input_scale as + # scalars, so realize those entries with ``[()]`` instead. get_shape = getattr(value, "get_shape", None) if get_shape is not None and len(get_shape()) == 0: return value[()] return value[:] +def _clear_checkpoint_fp8_pairs(module: nn.Module) -> int: + """Release checkpoint FP8 pairs left on parameters after conversion.""" + cleared = 0 + for param in module.parameters(): + if hasattr(param, _K3_CKPT_FP8_ATTR): + delattr(param, _K3_CKPT_FP8_ATTR) + cleared += 1 + return cleared + + @register_auto_model("KimiLinearForCausalLM") class KimiLinearForCausalLM(SpecDecOneEngineForCausalLM[KimiLinearModel, Any]): """Kimi K3 text core (KDA + MLA + MoE). @@ -3482,3 +3508,9 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: logger.info( f"Kimi K3: reading {n_mla} MLA q_a/q_b/o/g projections at FP8 block-scale" ) + + # Sub-switches may leave checkpoint pairs on BF16 parameters that were + # deliberately not converted. No later stage consumes them. + cleared_pairs = _clear_checkpoint_fp8_pairs(self.model) + if cleared_pairs: + logger.debug(f"Kimi K3: released {cleared_pairs} unconsumed checkpoint FP8 pairs") 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 3370ae3f2269..8b564563a92e 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 @@ -99,6 +99,35 @@ class _K3Config: activation_situ_linear_beta: float = 25.0 +def test_kimi_situ_betas_require_linear_beta(): + cfg = SimpleNamespace(activation_situ_beta=4.0, activation_situ_linear_beta=None) + + with pytest.raises(ValueError, match="require activation_situ_linear_beta"): + modeling_kimi_linear._resolve_kimi_situ_betas(cfg) + + +@pytest.mark.parametrize( + "situ_beta,situ_linear_beta", + [(0.0, 25.0), (4.0, 0.0), (-1.0, 25.0), (4.0, -1.0)], +) +def test_kimi_situ_betas_must_be_positive(situ_beta, situ_linear_beta): + cfg = SimpleNamespace( + activation_situ_beta=situ_beta, + activation_situ_linear_beta=situ_linear_beta, + ) + + with pytest.raises(ValueError, match="must be positive"): + modeling_kimi_linear._resolve_kimi_situ_betas(cfg) + + +def test_clear_checkpoint_fp8_pairs_releases_unconsumed_stashes(): + linear = torch.nn.Linear(2, 2, bias=False) + setattr(linear.weight, modeling_kimi_linear._K3_CKPT_FP8_ATTR, (torch.ones(1), torch.ones(1))) + + assert modeling_kimi_linear._clear_checkpoint_fp8_pairs(linear) == 1 + assert not hasattr(linear.weight, modeling_kimi_linear._K3_CKPT_FP8_ATTR) + + def _init_block_weights(block: KimiK3SparseMoeBlock, seed: int = 1234): """Fill gate/latent/shared weights and the MXFP4 expert bank.""" gen = torch.Generator(device="cpu").manual_seed(seed) @@ -1604,6 +1633,7 @@ def test_fp8_weight_read_prefers_the_checkpoint_pair(): setattr(linear.weight, _K3_CKPT_FP8_ATTR, (ckpt_fp8, ckpt_scale.float())) converted = FP8Linear.from_linear(linear) + assert not hasattr(linear.weight, _K3_CKPT_FP8_ATTR) expected_w, expected_s = FP8Linear.prepare_checkpoint_scale(ckpt_fp8, ckpt_scale.float()) assert _bitwise_equal(converted.weight, expected_w) assert _bitwise_equal(converted.weight_scale, expected_s) From d45ca60c589e209fa2bbf0014e44515963c310d4 Mon Sep 17 00:00:00 2001 From: Jonas Li <6110159+longlee0622@users.noreply.github.com> Date: Thu, 20 Aug 2026 00:07:39 -0700 Subject: [PATCH 9/9] [None][fix] Kimi K3: size CUTLASS SiTU alpha/beta with the shared EP partition Address review feedback on #17865. - Use `_compute_ep_partition` instead of `num_experts // ep_size` so the per-expert SiTU constants match `expert_size_per_partition` on uneven expert/EP splits, where moeOp.cpp's `swiglu_alpha must have num_experts_on_rank elements` check would otherwise fire. - Document why MEGAMOE_CUTEDSL intentionally has no SiTU kwarg branch. Signed-off-by: Jonas Li <6110159+longlee0622@users.noreply.github.com> --- .../_torch/models/modeling_kimi_linear.py | 21 ++++++++++++++++++- 1 file changed, 20 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 44c498b2a74d..931b09bdc19e 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -112,6 +112,7 @@ from ..distributed import AllReduce, AllReduceParams, AllReduceStrategy from ..model_config import ModelConfig from ..modules.fused_moe import ConfigurableMoE, create_moe +from ..modules.fused_moe.interface import _compute_ep_partition from ..modules.fused_moe.routing import DeepSeekV3MoeRoutingMethod from ..modules.gated_mlp import GatedMLP from ..modules.linear import Linear as TrtllmLinear @@ -1073,7 +1074,18 @@ def __init__( ) if routed_moe_model_config.moe_backend == "CUTLASS": - local_num_experts = self.num_experts // routed_moe_model_config.mapping.moe_ep_size + # Size the per-expert SiTU constants with the same ceil/floor + # partition the MoE backend uses for ``expert_size_per_partition``. + # A plain ``num_experts // ep_size`` is one element short on the + # first ``num_experts % ep_size`` ranks, which trips the + # ``swiglu_alpha must have num_experts_on_rank elements`` check in + # moeOp.cpp. K3's 384 experts divide evenly at EP8/EP16, so the + # mismatch is latent there but real for any uneven split. + local_num_experts, _, _ = _compute_ep_partition( + self.num_experts, + routed_moe_model_config.mapping.moe_ep_size, + routed_moe_model_config.mapping.moe_ep_rank, + ) device = torch.device("cuda", torch.cuda.current_device()) self.routed_situ_alpha = torch.full( (local_num_experts,), float(situ_beta), dtype=torch.float32, device=device @@ -1103,6 +1115,13 @@ def __init__( situ_beta=situ_beta, situ_linear_beta=situ_linear_beta, ) + # MEGAMOE_CUTEDSL has no branch on purpose. ``MegaMoECuteDsl`` resolves + # SiTU from the pretrained config in ``_resolve_activation_config`` + # (``activation=None`` -> "situ" when ``activation_situ_beta`` is + # present), and ``create_moe`` currently rejects the explicit + # ``activation``/``situ_beta``/``situ_linear_beta`` trio for anything + # other than ``MegaMoEDeepGemm``, so passing them here would raise. + # Unifying that plumbing is tracked in TRTLLM-15649. self.routed_experts = create_moe(**routed_moe_kwargs) if not isinstance(self.routed_experts, ConfigurableMoE): raise RuntimeError(