NVIDIA · dc3671 · Aug 20, 2026 · Aug 17, 2026
@@ -100,6 +100,28 @@ NP=1 ./mpi_launch.sh ./run.sh config_gen.yaml --model nvidia/NVIDIA-Nemotron-3-N
 NP=2 ./mpi_launch.sh ./run.sh config_ctx.yaml --model Qwen/Qwen3-Next-80B-A3B-Instruct --layer-indices 6,7 --no-enable-attention-dp --mamba-ssm-cache-dtype float16 --batch-size 4
 NP=2 ./mpi_launch.sh ./run.sh config_gen.yaml --model Qwen/Qwen3-Next-80B-A3B-Instruct --layer-indices 6,7 --no-enable-attention-dp --mamba-ssm-cache-dtype float16 --batch-size 512
 
+# Run Kimi K3 (KimiLinear: KDA + MLA hybrid), $K3 being a local checkpoint directory.
+# Layer indices are 0-based, so 4,5,6,7 is three KDA layers followed by one full-attention
+# (MLA) layer -- one full period of the 3:1 pattern. Layer 0 is the only dense layer.
+# --moe-backend must be TRTLLM: the SiTU routed experts support only TRTLLM and
+# MEGAMOE_DEEPGEMM, the config files default to CUTLASS, and of those two only TRTLLM is in
+# the balance-method whitelist -- MEGAMOE_DEEPGEMM additionally needs
+# --balance-method NotModified. GEN also builds a prefill runner, so
+# --moe-backend-for-prefill must be set to match.
+# The attn-residual snapshot stack grows one entry every attn_res_block_size (12) layers and
+# is seeded from layer_indices[0], so a slice at 0,1,2,3 measures a cheaper attn_res than one
+# at 48,49,50,51.
+NP=4 ./mpi_launch.sh ./run.sh config_ctx.yaml --model $K3 --layer-indices 4,5,6,7 --tokens-per-block 64 --moe-backend TRTLLM
+NP=4 ./mpi_launch.sh ./run.sh config_gen.yaml --model $K3 --layer-indices 4,5,6,7 --tokens-per-block 64 --moe-backend TRTLLM --moe-backend-for-prefill TRTLLM --seq-len-q 1
+
+# Run Kimi K3 with 3 draft tokens per step (seq_len_q = 1 + 3). --spec-max-draft-len makes
+# the cache manager allocate the KDA multi-token verify buffers; without it the run stops
+# with an error. No sampler exists here, so no draft is ever accepted.
+NP=4 ./mpi_launch.sh ./run.sh config_gen.yaml --model $K3 --layer-indices 4,5,6,7 --tokens-per-block 64 --moe-backend TRTLLM --moe-backend-for-prefill TRTLLM --batch-size 32 --seq-len-q 4 --spec-max-draft-len 3
+
+# Run Kimi K3 at the serving parallelism (DEP16, see examples/kimi_k3/eval_extra_llm_options.yaml)
+NP=16 ./mpi_launch.sh ./run.sh config_gen.yaml --model $K3 --layer-indices 4,5,6,7 --tokens-per-block 64 --moe-backend TRTLLM --moe-backend-for-prefill TRTLLM --seq-len-q 1 --moe-max-num-tokens 33024 --use-low-precision-moe-combine
+
 # Run with DeepEP A2A
 NP=4 ./mpi_launch.sh -x TRTLLM_FORCE_COMM_METHOD=DEEPEP ./run.sh config_ctx.yaml --moe-backend CUTEDSL
 NP=4 ./mpi_launch.sh -x TRTLLM_FORCE_COMM_METHOD=DEEPEP ./run.sh config_gen.yaml --moe-backend CUTEDSL
@@ -377,3 +399,25 @@ Limitations:
 2. Error `huggingface_hub.errors.HfHubHTTPError: 429 Client Error: Too Many Requests for url: https://huggingface.co/nvidia/DeepSeek-R1-0528-FP4-v2/resolve/main/config.json`.
 
    Please use a local model through the `--model` option, or follow Hugging Face's instructions: "We had to rate limit your IP. To continue using our service, create a HF account or login to your existing account, and make sure you pass a HF_TOKEN if you're using the API."
+
+3. Error `Routing results are not replaced` with a MoE backend that routes inside the kernel.
+
+   `--balance-method` overrides routing in Python, so it needs the top-k to be computed outside the MoE kernel. Under attention DP that is always the case; with `--no-enable-attention-dp` the top-k can be fused into the kernel instead. Either add `-x TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING=1` or set `--balance-method NotModified`.
+
+4. Error `Kimi K3 needs --spec-max-draft-len N for seq_len_q N+1`.
+
+   Kimi K3's multi-token verify path reads speculative KDA state buffers that the cache manager only allocates when it is built with a speculative config. Add `--spec-max-draft-len` equal to `seq_len_q - 1`. Note that the harness has no sampler, so no draft is ever accepted and the result is the zero-acceptance corner of the verify kernel, not a steady-state number.
+
+5. Error `AttributeError: 'KimiLinearConfig' object has no attribute 'n_routed_experts'` with `--scaled-from`.
+
+   Weak scaling rewrites `n_routed_experts`, which Kimi K3's config does not have (it uses `num_experts`). `--scaled-from` is not supported for this model.
+
+6. Numbers look plausible but a factor of ~10 too slow on a hybrid model.
+
+   Several model fast paths fall back silently. Check these lines in the rank-0 log:
+
+   - `KDA kernel dispatch: prefill=optimized decode=optimized verify=optimized`
+   - `fused prefill/decode/verify projections on <N> KDA layers` — `N` must equal the number of KDA layers in `--layer-indices`; `0` means every KDA layer is running the reference path (~70 us/layer of glue around a ~5 us kernel)
+   - `Mamba Cache (kda-replay) is allocated` — only with `--spec-max-draft-len`; without it the legacy per-step verify buffers are used instead
+
+   For reference, Kimi K3 on 4 GPUs (DEP4, one GB200 node) with `--layer-indices 4,5,6,7` and `--load-format DUMMY` runs ~60 ms per CTX iteration at `--batch-size 1 --seq-len-q 8193`, and ~4.2 ms per GEN iteration at `--batch-size 32 --seq-len-q 4 --seq-len-kv-cache 8193 --spec-max-draft-len 3`.
@@ -52,13 +52,27 @@ def comma_separated_floats(s):
 parser.add_argument("--kv-pool-headroom", type=int, default=1)
 parser.add_argument("--enable-swa-scratch-reuse", action="store_true")
 # Model init args
+parser.add_argument(
+    "--vision-config",
+    type=str,
+    choices=["none", "checkpoint"],
+    help="Which vision tower to build from a composite multimodal checkpoint. Default"
+    ' "none": drop it and profile the text decoder layers, which is all this tool runs.',
+)
 parser.add_argument("--load-format", type=str, choices=["AUTO", "DUMMY"])
 parser.add_argument("--max-num-tokens", type=int)
 parser.add_argument("--moe-backend", type=str)
 # No choices= (same as --moe-backend above): both flags feed MoeConfig.backend,
 # whose Literal is the authoritative list. A second list here can only drift.
 parser.add_argument("--moe-backend-for-prefill", type=str)
 parser.add_argument("--moe-max-num-tokens", type=int)
+parser.add_argument(
+    "--spec-max-draft-len",
+    type=int,
+    help="Draft tokens per generation step. Required by models whose multi-token verify"
+    " path needs cache-manager-allocated speculative buffers (Kimi K3 KDA). Must equal"
+    " seq_len_q - 1.",
+)
 group = parser.add_mutually_exclusive_group()
 group.add_argument(
     "--use-low-precision-moe-combine", action="store_true", dest="use_low_precision_moe_combine"
@@ -137,6 +151,8 @@ def comma_separated_floats(s):
     args.kv_cache_dtype = "auto"
 if args.mamba_ssm_cache_dtype is None:
     args.mamba_ssm_cache_dtype = "auto"
+if args.vision_config is None:
+    args.vision_config = "none"
 if args.load_format is None:
     args.load_format = "DUMMY"
 if args.max_num_tokens is None:
@@ -153,8 +169,27 @@ def comma_separated_floats(s):
     parser.error("Both --replay-start-iter and --replay-stop-iter must be provided or none")
 if args.replay_verify_metadata is None:
     args.replay_verify_metadata = True
+if args.spec_max_draft_len is not None:
+    if args.run_type != "GEN":
+        parser.error("--spec-max-draft-len only applies to --run-type GEN")
+    if args.spec_max_draft_len < 1:
+        parser.error("--spec-max-draft-len must be >= 1")
+    if args.seq_len_q_list != [args.spec_max_draft_len + 1]:
+        parser.error(
+            f"--spec-max-draft-len {args.spec_max_draft_len} implies a single seq_len_q"
+            f" {args.spec_max_draft_len + 1}, got {args.seq_len_q_list}"
+        )
 print(args)
 
+# Speculative decoding shape only: the harness has no sampler, so no drafts are
+# ever accepted. This exists so the cache manager allocates the multi-token
+# verify buffers (Kimi K3 KDA replay caches) that seq_len_q > 1 needs.
+spec_config = None
+if args.spec_max_draft_len is not None:
+    from tensorrt_llm.llmapi.llm_args import SADecodingConfig
+
+    spec_config = SADecodingConfig(max_draft_len=args.spec_max_draft_len)
+
 # MPI args
 rank = mpi_rank()
 world_size = mpi_world_size()
@@ -175,6 +210,8 @@ def comma_separated_floats(s):
     layer_indices=args.layer_indices,
     kv_pool_headroom=args.kv_pool_headroom,
     enable_swa_scratch_reuse=args.enable_swa_scratch_reuse,
+    spec_config=spec_config,
+    vision_config=args.vision_config,
 )
 attn_workspace = torch.empty((0,), device="cuda", dtype=torch.int8)
 logger.info("Layer-wise benchmarks: Create KV cache manager  ... Done")
@@ -200,6 +237,8 @@ def comma_separated_floats(s):
     mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype,
     use_low_precision_moe_combine=args.use_low_precision_moe_combine,
     use_cuda_graph=args.use_cuda_graph,
+    spec_config=spec_config,
+    vision_config=args.vision_config,
 )
 logger.info("Layer-wise benchmarks: Create runner  ... Done")
 
@@ -270,6 +309,7 @@ def comma_separated_floats(s):
             mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype,
             use_low_precision_moe_combine=args.use_low_precision_moe_combine,
             use_cuda_graph=False,
+            vision_config=args.vision_config,
         )
     logger.info("Layer-wise benchmarks: Create runner for prefill  ... Done")
 

@@ -549,6 +549,16 @@ def _swap_linear_to_fp8_weight_read(
     return 1
 
 
+def _has_weights(module: nn.Module) -> bool:
+    """False once ``modeling_utils.remove_weights()`` has stripped a module.
+
+    Post-load finalization walks every decoder layer, so it must skip layers
+    whose parameters were dropped — the layer-wise benchmarks keep only the
+    profiled slice resident.
+    """
+    return not getattr(module, "_weights_removed", False)
+
+
 def _convert_moe_mlps_to_fp8_weight_read(
     model: nn.Module, include_fused_gate_up: bool = True
 ) -> int:
@@ -563,6 +573,8 @@ def _convert_moe_mlps_to_fp8_weight_read(
     count = 0
 
     for layer in model.layers:
+        if not _has_weights(layer):
+            continue
         moe = getattr(layer, "block_sparse_moe", None)
         if moe is None:
             continue
@@ -623,7 +635,7 @@ def _convert_kda_projections_to_fp8_weight_read(model: nn.Module) -> int:
     count = 0
 
     for layer in model.layers:
-        if not getattr(layer, "is_kda", False):
+        if not getattr(layer, "is_kda", False) or not _has_weights(layer):
             continue
         mixer = getattr(getattr(layer, "self_attn", None), "mixer", None)
         if mixer is None:
@@ -693,7 +705,7 @@ def _convert_mla_projections_to_fp8_weight_read(model: nn.Module) -> int:
     for layer in model.layers:
         # MLA layers are the non-KDA layers (each layer is exactly one of the
         # two); their projections live on the KimiK3MLAAttention mixer.
-        if getattr(layer, "is_kda", False):
+        if getattr(layer, "is_kda", False) or not _has_weights(layer):
             continue
         mixer = getattr(getattr(layer, "self_attn", None), "mixer", None)
         if mixer is None:
@@ -2235,6 +2247,20 @@ def forward(
         prefix_sum = prefix_sum + hidden_states
         return prefix_sum, num_snapshots
 
+    def skip_forward(
+        self,
+        hidden_states: torch.Tensor,
+        block_residual: torch.Tensor,
+        attn_metadata: AttentionMetadata,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """No-op stand-in for ``forward``, matching ``DecoderLayer.skip_forward``.
+
+        ``modeling_utils.skip_forward()`` only drops a module's weights when it
+        finds this attribute, so without it the layer-wise benchmarks would
+        allocate all 93 layers instead of the profiled slice.
+        """
+        return hidden_states, block_residual
+
 
 # ---------------------------------------------------------------------------
 # Model.
@@ -2472,7 +2498,7 @@ def checkpoint_name_plan(
         # when the routed MoE is TP-sharded (moe_ep=1 -> ids 0..num_experts)).
         expert_jobs = []
         for layer_idx, layer in enumerate(self.model.layers):
-            if not getattr(layer, "is_moe", False):
+            if not getattr(layer, "is_moe", False) or not _has_weights(layer):
                 continue
             moe = layer.block_sparse_moe
             base = f"{prefix}model.layers.{layer_idx}.block_sparse_moe.experts"
@@ -2541,7 +2567,7 @@ def _load_trunk_params(
         mla_mixers = [
             layer.self_attn.mixer
             for layer in self.model.layers
-            if not getattr(layer, "is_kda", True)
+            if not getattr(layer, "is_kda", True) and _has_weights(layer)
         ]
         mla_kv_b_mixers = {id(mixer.kv_b_proj.weight): mixer for mixer in mla_mixers}
         mla_head_shard_linears = {}
@@ -2864,7 +2890,7 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None:
         # This must run after every KDA parameter is loaded and sharded.
         num_kda_fused = 0
         for layer in self.model.layers:
-            if getattr(layer, "is_kda", False):
+            if getattr(layer, "is_kda", False) and _has_weights(layer):
                 if not kda_fp8:
                     layer.self_attn.finalize_decode_weights()
                 num_kda_fused += int(
@@ -2919,7 +2945,7 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None:
                     # in BF16.
                     n_glue = 0
                     for layer in self.model.layers:
-                        if getattr(layer, "is_kda", False):
+                        if getattr(layer, "is_kda", False) and _has_weights(layer):
                             layer.self_attn.finalize_decode_weights_fp8()
                             n_glue += int(layer.self_attn._bfa_proj_weight is not None)
                     logger.info(

@@ -2,6 +2,12 @@
 
 from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Gate, Deepseekv3MoE
 from tensorrt_llm._torch.models.modeling_deepseekv4 import DeepseekV4Gate, DeepseekV4MoE
+from tensorrt_llm._torch.models.modeling_kimi_linear import (
+    KimiK3MoEGate,
+    KimiK3MoERuntime,
+    KimiKDARuntime,
+    KimiMLARuntime,
+)
 from tensorrt_llm._torch.models.modeling_nemotron_h import MLPLayer, NemotronHMOE
 from tensorrt_llm._torch.models.modeling_qwen3_next import (
     Qwen3NextGatedDeltaNet,
@@ -26,6 +32,13 @@ def mark_ranges():
     Qwen3NextSparseMoeBlock.forward = nvtx.annotate("Qwen3NextSparseMoeBlock")(
         Qwen3NextSparseMoeBlock.forward
     )
+    # Kimi K3. `KimiK3MLAAttention` overrides `MLA.forward`, so the range below
+    # is on its `KimiMLARuntime` wrapper. The gate is entered through
+    # `compute_logits`, not `forward`. Its MLPs are the shared `GatedMLP`.
+    KimiKDARuntime.forward = nvtx.annotate("KimiKDARuntime")(KimiKDARuntime.forward)
+    KimiMLARuntime.forward = nvtx.annotate("KimiMLARuntime")(KimiMLARuntime.forward)
+    KimiK3MoERuntime.forward = nvtx.annotate("KimiK3MoERuntime")(KimiK3MoERuntime.forward)
+    KimiK3MoEGate.compute_logits = nvtx.annotate("KimiK3MoEGate")(KimiK3MoEGate.compute_logits)
     MLA.forward = nvtx.annotate("MLA")(MLA.forward)
     Attention.forward = nvtx.annotate("Attention")(Attention.forward)
     MoE.forward = nvtx.annotate("MoE")(MoE.forward)