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>; @@ -5030,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, @@ -5292,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; @@ -5331,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 @@ -5619,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/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..cd41979bc72b 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -543,22 +543,36 @@ 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; + } } + 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), @@ -797,22 +811,36 @@ 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; + } } + 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/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_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/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/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/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/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/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 7f15e37fe0c7..931b09bdc19e 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 @@ -100,13 +112,14 @@ 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 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 @@ -306,6 +319,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. # --------------------------------------------------------------------------- @@ -500,12 +530,157 @@ 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: + 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) + @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 +689,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 +899,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.""" @@ -739,8 +1014,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 @@ -760,7 +1034,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,22 +1053,75 @@ 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": + # 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 + ) + self.routed_situ_beta = torch.full( + (local_num_experts,), + situ_linear_beta, + 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 # 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, ) + # 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( @@ -805,6 +1135,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 " @@ -913,15 +1257,54 @@ 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 config. Default split is EP-only; see ``_select_moe_tp_ep``.""" - supported_backends = {"TRTLLM", "MEGAMOE_DEEPGEMM"} + 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 and " - "MEGAMOE_DEEPGEMM backends; " + "Kimi K3 SiTU routed experts only support the CUTLASS, TRTLLM, " + "MEGAMOE_DEEPGEMM, and MEGAMOE_CUTEDSL backends; " f"got {model_config.moe_backend!r}." ) if model_config.moe_load_balancer is not None: @@ -966,7 +1349,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: @@ -2330,13 +2716,99 @@ 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. 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). @@ -2476,10 +2948,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 +3031,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 = _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 @@ -2574,8 +3050,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 +3076,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 +3192,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 +3217,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 +3313,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 +3337,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 +3363,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 +3387,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 +3397,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: @@ -2936,3 +3527,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/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..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 @@ -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,20 @@ 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. + _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 +1524,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 90828ccd6f5d..ecbdd5c59031 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,34 @@ 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) + # 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, @@ -4013,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/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 00bcd5c6d96d..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): @@ -88,7 +87,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/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/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() 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..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 @@ -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(), @@ -98,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) @@ -468,7 +498,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", ["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), @@ -485,7 +515,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): @@ -537,13 +568,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( @@ -623,6 +654,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 +682,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 +696,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 +913,836 @@ 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) + 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) + # 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/_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 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."""