[None][feat] modeling_v2: Kimi K3 MXFP4 target, sm_100, tp16 attention + moe tp4 x ep4

cpp/tensorrt_llm/kernels/kdaDecode/kdaDecodeLegacy.cu

@@ -275,6 +275,8 @@ __device__ __forceinline__ Sum2 block_reduce_sum2_for(float x, float y, float* s
         block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f;
         block_x = warp_reduce_sum(block_x);
         block_y = warp_reduce_sum(block_y);
+        // Lane 0 overwrites partials that other lanes of this warp read above.
+        __syncwarp();
         if (lane == 0)
         {
             scratch[0] = block_x;
@@ -349,6 +351,8 @@ __device__ __forceinline__ Sum2 block_reduce_sum2_active_for(float x, float y, f
         block_y = lane < kReduceWarps ? scratch[kReduceWarps + lane] : 0.0f;
         block_x = warp_reduce_sum(block_x);
         block_y = warp_reduce_sum(block_y);
+        // Lane 0 overwrites partials that other lanes of this warp read above.
+        __syncwarp();
         if (lane == 0)
         {
             scratch[0] = block_x;

cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu

@@ -411,9 +411,17 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c
                                     float2 f[2] = {__bfloat1622float2(v2[0]), __bfloat1622float2(v2[1])};
                                     if constexpr (FULL_N12)
                                     {
-                                        if (n == AN - 1 && lane == 0)
+                                        if (n == AN - 1)
                                         {
-                                            cute::arrive_barrier(plan.bar_consumed[chunk_slot]);
+                                            // The producer refills the slots through the async proxy (TMA): a
+                                            // cross-proxy fence orders this lane's generic-proxy reads before it.
+                                            asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
+                                            // Every lane's reads of the chunk's slots before lane 0 releases them.
+                                            __syncwarp();
+                                            if (lane == 0)
+                                            {
+                                                cute::arrive_barrier(plan.bar_consumed[chunk_slot]);
+                                            }
                                         }
                                     }
                                     tmem_st_32dp32bNx<4>(
@@ -492,6 +500,11 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c
                 }
                 if constexpr (!FULL_N12)
                 {
+                    // The producer refills the slots through the async proxy (TMA): a cross-proxy fence orders this
+                    // lane's generic-proxy reads before it.
+                    asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
+                    // Every lane's reads of the chunk's slots before lane 0 releases them to the producer.
+                    __syncwarp();
                     if (lane == 0)
                     {
                         cute::arrive_barrier(plan.bar_consumed[chunk_slot]);
@@ -551,6 +564,10 @@ __global__ void __launch_bounds__(BLK, 1) attn_res_fwd_online_v2_kernel(bf16_t c
                 {
                     cross_warp_tail(lane);
                 }
+                // ws_stats is one buffer reused every chunk. The barrier above orders this chunk's writes before the
+                // reads; this one orders the reads before the next chunk's writes, so a warp that has finished
+                // reading cannot overwrite its row while a slower warp still reads it.
+                cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 0);
                 float logit_n[N_CHUNK];
 #pragma unroll
                 for (int n = 0; n < N_CHUNK; n++)
@@ -972,6 +989,9 @@ __global__ void __launch_bounds__(BLK, 1)
                 }
                 cutlass::arch::NamedBarrier::sync(CONSUMER_THREADS, 1);
             }
+            // The producer refills the tile through the async proxy (TMA): a cross-proxy fence orders this thread's
+            // generic-proxy reads before it.
+            asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
             cute::arrive_barrier(plan.bar_consumed[slot]);
         }
     }
@@ -1027,12 +1047,18 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_single_cta_kernel(bf16
     bf16_t const* __restrict__ layer_res, bf16_t const* __restrict__ layer_res_add, bf16_t const* __restrict__ res_w,
     bf16_t const* __restrict__ rms_w, bf16_t const* __restrict__ output_rms_w, bf16_t* __restrict__ updated_layer_res,
     bf16_t* __restrict__ output, float* __restrict__ rsigma_out, float* __restrict__ probs_out,
-    float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens)
+    float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens, bool early_trigger)
 {
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
     if constexpr (ENABLE_PDL)
     {
         cudaGridDependencySynchronize();
+        // The dependent grid still waits for this one to complete before reading its
+        // output; triggering here lets it launch and stream its weights while this kernel runs.
+        if (early_trigger)
+        {
+            cudaTriggerProgrammaticLaunchCompletion();
+        }
     }
 
     constexpr int H = 7168;
@@ -1302,7 +1328,10 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_single_cta_kernel(bf16
 
     if constexpr (ENABLE_PDL)
     {
-        cudaTriggerProgrammaticLaunchCompletion();
+        if (!early_trigger)
+        {
+            cudaTriggerProgrammaticLaunchCompletion();
+        }
     }
 #else
     if (cute::thread0())
@@ -1312,11 +1341,17 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_single_cta_kernel(bf16
 #endif
 }
 
+// Decode handoff options for the s1 kernels (see AttnResFwdParams).
+struct S1Handoff
+{
+    bool early_trigger = false;
+};
+
 template <int N, bool FUSE_OUTPUT_RMS_NORM = false, bool FUSE_LAYER_ADD = false>
 static void launch_s1_single_cta(bf16_t const* block_residual, bf16_t const* layer_residual,
     bf16_t const* layer_residual_add, bf16_t const* res_weight, bf16_t const* rms_weight,
     bf16_t const* output_rms_weight, bf16_t* updated_layer_residual, bf16_t* output, float* rsigma, float* probs,
-    float* logits, float rms_eps, float output_rms_eps, int num_tokens, cudaStream_t stream)
+    float* logits, float rms_eps, float output_rms_eps, int num_tokens, cudaStream_t stream, S1Handoff handoff = {})
 {
     if (tensorrt_llm::common::getEnvEnablePDL())
     {
@@ -1332,14 +1367,14 @@ static void launch_s1_single_cta(bf16_t const* block_residual, bf16_t const* lay
         config.numAttrs = 1;
         cudaLaunchKernelEx(&config, kernel, block_residual, layer_residual, layer_residual_add, res_weight, rms_weight,
             output_rms_weight, updated_layer_residual, output, rsigma, probs, logits, rms_eps, output_rms_eps,
-            num_tokens);
+            num_tokens, handoff.early_trigger);
     }
     else
     {
         attn_res_fwd_s1_single_cta_kernel<N, FUSE_OUTPUT_RMS_NORM, FUSE_LAYER_ADD, false>
             <<<num_tokens, 256, 0, stream>>>(block_residual, layer_residual, layer_residual_add, res_weight, rms_weight,
                 output_rms_weight, updated_layer_residual, output, rsigma, probs, logits, rms_eps, output_rms_eps,
-                num_tokens);
+                num_tokens, false);
     }
 }
 
@@ -1361,13 +1396,22 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_splitk_kernel(bf16_t c
     bf16_t const* __restrict__ layer_res, bf16_t const* __restrict__ layer_res_add, bf16_t const* __restrict__ res_w,
     bf16_t const* __restrict__ rms_w, bf16_t const* __restrict__ output_rms_w, bf16_t* __restrict__ updated_layer_res,
     bf16_t* __restrict__ output, float* __restrict__ rsigma_out, float* __restrict__ probs_out,
-    float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens)
+    float* __restrict__ logits_out, float rms_eps, float output_rms_eps, int num_tokens, bool early_trigger)
 {
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
-    if constexpr (ENABLE_PDL)
+    auto const wait_for_previous_grid = [&]
     {
-        cudaGridDependencySynchronize();
-    }
+        if constexpr (ENABLE_PDL)
+        {
+            cudaGridDependencySynchronize();
+            // See attn_res_fwd_s1_single_cta_kernel.
+            if (early_trigger)
+            {
+                cudaTriggerProgrammaticLaunchCompletion();
+            }
+        }
+    };
+    wait_for_previous_grid();
 
     namespace cg = cooperative_groups;
     constexpr int H = 7168;
@@ -1420,6 +1464,7 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_splitk_kernel(bf16_t c
 
     float sq[N] = {};
     float dot[N] = {};
+    // Every candidate in one pass, so all loads are in flight together.
 #pragma unroll
     for (int ki = tid; ki < K_PER_CTA; ki += THREADS)
     {
@@ -1619,7 +1664,10 @@ __global__ void __launch_bounds__(256, 1) attn_res_fwd_s1_splitk_kernel(bf16_t c
 
     if constexpr (ENABLE_PDL)
     {
-        cudaTriggerProgrammaticLaunchCompletion();
+        if (!early_trigger)
+        {
+            cudaTriggerProgrammaticLaunchCompletion();
+        }
     }
 #else
     if (cute::thread0())
@@ -1633,13 +1681,14 @@ template <int N, int GROUPS = 8, bool FUSE_OUTPUT_RMS_NORM = false, bool FUSE_LA
 static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_residual,
     bf16_t const* layer_residual_add, bf16_t const* res_weight, bf16_t const* rms_weight,
     bf16_t const* output_rms_weight, bf16_t* updated_layer_residual, bf16_t* output, float* rsigma, float* probs,
-    float* logits, float rms_eps, float output_rms_eps, int num_tokens, cudaStream_t stream)
+    float* logits, float rms_eps, float output_rms_eps, int num_tokens, cudaStream_t stream, S1Handoff handoff = {})
 {
     constexpr int K_PER_CTA = 7168 / GROUPS;
     constexpr int WARPS = 8;
     constexpr size_t smem_size
         = (size_t) N * K_PER_CTA * sizeof(float) + (size_t) WARPS * N * sizeof(float2) + (size_t) N * sizeof(float);
     bool const enable_pdl = tensorrt_llm::common::getEnvEnablePDL();
+    bool early_trigger = enable_pdl && handoff.early_trigger;
     auto kernel = enable_pdl ? &attn_res_fwd_s1_splitk_kernel<N, GROUPS, FUSE_OUTPUT_RMS_NORM, FUSE_LAYER_ADD, true>
                              : &attn_res_fwd_s1_splitk_kernel<N, GROUPS, FUSE_OUTPUT_RMS_NORM, FUSE_LAYER_ADD, false>;
     static std::once_flag attrs_set[2][64];
@@ -1659,7 +1708,7 @@ static void launch_s1_splitk(bf16_t const* block_residual, bf16_t const* layer_r
     void* args[] = {const_cast<bf16_t**>(&block_residual), const_cast<bf16_t**>(&layer_residual),
         const_cast<bf16_t**>(&layer_residual_add), const_cast<bf16_t**>(&res_weight), const_cast<bf16_t**>(&rms_weight),
         const_cast<bf16_t**>(&output_rms_weight), &updated_layer_residual, &output, &rsigma, &probs, &logits, &rms_eps,
-        &output_rms_eps, &num_tokens};
+        &output_rms_eps, &num_tokens, &early_trigger};
     cudaLaunchConfig_t config{};
     // One cluster per token. clusterDim stays GROUPS, so clusters are contiguous
     // spans of the grid and each token's DSM exchange stays within its own.
@@ -1868,18 +1917,22 @@ static void launchAttnResDecodeRmsNorm(AttnResFwdParams const& params, cudaStrea
 
     auto const* layer_residual_add = FUSE_LAYER_ADD ? params.layerResidualAdd : nullptr;
     auto* updated_layer_residual = FUSE_LAYER_ADD ? params.updatedLayerResidual : nullptr;
+    // The dependent kernel launches (and streams its weights) while this one runs: it waits for
+    // this grid before reading the output.
+    constexpr bool early_trigger = true;
+    S1Handoff const handoff{early_trigger};
 
     if constexpr (N <= 4)
     {
         launch_s1_single_cta<N, true, FUSE_LAYER_ADD>(params.blockResidual, params.layerResidual, layer_residual_add,
             params.resWeight, params.rmsWeight, params.outputRmsWeight, updated_layer_residual, params.output, nullptr,
-            nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream);
+            nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream, handoff);
     }
     else
     {
         launch_s1_splitk<N, 8, true, FUSE_LAYER_ADD>(params.blockResidual, params.layerResidual, layer_residual_add,
             params.resWeight, params.rmsWeight, params.outputRmsWeight, updated_layer_residual, params.output, nullptr,
-            nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream);
+            nullptr, nullptr, params.rmsEps, params.outputRmsEps, params.seqLen, stream, handoff);
     }
 }
 

cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwdPersistentFused.cu

@@ -650,6 +650,11 @@ __global__ void __launch_bounds__(BLK, 1)
                     default: __builtin_unreachable();
                     }
                 }
+                // The producer refills the slots through the async proxy (TMA): a cross-proxy fence orders this lane's
+                // generic-proxy reads before it.
+                asm volatile("fence.proxy.async.shared::cta;" ::: "memory");
+                // Every lane's reads of the chunk's slots before lane 0 releases them to the producer.
+                __syncwarp();
                 if (lane == 0)
                 {
                     mbarrier_arrive(plan.bar_consumed[chunk_slot]);

tensorrt_llm/_torch/_experimental/modeling_v2/README.md

@@ -181,8 +181,10 @@ Perf is measured, never gated.
 
 ## Status of every record in this tree
 
-**The catalog is fully certified on sm_103. The targets construct but have
-never executed.**
+**The catalog is certified: 19 entries on sm_103, and 18 on sm_100 (B200 /
+GB200): the 16 Kimi K3 entries, where their first caller runs, and
+`cublas_mm` and `flashinfer_rmsnorm`. The targets construct but have never
+executed.**
 
 Two things voided every receipt in the move: each catalog test file was
 rewritten, and the targets moved from sm_100 (B200) to sm_103 (GB300), where

tensorrt_llm/_torch/_experimental/modeling_v2/_router_index.py

@@ -59,6 +59,7 @@
 MODELING_V2_ROUTERS = {
     "GptOssForCausalLM": "models.gpt_oss.routing",
     "DeepseekV3ForCausalLM": "models.deepseek_v3.routing",
+    "KimiK3ForConditionalGeneration": "models.kimi_k3_vl.routing",
 }
 
 #: Backends whose model construction reaches ``modeling_v2_resolve``. The

tensorrt_llm/_torch/_experimental/modeling_v2/catalog/activation/k3_situ_mul.md

@@ -0,0 +1,105 @@
+---
+receipts:
+  sm_100: {status: passed, tests: 4}
+---
+
+# k3_situ_mul
+
+**Wraps** `torch.ops.trtllm.k3_situ_mul` (one call).
+
+## Semantics
+
+The SiTU-gated multiply of the Kimi K3 dense MLP (`SituAndMul` in `tensorrt_llm/_torch/modules/situ.py`) for at
+most 8 tokens, as one CuTe DSL kernel that lets the next kernel launch at once under programmatic dependent launch.
+The last dim of the input holds the gate half followed by the up half:
+
+```
+K = gu.shape[1] // 2;   g = gu[m, k],  u = gu[m, K + k]            0 <= m < M <= 8, 0 <= k < K, in fp32
+a         = (beta * tanh(g / beta)) * sigmoid(g)                    sigmoid(g) = 1 / (1 + exp(-g))
+v         = linear_beta * tanh(u / linear_beta)                     if linear_beta is not None, else v = u
+out[m, k] = bf16(a * v)
+```
+
+The arithmetic is fp32 in this order, with IEEE division and the full-precision `exp` and `tanh`, and the result
+is rounded to bf16 once, round-to-nearest-even. That is the order of `SituAndMul`'s eager path. The `exp` and
+`tanh` implementations are not torch's or Triton's, so the result is not bit-identical to `SituAndMul` (eager or
+fused); they agree within the tolerance below. Rows are independent: a token's output row has the same bits
+whatever `M` is.
+
+Fusion boundary: the single call computes the activation only. The caller owns the gate_up projection that
+produces `gu` and the down projection that consumes the output. There is no quantization of the output.
+
+## Signature
+
+```python
+def k3_situ_mul(gu: torch.Tensor, beta: float = 1.0, linear_beta: float | None = None) -> torch.Tensor
+```
+
+| Argument | Shape | Dtype | Layout | Device |
+|---|---|---|---|---|
+| `gu` | `[M, 2 * K]`, gate half first, `1 <= M <= 8`, `K % 8 == 0` | bf16 | contiguous, 16-byte-aligned | CUDA |
+| `beta` | scalar | Python float | — | — |
+| `linear_beta` | scalar or `None` | Python float, not 0 (the wrapper asserts) | — | — |
+| returns | `[M, K]` | bf16 | newly allocated, contiguous | `gu`'s device |
+
+Every element of the returned tensor is written; `gu` is not mutated.
+
+### Certified arguments
+
+The dense MLP of Kimi K3 at TP16 (gate_up output `[M, 4224]`, activation `[M, 2112]`) at every `M` in 1..8, with
+`(beta, linear_beta)` = `(4.0, 25.0)` (the Kimi K3 checkpoint's) and `(1.0, None)` (the defaults). Inputs are
+`gu ~ 2 * N(0, 1)`, rounded to bf16. The gate per cell: `max|y - ref| / max|ref| <= 8e-3` against the formula
+above evaluated in fp64 on the same bf16 input, identical bits on a repeated call, and each `M`'s rows
+bit-identical to the same rows of the 8-row call. Also certified, with both settings at `M` 1 and 8: calls
+captured in a CUDA graph after an eager call per key, replayed with `gu` rewritten in place, return the bits of
+eager calls on the new `gu`. Other widths and `beta` / `linear_beta` values are accepted but not certified.
+
+## Metadata consumed
+
+None. The op reads no attention metadata, KV cache or module state; every input is an argument, and the kernel
+runs on the current CUDA stream of `gu`'s device.
+
+One process-global cache sits behind it, and it is result-neutral: the compiled kernel per key
+`(K, linear_beta is not None, PDL on/off)`. The first call with a new key compiles the kernel with the CuTe DSL,
+which costs host time; `M`, `beta` and the value of `linear_beta` are runtime arguments and never recompile. That
+first call refuses to run under CUDA-graph capture: it raises `RuntimeError` ("run once per shape outside
+CUDA-graph capture first") before launching anything. Call every key once eagerly, then capture; captured and
+later calls reuse the compiled kernel. The PDL setting is read from `TRTLLM_ENABLE_PDL` on every call, so
+changing it within a process adds a key.
+
+## Preconditions
+
+The op checks these and raises `ValueError` ("k3_situ_mul: unsupported call ...") when one fails:
+
+- `gu` is a 2-D contiguous CUDA bf16 tensor with `1 <= M <= 8` rows.
+- `gu.shape[1] % 16 == 0`, so each half is a whole number of 16-byte vectors.
+- `gu` starts at a 16-byte-aligned address.
+
+Certified refusals: 0 and 9 rows, a width of 4216, a `gu` starting 2 bytes past a 16-byte boundary, fp16, a
+row-strided `gu` and a 1-D `gu`.
+
+The wrapper adds one check, because the op does not fail on it: `linear_beta` must not be `0.0`. The op hands the
+kernel `linear_beta or 1.0`, so `linear_beta=0.0` would run as `1.0` (`v = tanh(u)`) with no error, where the
+formula gives `0 * tanh(u / 0)`. The wrapper raises `AssertionError` before calling the op (certified).
+
+Not checked by the op:
+
+- `K >= 8`; an empty `gu` (width 0) passes the op's check.
+- The architecture. The kernel needs no SM 10.x feature (no tcgen05, no clusters; `griddepcontrol` needs SM 9.0
+  or newer), but only sm_100 has been measured.
+- The CuTe DSL (`cutlass`) and `cuda-python` (`cuda.bindings`) are importable. The op imports the kernel module
+  after its check, so without them a call raises `ImportError`.
+
+## Notes
+
+- Programmatic dependent launch (PDL). With `TRTLLM_ENABLE_PDL` unset or `1` the kernel is launched with PDL; any
+  other value launches it without. There is no `trigger_early` argument: every CTA executes
+  `griddepcontrol.launch_dependents` first, then `griddepcontrol.wait`, and only then reads `gu`. So the next
+  kernel on the stream, if launched with PDL, launches at once and can stream its own weights while this kernel
+  and its predecessor run (the CTM GEMVs do); it must execute `griddepcontrol.wait`
+  (`cudaGridDependencySynchronize`) before it reads the output, as those GEMVs do before they read their
+  activation. Kernels launched without PDL, torch's included, start after this one completes as usual. None of
+  this changes results.
+- Two identical calls return identical bits.
+- Grid: `ceil(K / 1024)` x 8 CTAs of 128 threads, 8 columns per thread; rows past `M` exit at once.
+- Only sm_100 has been measured.