[None][feat] bring up Kimi K3 NVFP4 with CUTLASS and cuteDSL MegaMoE SiTU

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

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 <typename T,                   /*The type used for activations/scales/compute*/

cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h

@@ -137,8 +137,8 @@ struct ActivationParams
     explicit ActivationParams(ActivationType activation_type)
         : activation_type(activation_type)
     {
-        TLLM_CHECK_WITH_INFO(activation_type != ActivationType::SwigluBias,
-            "SwigluBias is not supported in ActivationParams without swiglu_alpha and swiglu_beta");
+        TLLM_CHECK_WITH_INFO(activation_type != ActivationType::SwigluBias && activation_type != ActivationType::SiTu,
+            "SwigluBias and SiTu are not supported in ActivationParams without alpha and beta");
     }
 
     ActivationParams(
@@ -148,6 +148,10 @@ struct ActivationParams
         , swiglu_beta(swiglu_beta)
         , swiglu_limit(swiglu_limit)
     {
+        TLLM_CHECK_WITH_INFO(activation_type != ActivationType::SiTu || (swiglu_alpha && swiglu_beta),
+            "SiTu requires both alpha and beta activation parameters");
+        TLLM_CHECK_WITH_INFO(
+            activation_type != ActivationType::SiTu || !swiglu_limit, "SiTu does not support a clamp limit");
     }
 
     // TODO Port everything properly and get rid of these implicit conversions

cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu

@@ -2284,6 +2284,8 @@ void doGatedActivation(ActivationOutputType* output, GemmOutputType const* gemm_
         ? &doGatedActivationKernel<ActivationOutputType, GemmOutputType, GLUAdaptor<cutlass::epilogue::thread::GELU>>
         : activation_type == ActivationType::SwigluBias
         ? &doGatedActivationKernel<ActivationOutputType, GemmOutputType, SwigluBiasAdaptor>
+        : activation_type == ActivationType::SiTu
+        ? &doGatedActivationKernel<ActivationOutputType, GemmOutputType, SiTuAdaptor>
         : nullptr;
     TLLM_CHECK_WITH_INFO(fn != nullptr, "Invalid activation type");
     fn<<<blocks, threads, 0, stream>>>(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<T, GemmOutputType, ScaleBiasType, SwigluBiasAdaptor,
                         decltype(block_scaling_type)::value, num_rows_per_cta_v, false, kWriteFp8>;
+                case ActivationType::SiTu:
+                    return &doActivationKernel<T, GemmOutputType, ScaleBiasType, SiTuAdaptor,
+                        decltype(block_scaling_type)::value, num_rows_per_cta_v, false, kWriteFp8>;
                 case ActivationType::Relu2:
                     return &doActivationKernel<T, GemmOutputType, ScaleBiasType,
                         IdentityAdaptor<cutlass::epilogue::thread::Relu2>, decltype(block_scaling_type)::value,
@@ -2963,6 +2968,8 @@ void doActivationDynamic(T* output, GemmOutputType const* gemm_result, float con
                 case ActivationType::SwigluBias:
                     return &doActivationKernel<T, GemmOutputType, ScaleBiasType, SwigluBiasAdaptor, NVFP4_TYPE, kRows,
                         true>;
+                case ActivationType::SiTu:
+                    return &doActivationKernel<T, GemmOutputType, ScaleBiasType, SiTuAdaptor, NVFP4_TYPE, kRows, true>;
                 case ActivationType::Relu2:
                     return &doActivationKernel<T, GemmOutputType, ScaleBiasType,
                         IdentityAdaptor<cutlass::epilogue::thread::Relu2>, 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 <int BLOCK_SIZE, int NUM_ROUTING_SAMPLES>
 __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<std::string, std::pair<size_t, size_t>> 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<std::string, std::pair<size_t, size_t>> out_map;
@@ -5331,14 +5352,14 @@ std::map<std::string, std::pair<size_t, size_t>> 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<size_t>(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<float*>(workspace_ptr_char + alphaWorkspace.second);
+        auto* beta = reinterpret_cast<float*>(workspace_ptr_char + betaWorkspace.second);
+        constexpr int kThreadsPerBlock = 128;
+        populateProfilerSiTuParamsKernel<<<ceilDiv(mNumExpertsPerNode, kThreadsPerBlock), kThreadsPerBlock, 0,
+            stream>>>(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,

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<float>::infinity();
+
+    template <class T>
+    __device__ T operator()(T const& gate, T const& linear) const
+    {
+        cutlass::epilogue::thread::Sigmoid<T> sigmoid{};
+        cutlass::epilogue::thread::Tanh<T> tanhFn{};
+        return tanhFn(gate * (1.0f / alpha)) * alpha * sigmoid(gate) * tanhFn(linear * (1.0f / beta)) * beta;
+    }
+};
+
 } // namespace kernels::cutlass_kernels
 
 TRTLLM_NAMESPACE_END

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<float const*>(swiglu_alpha.has_value() ? swiglu_alpha.value().const_data_ptr() : nullptr),
             reinterpret_cast<float const*>(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<float const*>(swiglu_alpha.has_value() ? swiglu_alpha.value().const_data_ptr() : nullptr),
             reinterpret_cast<float const*>(swiglu_beta.has_value() ? swiglu_beta.value().const_data_ptr() : nullptr),

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

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

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