[TRTLLM-14705][fix] Kimi K3 B200 enablement: MLA decode dispatch fix, L0 wiring, docs

docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md

@@ -10,7 +10,15 @@ This guide uses Slurm and the `trtllm-llmapi-launch` multi-node launcher. The co
 
 ## Prerequisites
 
-* GPU: NVIDIA Blackwell GPUs. DEP16 and TEP16 use 16 GPUs; the TEP8 recipe uses 8 GPUs. These deployment recipes were validated on GB300 NVL GPUs. The repository's Slurm examples assume 4 GPUs per node. Other GPU architectures are not currently supported.
+* GPU: NVIDIA Blackwell GPUs. DEP16 and TEP16 use 16 GPUs; the TEP8 recipe uses 8 GPUs. These deployment recipes were validated on GB300 NVL GPUs. The repository's Slurm examples assume 4 GPUs per node. Per-GPU memory requirements differ per recipe, and are set by the attention layout rather than by the GPU count:
+
+  | Recipe | Attention layout | Per-rank weights | Requires |
+  | :-- | :-- | --: | :-- |
+  | DEP16 | attention-DP (replicated) | 210 GB | GB300-class per-GPU memory |
+  | TEP8 | attention-TP, EP8 | 213 GB | GB300-class per-GPU memory |
+  | TEP16 | attention-TP, EP16 | 115 GB | validated on GB200 (`SM100`) |
+
+  DEP16 replicates the BF16 non-expert weights on every rank (114 GB per rank) on top of the MXFP4 routed experts at 16-way expert parallelism (90 GB per rank). TEP16 shards those non-expert weights instead, which is what brings it within `SM100` per-GPU memory; TEP8 does not fit because its 8-way expert share alone is 181 GB per rank. On B200 (`SM100`), Kimi K3 is functionally supported at the kernel and module level and covered by unit tests in CI. Other GPU architectures are not supported.
 * Multi-node launcher: Slurm with the pyxis/enroot container plugin (or an equivalent MPI launcher) to start one rank per GPU across the nodes.
 * High-speed inter-node interconnect (e.g., NVLink/InfiniBand) for the expert-parallel traffic.
 * Shared filesystem visible to all nodes for the repository, the model weights, and the configuration file.
@@ -26,7 +34,7 @@ The checkpoint and the configuration file must live on a shared filesystem visib
 
 ## Feature Support Notes
 
-* **Blackwell only.** NVIDIA Blackwell GPUs are supported. The configurations and results in this guide were validated on NVIDIA GB300 NVL GPUs. Support for other GPU architectures may be added in a future release.
+* **Blackwell only.** NVIDIA Blackwell GPUs are supported. The performance results in this guide were validated on NVIDIA GB300 NVL GPUs. Kimi K3 kernels and modules are also functional on B200 (`SM100`) and covered by unit tests in CI, and the TEP16 deployment is validated end-to-end on GB200 (`SM100`); DEP16 and TEP8 require GB300-class per-GPU memory (see Prerequisites). Support for other GPU architectures may be added in a future release.
 * **High-throughput and low-latency deployments are provided.** DEP16 (`enable_attention_dp: true`, `moe_expert_parallel_size: 16`) is the high-throughput deployment. TEP16 (`enable_attention_dp: false`, `moe_expert_parallel_size: 16`) is the low-latency deployment. An 8-GPU deployment, TEP8 (`enable_attention_dp: false`, `moe_expert_parallel_size: 8`), is also provided. Select the deployment and concurrency appropriate for your workload.
 * **CUDA graphs and the overlap scheduler are enabled.** The performance-sweep recipes set `disable_overlap_scheduler: false` and enable CUDA graphs. DEP16 additionally sets `cuda_graph_config.enable_padding: true`.
 * **Chunked prefill is supported and enabled** (`enable_chunked_prefill: true`), so prompts longer than `max_num_tokens` are scheduled across multiple steps.
@@ -44,7 +52,7 @@ python3 scripts/build_wheel.py --cuda_architectures 103-real --skip_building_whe
 .venv-3.12/bin/python -m pip install -e .
 ```
 
-`build_wheel.py` creates the virtual environment at the repository root, named after the container's Python version: `.venv-3.12` for the current containers (Python 3.12). If your container ships a different Python, substitute the matching `.venv-<major>.<minor>` path in the commands on this page. Adjust `--cuda_architectures` to the target GPUs (`103-real` for GB300). The multi-node jobs below run TensorRT LLM from this in-place environment, so build and install with the repository at the same path the jobs use.
+`build_wheel.py` creates the virtual environment at the repository root, named after the container's Python version: `.venv-3.12` for the current containers (Python 3.12). If your container ships a different Python, substitute the matching `.venv-<major>.<minor>` path in the commands on this page. Adjust `--cuda_architectures` to the target GPUs (`103-real` for GB300, `100-real` for B200). A `103-real` build also runs on B200 (the Kimi K3 kernels compile for the `100f` family) but omits the `sm100a`-specific batched-GEMM kernels, so build with `100-real` when targeting B200. The multi-node jobs below run TensorRT LLM from this in-place environment, so build and install with the repository at the same path the jobs use.
 
 Kimi K3 additionally depends on `fla` and `einops`, installed into the same in-place environment (these dependencies might be removed in future releases, replaced by other kernels):
 
@@ -271,6 +279,34 @@ The job writes progress and results to `kimi-k3-eval-<job-id>.log` in the submis
 
 The expected average accuracy is approximately 96.47. Small differences (roughly ±0.5 points) are possible with different checkpoint or dependency revisions.
 
+### TEP16 on GB200
+
+The same job runs the TEP16 layout with `--parallel tep`, which rewrites a per-job copy of the evaluation YAML with `enable_attention_dp: false` and raises `max_batch_size` from 32 to 128 (with attention-DP off every rank serves the same global batch instead of its own, so the batch size is raised to recover eval concurrency):
+
+```bash
+sbatch --account <account> --partition batch --qos <qos> --time 04:00:00 \
+    examples/kimi_k3/run_eval_kimi_k3.sbatch \
+    --model /path/to/kimi-k3-checkpoint \
+    --image /path/to/tensorrt-llm-container.sqsh \
+    --task gsm8k --parallel tep
+```
+
+The batch script declares `--nodes=4 --ntasks-per-node=4 --gpus-per-node=4`, and takes `--account`, `--partition` and `--qos` from the submitting command line. Export `KIMI_K3_ROUTER_BF16=0` before submitting: with attention-DP off the MoE router gate defaults to its BF16 fast path, which can flip borderline expert picks, so the reference scores above are only comparable with that path disabled. `KIMI_K3_FP8_WEIGHT_READ` defaults to `0`, which is the precision the reference scores were measured at.
+
+Measured on 16 GB200 GPUs (4 nodes, `100-real` build, 184.31 GiB per GPU), with the checkpoint's native MXFP4 routed experts:
+
+| Filter | Exact match |
+| :-- | --: |
+| Flexible extract | 96.82 |
+| Strict match | 96.74 |
+
+| Per-rank memory | Value |
+| :-- | --: |
+| Weights | 106.67 GiB |
+| Non-torch (NCCL, CUDA graphs) | 15.32 GiB |
+| Peak during profiling | 125.96 GiB |
+| KV cache at `free_gpu_memory_fraction: 0.25` | 15.60 GiB |
+
 ## Benchmarking Performance
 
 ### Run the End-to-End Performance Sweep

docs/source/models/supported-models.md

@@ -32,6 +32,8 @@ The following is a table of supported models for the PyTorch backend:
 | `HunYuanMoEForCausalLM` [^5]         | Hunyuan MoE                        | `tencent/Hunyuan-A13B-Instruct`              |
 | `InternLM3ForCausalLM` [^5]          | InternLM3                          | `internlm/internlm3-8b-instruct`             |
 | `KimiK25ForConditionalGeneration`    | Kimi-K2.5                          | `moonshotai/Kimi-K2.5`                       |
+| `KimiK3ForConditionalGeneration` [^15]| Kimi-K3                            | `moonshotai/Kimi-K3`                         |
+| `KimiLinearForCausalLM` [^15]        | Kimi-K3 (text decoder)             | `moonshotai/Kimi-K3`                         |
 | `LagunaForCausalLM`                  | Laguna-XS                          | `poolside/laguna-XS.2`                       |
 | `LlamaForCausalLM`                   | Llama 3.1, Llama 3, Llama 2, LLaMA | `meta-llama/Meta-Llama-3.1-70B`              |
 | `Llama4ForConditionalGeneration`     | Llama 4                            | `meta-llama/Llama-4-Scout-17B-16E-Instruct`  |
@@ -99,6 +101,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl
 [^12]: Supports text, image, and video inputs over the block-sparse attention path. The published MXFP8 checkpoint is dequantized on load so the runtime sees an effectively BF16 model. The text decoder is also usable standalone (text-only) via the `MiniMaxM3SparseForCausalLM` architecture. KV cache reuse and MTP are not supported on the sparse-attention path in this release.
 [^13]: The Cosmos 3 family also supports visual generation through the VisualGen API. See [Visual Generation Models](#visual-generation-models).
 [^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.
 
 # Encoder-Decoder Feature Support Matrix (PyTorch Backend)
 

examples/kimi_k3/README.md

@@ -5,8 +5,15 @@ start and configuration for GSM8K evaluation.
 
 ## Hardware support
 
-Only NVIDIA Blackwell GPUs are currently supported and tested. Support for
-other GPU architectures may be added in a future release.
+Only NVIDIA Blackwell GPUs (`SM100` family) are supported. The performance
+results in this example were measured on GB300 NVL (`SM103`). Per-GPU memory
+requirements differ per recipe, and are set by the attention layout rather
+than by the GPU count: DEP16 needs 210 GB per rank and TEP8 needs 213 GB per
+rank, so both require GB300-class per-GPU memory, while TEP16 needs 115 GB
+per rank and is validated end-to-end on GB200 (`SM100`) at 16 GPUs. B200
+(`SM100`) is functionally supported at the kernel and module level and
+covered by unit tests in CI. See Current limitations. Support for other GPU
+architectures may be added in a future release.
 
 ## Prerequisites
 
@@ -27,12 +34,14 @@ other GPU architectures may be added in a future release.
   for `.venv-3.12` in every command below and export
   `TRTLLM_VENV=/path/to/repo/.venv-<major>.<minor>` when submitting the
   Slurm jobs (they default to the repository-root `.venv-3.12`). Adjust
-  `--cuda_architectures` to the target GPUs (`103-real` for GB300).
+  `--cuda_architectures` to the target GPUs (`103-real` for GB300,
+  `100-real` for B200).
 - A complete Hugging Face-format Kimi K3 checkpoint and tokenizer, e.g.
   [moonshotai/Kimi-K3](https://huggingface.co/moonshotai/Kimi-K3) downloaded
   from the Hugging Face Hub (the example scripts take a local filesystem
   path).
-- A Slurm cluster with 16 NVIDIA Blackwell GPUs and a TensorRT-LLM container
+- A Slurm cluster with 16 NVIDIA Blackwell GPUs (GB300-class per-GPU memory
+  for DEP16 and TEP8; see Hardware support) and a TensorRT-LLM container
   image. The image passed as `--image` below must already provide
   TensorRT-LLM's runtime dependencies, that is, a release-style TensorRT-LLM
   container; a build or devel image without them does not work. The Slurm
@@ -214,6 +223,17 @@ decoding requires the default cache manager, which cannot reuse blocks.
 ## Current limitations
 
 - Pipeline parallelism is not supported.
+- **DEP16 and TEP8 require GB300-class per-GPU memory.** Under attention-DP
+  the BF16 non-expert weights are replicated on every rank (114 GB per rank),
+  which together with the MXFP4 routed experts at 16-way expert parallelism
+  (90 GB per rank) needs 210 GB per rank for DEP16; TEP8 needs 213 GB per
+  rank because its 8-way expert share alone is 181 GB. Neither fits `SM100`
+  per-GPU memory. Use TEP16, which shards the non-expert weights and needs
+  115 GB per rank — validated end-to-end on GB200 (16 GPUs, 4 nodes), see the
+  deployment guide. Note that the FP8 weight-read path (TRTLLM-14765) does
+  not lift the DEP16 requirement: the conversion runs after the weights are
+  already resident in BF16, so it lowers the steady-state footprint but
+  leaves the load-time peak unchanged.
 - **Known performance limitation at DEP16 saturation** (attention-DP +
   EP16 throughput recipe): the 8K/1K serving sweep loses several percent
   of output throughput at concurrency ≥ 128 (up to ~15% at concurrency

tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py

@@ -61,6 +61,34 @@ def _select_mla_generation_backend(quant_config: Optional[QuantConfig]) -> str:
     return backend
 
 
+def _validate_mla_generation_backend(backend: str, num_heads: int) -> None:
+    """Fail fast when `backend` can never run at this per-rank head count.
+
+    FlashInfer's `trtllm_batch_decode_with_kv_cache_mla` rejects
+    `64 < num_heads_q < 128` for every batch shape, and the per-batch policy
+    never demotes away from an explicit `trtllm-gen` selection (the
+    FP8-KV-cache override or a `TLLM_K3_MLA_GEN_BACKEND=trtllm-gen` request).
+    Without this check the conflict only surfaces as a FlashInfer error deep
+    in attention warmup.
+
+    The bound mirrors FlashInfer's validation gate verbatim: it predicts
+    FlashInfer's rejection, it is not a verified support claim for the head
+    counts outside the range. For Kimi K3 the open side above 128 is
+    unreachable anyway — per-rank Q heads never exceed 96 (all heads
+    replicated under attention-DP, `96 / tp_size` under TEP head sharding).
+    """
+    if backend == "trtllm-gen" and 64 < num_heads < 128:
+        raise ValueError(
+            "Kimi K3 MLA: the trtllm-gen generation backend cannot run with "
+            f"{num_heads} query heads per rank (trtllm-gen MLA decode rejects "
+            "64 < num_heads_q < 128; under attention-DP every rank keeps all "
+            "heads). trtllm-gen was selected explicitly — by the FP8-KV-cache "
+            f"override or by {_KIMI_K3_MLA_GEN_BACKEND_ENV}=trtllm-gen. Use "
+            "tensor-parallel head sharding (TEP) so each rank has <= 64 "
+            "heads, or a BF16 KV cache with the default cute-dsl backend."
+        )
+
+
 def _kimi_k3_mla_decode_backend_policy(
     requested_backend: str,
     metadata: TrtllmAttentionMetadata,
@@ -298,6 +326,11 @@ def __init__(
         # Only the absorbed-generation backend (mqa) requests CuTe-DSL, so
         # only it needs K3's per-batch fallback policy; mha keeps the
         # default trtllm-gen selection.
+        # Validate here rather than in _select_mla_generation_backend: the
+        # per-rank head count (replicated under attention-DP, sharded under
+        # TEP) is only authoritative once the base MLA module has built its
+        # generation backend.
+        _validate_mla_generation_backend(self.mqa.flashinfer_mla_backend, self.mqa.num_heads)
         self.mqa.mla_backend_policy = partial(
             _kimi_k3_mla_decode_backend_policy,
             num_heads=self.mqa.num_heads,

tests/integration/test_lists/test-db/l0_b200.yml

@@ -106,12 +106,23 @@ l0_b200:
   - unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py
   - unittest/_torch/modules/kimi_kda/test_kda_decode_op.py
   - unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py
+  # CuTe MTP verify kernel parity (trtllm::kda_mtp_decode).
+  - unittest/_torch/modules/kimi_kda/test_kda_mtp_decode_cute_parity.py
   - unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py
   # GPU KDA disagg transfer + peer-validation (cpu_only cases skipped by the
   # stage's "not cpu_only" markexpr and run on the CPU-Generic stage instead).
   - unittest/disaggregated/test_kda_mamba_transfer.py
   - unittest/_torch/moe/test_kimi_k3_mlp.py
   - unittest/_torch/moe/test_kimi_k3_situ_moe.py
+  # Kimi K3 MLA backend selection + config/checkpoint plumbing (monkeypatched,
+  # HW-light). Collected nowhere else: the CPU-Generic stage only collects files
+  # carrying the literal cpu_only marker, and these carry none.
+  - unittest/_torch/modules/test_kimi_k3_mla_backend.py
+  - unittest/_torch/modeling/test_kimi_k3_config_routing.py
+  - unittest/_torch/modeling/test_kimi_linear_checkpoint.py
+  - unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py
+  # CPU self-test of the Kimi K3 disagg parity harness comparison logic.
+  - test_kimi_k3_specdec.py::test_kimi_k3_disagg_parity_selftest
   - unittest/_torch/custom_ops/test_deepseek_v4_q_norm.py TIMEOUT (15)
   # ------------- modules (non-MoE) ---------------
   - unittest/_torch/modules/test_mla_helix.py
@@ -344,6 +355,9 @@ l0_b200:
   # Covered by H100 pre_merge for primary HW-agnostic signal; keep B200 runtime
   # canary in post_merge for CUDA IPC / virtual memory / profiling paths.
   - unittest/_torch/misc
+  # Kimi K3 MoE routing shape on the TRTLLM backend (random weights, no
+  # checkpoint); ~13 min, too heavy for pre_merge. Mirrors l0_gb300_multi_gpus.
+  - unittest/_torch/moe/test_moe_backend.py -k "kimi_k3_shape"
   - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_w4a8_mxfp4[fp8-latency]
   - accuracy/test_llm_api_pytorch.py::TestQwen3_8B::test_w4a8_mxfp4[mxfp8-latency]
   - accuracy/test_llm_api_pytorch.py::TestQwen3_30B_A3B::test_nvfp4[latency_moe_cutlass-torch_compile=False]

tests/integration/test_lists/test-db/l0_dgx_b200.yml

@@ -362,3 +362,7 @@ l0_dgx_b200:
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV32::test_nvfp4_attn_multi_gpus TIMEOUT (60)
   # ---- FP8 per-tensor (QDQ) moved to post-merge; block-scale/W4A8 stay in pre ----
   - unittest/_torch/moe/test_moe_module.py::test_configurable_moe_multi_gpu -k "CUTLASS and FP8 and not FP8_BLOCK_SCALES and not W4A8 and not MXFP8"
+  # Kimi K3 Tier-1: truncated 4-layer (KDA + first full-attention) SA spec-dec
+  # logits parity on TP=4 (skips cleanly when the Kimi-K3 checkpoint is not
+  # staged under LLM_MODELS_ROOT).
+  - test_kimi_k3_specdec.py::test_kimi_k3_sa_specdec_logits_parity TIMEOUT (45)