[None][feat] modeling_v2 catalog: Kimi K3 decode, KDA / MLA and generic-path kernels as entries (sm_100)

cpp/tensorrt_llm/kernels/kimiK3AttnRes/attnResFwd.cu

@@ -1046,12 +1046,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;
@@ -1321,7 +1327,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())
@@ -1331,11 +1340,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())
     {
@@ -1351,14 +1366,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);
     }
 }
 
@@ -1380,13 +1395,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;
@@ -1439,6 +1463,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)
     {
@@ -1638,7 +1663,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())
@@ -1652,13 +1680,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];
@@ -1678,7 +1707,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.
@@ -1887,18 +1916,19 @@ 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;
+    S1Handoff const handoff{params.earlyTrigger};
 
     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/attnResFwd.h

@@ -53,6 +53,7 @@ struct AttnResFwdParams
     int hiddenSize;                        // H
     float rmsEps;
     float outputRmsEps;
+    bool earlyTrigger = false;
 };
 
 //! Launches the fused attention-residual forward on the supplied stream.

cpp/tensorrt_llm/thop/attnResOp.cpp

@@ -125,7 +125,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor> attn_res_fwd(
 }
 
 at::Tensor attn_res_rmsnorm_fwd(at::Tensor layer_residual, at::Tensor block_residual, at::Tensor res_weight,
-    at::Tensor rms_weight, at::Tensor output_rms_weight, double rms_eps, double output_rms_eps)
+    at::Tensor rms_weight, at::Tensor output_rms_weight, double rms_eps, double output_rms_eps, bool early_trigger)
 {
     TORCH_CHECK(layer_residual.dim() == 3, "attn_res_rmsnorm_fwd: layer_residual must be [T, B, H]");
     TORCH_CHECK(block_residual.dim() == 4, "attn_res_rmsnorm_fwd: block_residual must be [K, T, B, H]");
@@ -179,6 +179,7 @@ at::Tensor attn_res_rmsnorm_fwd(at::Tensor layer_residual, at::Tensor block_resi
     params.hiddenSize = H;
     params.rmsEps = static_cast<float>(rms_eps);
     params.outputRmsEps = static_cast<float>(output_rms_eps);
+    params.earlyTrigger = early_trigger;
 
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
     kernels::kimiK3AttnRes::invokeAttnResRmsNormFwd(params, stream);
@@ -188,7 +189,7 @@ at::Tensor attn_res_rmsnorm_fwd(at::Tensor layer_residual, at::Tensor block_resi
 
 std::tuple<at::Tensor, at::Tensor> attn_res_add_rmsnorm_fwd(at::Tensor layer_residual, at::Tensor layer_residual_add,
     at::Tensor block_residual, at::Tensor res_weight, at::Tensor rms_weight, at::Tensor output_rms_weight,
-    double rms_eps, double output_rms_eps)
+    double rms_eps, double output_rms_eps, bool early_trigger)
 {
     TORCH_CHECK(layer_residual.dim() == 3, "attn_res_add_rmsnorm_fwd: layer_residual must be [T, B, H]");
     TORCH_CHECK(layer_residual_add.sizes() == layer_residual.sizes(),
@@ -248,6 +249,7 @@ std::tuple<at::Tensor, at::Tensor> attn_res_add_rmsnorm_fwd(at::Tensor layer_res
     params.hiddenSize = H;
     params.rmsEps = static_cast<float>(rms_eps);
     params.outputRmsEps = static_cast<float>(output_rms_eps);
+    params.earlyTrigger = early_trigger;
 
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
     kernels::kimiK3AttnRes::invokeAttnResAddRmsNormFwd(params, stream);
@@ -364,11 +366,11 @@ TORCH_LIBRARY_FRAGMENT(trtllm, m)
     m.def(
         "attn_res_rmsnorm_fwd(Tensor layer_residual, Tensor block_residual, "
         "Tensor res_weight, Tensor rms_weight, Tensor output_rms_weight, "
-        "float rms_eps, float output_rms_eps) -> Tensor");
+        "float rms_eps, float output_rms_eps, bool early_trigger=False) -> Tensor");
     m.def(
         "attn_res_add_rmsnorm_fwd(Tensor layer_residual, Tensor layer_residual_add, "
         "Tensor block_residual, Tensor res_weight, Tensor rms_weight, Tensor output_rms_weight, "
-        "float rms_eps, float output_rms_eps) -> (Tensor, Tensor)");
+        "float rms_eps, float output_rms_eps, bool early_trigger=False) -> (Tensor, Tensor)");
     m.def(
         "attn_res_add_rmsnorm_persistent_fwd(Tensor layer_residual, Tensor? layer_residual_add, "
         "Tensor block_residual, Tensor res_weight, Tensor rms_weight, Tensor output_rms_weight, "

docs/source/features/torch_compile_and_piecewise_cuda_graph.md

@@ -139,6 +139,18 @@ the Multimodal Model Family.
 
 ## Development Guide
 
+Executor host-upload buffers must keep stable device addresses across capture and replay. Their internal
+`DeviceCopyBuffer` owns the tensor and last-uploaded host snapshot; every cached writer must use the same
+owner's `update(host)` method. A Mamba metadata CUDA alias shares its cache manager's owner, rather than
+creating another cache over the same storage. Returned tensors are borrowed views, not immutable tensors;
+writing them directly bypasses the cached-upload contract.
+
+Changed updates use independent host staging, pinned according to the owner's configuration, and write only
+the supplied prefix. Shorter updates preserve the device tail and its cached values; unchanged prefixes
+enqueue no upload. The C++ Mamba manager uses `force=True` to retain its unconditional-upload behavior while
+keeping aliases' shared snapshot coherent. `TLLM_DEBUG_MODE=1` verifies skipped values outside CUDA graph
+capture; capture does not add a device synchronization for this check.
+
 ### Background Knowledge
 
 Currently, TRT-LLM mainly relies on torch.compile **fullgraph** mode to enable Piecewise CUDA Graph feature, which means all the operations in the model must be recognized by torch.compile.

docs/source/models/supported-models.md

@@ -97,7 +97,7 @@ statuses for the same architecture in the two matrices.
 [^14]: Requires `transformers>=5.7.0`: MiniCPM-V 4.6 was upstreamed into transformers as a native model type (`minicpmv4_6`) and the checkpoint ships no remote code (`auto_map`) to fall back on. The Qwen3.5-hybrid text tower runs in BF16. Image, video, and text inputs are supported in this release (video reuses the same NaViT-packed vision path as image via `MiniCPMV4_6InputProcessor`).
 [^15]: Kimi K3 is only supported on NVIDIA Blackwell GPUs (`SM100` family); see the [Kimi K3 deployment guide](../deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md). Which recipes fit is set by the attention layout rather than by the GPU count: DEP16 (`enable_attention_dp: true`) replicates the BF16 non-expert weights on every rank (114 GB) on top of the MXFP4 routed experts at 16-way expert parallelism (90 GB), needing 210 GB per rank, so it requires GB300-class per-GPU memory. TEP16 (`enable_attention_dp: false`) shards those non-expert weights instead and needs 115 GB per rank; it is validated end-to-end on GB200 (`SM100`) at 16 GPUs. On B200 (`SM100`), kernel and module support is functional and covered by CI. Note that the FP8 weight-read path (TRTLLM-14765) does not relax the DEP16 requirement: the conversion runs after the weights are already resident in BF16, so it lowers the steady-state footprint but not the load-time peak.
 [^16]: Guided decoding for `Glm4MoeForCausalLM` is currently supported only with the `xgrammar` backend; `llguidance` is not working.
-[^17]: Kimi K3 has no MTP or EAGLE-3 head, and its DSpark checkpoints are not compatible with plain `DFlash`.
+[^17]: Kimi K3 has no MTP or EAGLE-3 head, and its DSpark checkpoints are not compatible with plain `DFlash`. Speculators-format DSpark checkpoints require `aux_hidden_state_layer_ids` to be a nonempty list of positive integers; booleans, floats and strings are rejected rather than coerced.
 [^18]: NGram and standalone Suffix Automaton (SA) use model-free drafting on the PyTorch backend, so they are not listed in individual entries. This does not imply universal end-to-end support: compatibility depends on each model's multi-token verification and cache-management paths and may be untested or explicitly restricted.
 [^19]: KV cache reuse for hybrid recurrent-attention models requires an explicit recurrent-state snapshot policy, such as `kv_cache_config.mamba_state_config.periodic_snapshot_interval`; the model default disables reuse when no snapshot policy is configured.
 [^20]: Supports text, image, and video inputs, MTP (including with attention data parallelism), and FP8 KV cache. Image and video inputs require `transformers==5.17.0`. Beam search is not supported. See the [GLM-5.3-Flash deployment guide](../deployment-guide/deployment-guide-for-glm-5.3-flash-on-trtllm.md) for setup and feature-specific requirements.

tensorrt_llm/_device_copy.py

@@ -0,0 +1,45 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+
+import torch
+
+_verify_skipped_device_copies = os.environ.get("TLLM_DEBUG_MODE", "")[0:1] == "1"
+
+
+class DeviceCopyBuffer:
+    def __init__(self, tensor: torch.Tensor, *, pin_memory: bool = False):
+        self._tensor = tensor
+        self._pin_memory = pin_memory
+        self._last_host_values = None
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        return self._tensor
+
+    def update(self, host: torch.Tensor, *, force: bool = False) -> None:
+        flat = host.reshape(-1)
+        n = flat.numel()
+        last = self._last_host_values
+        if not force and last is not None and last.numel() >= n and torch.equal(last[:n], flat):
+            if _verify_skipped_device_copies:
+                if self._tensor.is_cuda:
+                    if torch.cuda.is_current_stream_capturing():
+                        return
+                    torch.cuda.synchronize(self._tensor.device)
+                held = self._tensor.view(-1)[:n].cpu()
+                assert torch.equal(held, flat.to(held.dtype)), (
+                    f"DeviceCopyBuffer skipped a copy, but the buffer holds {held.tolist()} "
+                    f"instead of {flat.tolist()}: something else wrote it"
+                )
+            return
+        staging = torch.empty_like(flat, device="cpu", pin_memory=self._pin_memory)
+        staging.copy_(flat)
+        self._tensor.view(-1)[:n].copy_(staging, non_blocking=True)
+        if last is None or last.numel() <= n:
+            self._last_host_values = flat.clone()
+        else:
+            last = last.clone()
+            last[:n] = flat
+            self._last_host_values = last

tensorrt_llm/_torch/_experimental/modeling_v2/README.md

@@ -279,7 +279,9 @@ Perf is measured, never gated.
 
 **Every entry the gpt-oss and deepseek default targets call is certified on
 sm_100 and sm_103, and both targets are gated on both.** The qwen4_exp target
-and the entries it owns are certified on sm_103.
+and the entries it owns are certified on sm_103. Kimi K3's sm_100 operation
+receipts remain in its adjacent records; its generic-operation regressions
+exercise the relocated entries.
 
 The default targets' 21 entries -- 20 shared, 1 (`fused_moe`) owned by deepseek's default target -- each drive
 their cells through `__call__` and `reference`. Where that runs depends on the