From d08f7fa7f039c717618f5a3406a0daebfa887d2a Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Mon, 24 Aug 2026 17:03:47 -0700 Subject: [PATCH 1/7] [TRTLLM-14705][test] Wire Kimi K3 unit tests and Tier-1 logits parity into B200 L0 lists l0_b200 pre-merge: add the fused CUDA decode op and CuTe MTP verify parity suites (previously GB300-only; SM103-only selector cases self-skip on B200), plus the Kimi K3 MLA-backend/config/checkpoint plumbing tests and the disagg-parity CPU selftest, none of which were collected by any stage. l0_b200 post-merge: the ~13 min test_moe_backend kimi_k3_shape case, mirroring l0_gb300_multi_gpus. l0_dgx_b200 4-GPU post-merge: the Tier-1 truncated-checkpoint SA spec-dec logits-parity test (TP=4). All entries validated green on an 8xB200 x86 node with a 100-real build. Signed-off-by: Michal Guzek --- tests/integration/test_lists/test-db/l0_b200.yml | 16 ++++++++++++++++ .../test_lists/test-db/l0_dgx_b200.yml | 4 ++++ 2 files changed, 20 insertions(+) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 4a217b4294c5..912f9532f02f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -106,12 +106,25 @@ 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 + # Fused CUDA decode op and CuTe MTP verify kernel (SM103-only selector cases + # self-skip on B200). + - unittest/_torch/modules/kimi_kda/test_kda_decode_op.py + - 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 +357,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] diff --git a/tests/integration/test_lists/test-db/l0_dgx_b200.yml b/tests/integration/test_lists/test-db/l0_dgx_b200.yml index fb4bf93e7a9b..405a83d4db58 100644 --- a/tests/integration/test_lists/test-db/l0_dgx_b200.yml +++ b/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) From 2b4a6ae6139772b81ae4bf73114afec6ac858f6d Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Mon, 24 Aug 2026 17:03:59 -0700 Subject: [PATCH 2/7] [TRTLLM-14705][doc] Document Kimi K3 B200 support and limitations Add KimiK3ForConditionalGeneration / KimiLinearForCausalLM to the supported-models matrix with a Blackwell footnote; state B200 (SM100) kernel/module-level support and the full-model memory limitation in the deployment guide and example README; add 100-real build guidance for B200 targets. Signed-off-by: Michal Guzek --- .../deployment-guide-for-kimi-k3-on-trtllm.md | 6 +++--- docs/source/models/supported-models.md | 3 +++ examples/kimi_k3/README.md | 21 +++++++++++++++---- 3 files changed, 23 insertions(+), 7 deletions(-) diff --git a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md index 2e234659c34c..d76a1551fe8e 100644 --- a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md @@ -10,7 +10,7 @@ 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. These recipes require GB300-class per-GPU memory: on B200 (`SM100`), Kimi K3 is functionally supported at the kernel and module level, but the deployments in this guide exceed the 180 GB per-GPU memory (the BF16 non-expert weights take ~109 GB per rank, plus ~87 GB per rank of MXFP4 routed experts at 16-way expert parallelism); full-model serving on B200 would require a wider expert-parallel group (>=32 GPUs, not validated). 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 +26,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 configurations and 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; the full-model deployments in this guide exceed B200 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 +44,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-.` 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-.` 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): diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 7e8fb78e6e3a..947911993643 100644 --- a/docs/source/models/supported-models.md +++ b/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). Full-model deployment recipes are validated on GB300 NVL (`SM103`); see the [Kimi K3 deployment guide](../deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md). On B200 (`SM100`), kernel and module support is functional and covered by CI, but full-model serving is not supported at 16-way expert parallelism: the BF16 non-expert weights (~109 GB per rank) plus the MXFP4 routed experts (~87 GB per rank) exceed the 180 GB per-GPU memory. A wider expert-parallel group (>=32 GPUs, not validated) or the default-off FP8 weight-read path (TRTLLM-14765) would be required. # Encoder-Decoder Feature Support Matrix (PyTorch Backend) diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index 14d6023cd713..9c809b225cbe 100644 --- a/examples/kimi_k3/README.md +++ b/examples/kimi_k3/README.md @@ -5,8 +5,12 @@ 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 full-model +16-GPU deployments in this example were validated on GB300 NVL (`SM103`). +B200 (`SM100`) is functionally supported at the kernel and module level and +covered by unit tests in CI; full-model serving on B200 is not supported at +16-way expert parallelism (see Current limitations). Support for other GPU +architectures may be added in a future release. ## Prerequisites @@ -27,12 +31,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-.` 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; + 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 +220,13 @@ decoding requires the default cache manager, which cannot reuse blocks. ## Current limitations - Pipeline parallelism is not supported. +- **B200 full-model serving is not supported.** Kimi K3 kernels and modules + are functional on B200 (`SM100`) and covered by unit tests in CI, but at + 16-way expert parallelism the BF16 non-expert weights (~109 GB per rank) + plus the MXFP4 routed experts (~87 GB per rank) exceed the 180 GB per-GPU + memory. Serving would require a wider expert-parallel group (>=32 GPUs, + not validated) or the default-off FP8 weight-read path, tracked as + TRTLLM-14765. - **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 From ef04658387ab707ce66cedcb37c760464ba8a7fc Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Mon, 24 Aug 2026 21:56:48 -0700 Subject: [PATCH 3/7] [TRTLLM-14705][fix] Fail fast when Kimi K3 MLA selects trtllm-gen with an unsupported head count MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The FP8-KV-cache override (and an explicit TLLM_K3_MLA_GEN_BACKEND request) select trtllm-gen before the per-batch policy runs, and the policy passes explicit selections through — so with 64 < per-rank heads < 128 (e.g. 96 under attention-DP) every generation batch would crash inside FlashInfer at attention warmup, since no batch shape can run trtllm-gen at those head counts and CuTe-DSL cannot serve FP8 KV device scales either. Validate at module construction, where the per-rank head count is authoritative, and raise a targeted error naming the conflict and the remedies (TEP sharding or BF16 KV cache). Addresses the CodeRabbit review finding on PR #18164. Signed-off-by: Michal Guzek --- .../kimi_k3_mla/kimi_k3_mla_attention.py | 27 +++++++++++++ .../modules/test_kimi_k3_mla_backend.py | 39 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py index 03bcc9004801..5b229c772078 100644 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py +++ b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py @@ -61,6 +61,28 @@ 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. + """ + 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 +320,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, diff --git a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py index 67ec3ea13a4b..12a206e1e541 100644 --- a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py +++ b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py @@ -10,6 +10,7 @@ _KIMI_K3_MLA_GEN_BACKEND_ENV, _kimi_k3_mla_decode_backend_policy, _select_mla_generation_backend, + _validate_mla_generation_backend, ) from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -53,6 +54,44 @@ def test_select_kimi_k3_mla_generation_backend_uses_trtllm_gen_for_fp8_kv_cache( assert _select_mla_generation_backend(quant_config) == "trtllm-gen" +@pytest.mark.parametrize( + ("backend", "num_heads"), + [ + ("cute-dsl", 96), + ("trtllm-gen", 6), + ("trtllm-gen", 64), + ("trtllm-gen", 128), + ], +) +def test_validate_mla_generation_backend_accepts_runnable_configs( + backend: str, num_heads: int +) -> None: + _validate_mla_generation_backend(backend, num_heads) + + +@pytest.mark.parametrize("num_heads", [65, 96, 127]) +def test_validate_mla_generation_backend_rejects_trtllm_gen_mid_head_counts( + num_heads: int, +) -> None: + """trtllm-gen with 64 < H < 128 per-rank heads can never run any batch.""" + with pytest.raises(ValueError, match="query heads per rank"): + _validate_mla_generation_backend("trtllm-gen", num_heads) + + +def test_fp8_kv_cache_with_attention_dp_head_count_fails_fast( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The FP8-KV override forces trtllm-gen; with 96 per-rank heads + (attention-DP replication) construction must fail fast rather than + crash inside FlashInfer at attention warmup.""" + monkeypatch.delenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, raising=False) + quant_config = QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8) + + backend = _select_mla_generation_backend(quant_config) + with pytest.raises(ValueError, match="FP8-KV-cache"): + _validate_mla_generation_backend(backend, num_heads=96) + + @pytest.mark.parametrize( ( "requested_backend", From ce4926b55a4888926f560938d5369194cafdc01d Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Tue, 25 Aug 2026 10:03:04 -0700 Subject: [PATCH 4/7] [TRTLLM-14705][test] Add constructor-level FP8-KV attention-DP fail-fast regression test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Construct a real KimiK3MLAAttention (skip_create_weights_in_init, K3 geometry, FP8 KV cache, attention-DP mapping) and assert that __init__ itself raises — pinning the _validate_mla_generation_backend call site, not just the helper. The helper-level tests would stay green if the constructor stopped invoking the validation. Addresses the CodeRabbit coverage finding on PR #18164. Signed-off-by: Michal Guzek --- .../modules/test_kimi_k3_mla_backend.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py index 12a206e1e541..33f32cf693c5 100644 --- a/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py +++ b/tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py @@ -5,13 +5,17 @@ from typing import Optional import pytest +import torch +from tensorrt_llm._torch.model_config import ModelConfig +from tensorrt_llm._torch.modules.kimi_k3_mla import KimiK3MLAAttention from tensorrt_llm._torch.modules.kimi_k3_mla.kimi_k3_mla_attention import ( _KIMI_K3_MLA_GEN_BACKEND_ENV, _kimi_k3_mla_decode_backend_policy, _select_mla_generation_backend, _validate_mla_generation_backend, ) +from tensorrt_llm.mapping import Mapping from tensorrt_llm.models.modeling_utils import QuantAlgo, QuantConfig @@ -92,6 +96,43 @@ def test_fp8_kv_cache_with_attention_dp_head_count_fails_fast( _validate_mla_generation_backend(backend, num_heads=96) +@pytest.mark.skipif( + not torch.cuda.is_available(), + reason="KimiK3MLAAttention builds real TRTLLM attention backends", +) +def test_kimi_k3_mla_construction_fails_fast_for_fp8_kv_with_attention_dp( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Constructor-level regression for the FP8-KV + attention-DP conflict: + the override selects trtllm-gen, attention-DP keeps all 96 query heads + on every rank, and `KimiK3MLAAttention.__init__` itself must raise — + before attention warmup ever runs.""" + monkeypatch.delenv(_KIMI_K3_MLA_GEN_BACKEND_ENV, raising=False) + model_config = ModelConfig( + skip_create_weights_in_init=True, + quant_config=QuantConfig(kv_cache_quant_algo=QuantAlgo.FP8), + mapping=Mapping(world_size=4, tp_size=4, rank=0, enable_attention_dp=True), + ) + + with pytest.raises(ValueError, match="query heads per rank"): + # Kimi K3 MLA geometry (96 Q heads); small max positions keep the + # identity-RoPE table allocation negligible. + KimiK3MLAAttention( + hidden_size=7168, + num_heads=96, + q_lora_rank=1536, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128, + rms_norm_eps=1e-6, + dtype=torch.bfloat16, + layer_idx=0, + max_position_embeddings=256, + model_config=model_config, + ) + + @pytest.mark.parametrize( ( "requested_backend", From 6e802172a7e74892f6ec34935531a2e4fd46d559 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Tue, 25 Aug 2026 10:03:32 -0700 Subject: [PATCH 5/7] [TRTLLM-14705][chore] Clarify that the MLA head-count gate mirrors FlashInfer, not a support matrix MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 64 < num_heads_q < 128 bound in the decode policy and the construction-time validation copies FlashInfer's own validation gate verbatim — it predicts FlashInfer's rejection and makes no support claim for head counts outside the range. Note the Kimi K3 invariant that makes the open upper side unreachable: per-rank Q heads never exceed 96 (all heads replicated under attention-DP, 96 / tp_size under TEP head sharding). Addresses zhaoyangwang-nvidia's review question on PR #18164. Comment-only change. Signed-off-by: Michal Guzek --- .../_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py index 5b229c772078..7e1682366863 100644 --- a/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py +++ b/tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py @@ -70,6 +70,12 @@ def _validate_mla_generation_backend(backend: str, num_heads: int) -> None: 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( From 39cae851bd49e6cc605b151d9c702205b989c291 Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Fri, 28 Aug 2026 10:59:19 -0700 Subject: [PATCH 6/7] [TRTLLM-15786][doc] Correct Kimi K3 SM100 memory limits; TEP16 validated on GB200 The B200 limitation text added earlier in this PR states that full-model serving needs ">=32 GPUs or the default-off FP8 weight-read path (TRTLLM-14765)". Measurements on GB200 (SM100) under TRTLLM-15786 show both remedies are wrong, and that the real constraint is the attention layout rather than the GPU count. - TEP16 (enable_attention_dp: false) shards the non-expert weights instead of replicating them and needs 115 GB per rank, so it fits SM100 per-GPU memory at the same 16 GPUs. Validated end-to-end on 16 GB200 GPUs: 106.67 GiB of weights, 125.96 GiB peak, and GSM8K 96.82 flexible / 96.74 strict against the GB300 reference of 96.51 / 96.44. - DEP16 (210 GB per rank) and TEP8 (213 GB per rank, its 8-way expert share alone being 181 GB) are the recipes that require GB300-class memory. The DEP16 conclusion is unchanged; only the stated remedy is. - The FP8 weight-read path does not lift the DEP16 requirement. _finalize_weight_load runs after _load_trunk_params and _load_expert_slices, so the conversion happens once every parameter is already resident in BF16: it lowers the steady-state footprint but leaves the load-time peak unchanged. This is documented in the repo already, in the empty_placeholder docstring in modeling_kimi_linear.py. Also corrects the per-rank figures. Parsing all 96 safetensors shards gives 114 GB of replicated non-expert weights and 90 GB of routed experts at EP16, not ~109 GB and ~87 GB; those two values sum to the correct total but are labelled GB where the underlying numbers are GiB. Adds the TEP16 reproduction to the deployment guide's accuracy section, including the two environment switches that are not obvious from the batch script: KIMI_K3_ROUTER_BF16=0 (with attention-DP off the MoE router gate takes a BF16 fast path that can flip borderline expert picks, so scores are only comparable to the DEP16 reference with it disabled) and the KIMI_K3_FP8_WEIGHT_READ=0 default that the published numbers were measured at. run_eval_kimi_k3.sbatch takes --account, --partition and --qos from the submitting command line. Docs only; no functional change. Signed-off-by: Michal Guzek --- .../deployment-guide-for-kimi-k3-on-trtllm.md | 40 ++++++++++++++++++- docs/source/models/supported-models.md | 2 +- examples/kimi_k3/README.md | 35 +++++++++------- 3 files changed, 60 insertions(+), 17 deletions(-) diff --git a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md index d76a1551fe8e..4d543f9deee2 100644 --- a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md +++ b/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. These recipes require GB300-class per-GPU memory: on B200 (`SM100`), Kimi K3 is functionally supported at the kernel and module level, but the deployments in this guide exceed the 180 GB per-GPU memory (the BF16 non-expert weights take ~109 GB per rank, plus ~87 GB per rank of MXFP4 routed experts at 16-way expert parallelism); full-model serving on B200 would require a wider expert-parallel group (>=32 GPUs, not validated). Other GPU architectures are not 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. Kimi K3 kernels and modules are also functional on B200 (`SM100`) and covered by unit tests in CI; the full-model deployments in this guide exceed B200 per-GPU memory (see Prerequisites). 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. @@ -271,6 +279,34 @@ The job writes progress and results to `kimi-k3-eval-.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 --partition batch --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 diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 947911993643..1d21af9b0a69 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -101,7 +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). Full-model deployment recipes are validated on GB300 NVL (`SM103`); see the [Kimi K3 deployment guide](../deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md). On B200 (`SM100`), kernel and module support is functional and covered by CI, but full-model serving is not supported at 16-way expert parallelism: the BF16 non-expert weights (~109 GB per rank) plus the MXFP4 routed experts (~87 GB per rank) exceed the 180 GB per-GPU memory. A wider expert-parallel group (>=32 GPUs, not validated) or the default-off FP8 weight-read path (TRTLLM-14765) would be required. +[^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) diff --git a/examples/kimi_k3/README.md b/examples/kimi_k3/README.md index 9c809b225cbe..27f75b8d2b70 100644 --- a/examples/kimi_k3/README.md +++ b/examples/kimi_k3/README.md @@ -5,11 +5,14 @@ start and configuration for GSM8K evaluation. ## Hardware support -Only NVIDIA Blackwell GPUs (`SM100` family) are supported. The full-model -16-GPU deployments in this example were validated on GB300 NVL (`SM103`). -B200 (`SM100`) is functionally supported at the kernel and module level and -covered by unit tests in CI; full-model serving on B200 is not supported at -16-way expert parallelism (see Current limitations). Support for other GPU +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 @@ -37,8 +40,8 @@ architectures may be added in a future release. [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 (GB300-class per-GPU memory; - see Hardware support) 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 @@ -220,13 +223,17 @@ decoding requires the default cache manager, which cannot reuse blocks. ## Current limitations - Pipeline parallelism is not supported. -- **B200 full-model serving is not supported.** Kimi K3 kernels and modules - are functional on B200 (`SM100`) and covered by unit tests in CI, but at - 16-way expert parallelism the BF16 non-expert weights (~109 GB per rank) - plus the MXFP4 routed experts (~87 GB per rank) exceed the 180 GB per-GPU - memory. Serving would require a wider expert-parallel group (>=32 GPUs, - not validated) or the default-off FP8 weight-read path, tracked as - TRTLLM-14765. +- **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 From 6d57e79c73a44dfd2070400292af280fb95e340d Mon Sep 17 00:00:00 2001 From: Michal Guzek Date: Fri, 28 Aug 2026 13:50:22 -0700 Subject: [PATCH 7/7] [TRTLLM-14705][test] Drop duplicate test_kda_decode_op.py L0 entry Upstream PR #17870 wired test_kda_decode_op.py into the same l0_b200 pre-merge block; the rebase auto-merged both copies without a textual conflict and trt-test-db's L0 verification rejects the duplicate test/condition pair. Keep upstream's entry; ours now adds only the CuTe MTP verify parity suite. Signed-off-by: Michal Guzek --- tests/integration/test_lists/test-db/l0_b200.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 912f9532f02f..a3e8504ad21f 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -106,9 +106,7 @@ 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 - # Fused CUDA decode op and CuTe MTP verify kernel (SM103-only selector cases - # self-skip on B200). - - unittest/_torch/modules/kimi_kda/test_kda_decode_op.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