From 4755580fe2a85ac2296da1dc01cf60f723ff943d Mon Sep 17 00:00:00 2001 From: Guiju Zhang <7135567+cascade812@users.noreply.github.com> Date: Tue, 1 Sep 2026 22:25:28 -0700 Subject: [PATCH 1/4] kimi feat matrix Signed-off-by: Guiju Zhang <7135567+cascade812@users.noreply.github.com> --- docs/source/models/supported-models.md | 2 + tests/integration/defs/accuracy/test_kimi3.py | 322 ++++++++++++++++++ .../defs/accuracy/test_llm_api_pytorch.py | 101 ------ .../test_lists/qa/llm_function_multinode.txt | 6 +- 4 files changed, 327 insertions(+), 104 deletions(-) create mode 100644 tests/integration/defs/accuracy/test_kimi3.py diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 3b8df58bfe7f..847dc1e36adb 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -82,6 +82,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl | `Qwen3_5MoeForCausalLM` | Yes | Yes | Yes | Yes | Yes | Yes | No | No | No | Yes | Untested | Yes | N/A | Untested | Untested | | `Llama4ForConditionalGeneration` | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | No | Yes | Yes | Untested | N/A | Yes | Yes | | `GptOssForCausalLM` | Yes | Yes | Yes | Yes | Yes | No | Yes | No | Yes | Yes | Yes | Yes | N/A | Yes | Yes | +| `KimiK3ForConditionalGeneration` [^15] [^17] | Yes | Yes | Yes | Yes | Yes | No | No | No | No | Yes | No | Yes | N/A | Yes | Yes | | `Glm4MoeLiteForCausalLM` [^5] | Yes | Yes | Untested | Untested | Yes | No | No | No | No | Yes | Untested | Untested | N/A | Untested | Untested | | `NemotronHForCausalLM` | Yes | Yes | Yes | Yes | Yes | Yes | No | No | No | Yes | Yes | Yes | N/A | Untested | Untested | | `Gemma4ForConditionalGeneration` | Untested | Yes | Untested | No | Yes | Yes | No | No | No | Yes | Untested | No | Yes | Untested | Untested | @@ -103,6 +104,7 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl [^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. +[^16]: Kimi K3 supports suffix-automaton and DSpark speculative decoding. It has no MTP or EAGLE-3 head, and its DSpark checkpoints are not compatible with the plain `DFlash` mode represented by this matrix. # Encoder-Decoder Feature Support Matrix (PyTorch Backend) diff --git a/tests/integration/defs/accuracy/test_kimi3.py b/tests/integration/defs/accuracy/test_kimi3.py new file mode 100644 index 000000000000..dbc2d01a4b65 --- /dev/null +++ b/tests/integration/defs/accuracy/test_kimi3.py @@ -0,0 +1,322 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re + +import pytest +import torch + +from tensorrt_llm import LLM +from tensorrt_llm._torch.configs import KimiK3Config, KimiLinearConfig +from tensorrt_llm._torch.pyexecutor.config_utils import load_pretrained_config +from tensorrt_llm.llmapi import ( + CudaGraphConfig, + KvCacheConfig, + MambaStateConfig, + MoeConfig, + SADecodingConfig, + SamplingParams, + SchedulingParams, +) +from tensorrt_llm.sampling_params import GuidedDecodingParams, LogitsProcessor + +from ..conftest import llm_models_root, skip_pre_blackwell +from .accuracy_core import GSM8K, LlmapiAccuracyTestHarness, assert_acceptance_length + + +class _ForceTokenLogitsProcessor(LogitsProcessor): + def __init__(self, forced_token_id: int) -> None: + self._forced_token_id = forced_token_id + + def __call__( + self, + req_id: int, + logits: torch.Tensor, + token_ids: list[list[int]], + stream_ptr: int | None, + client_id: int | None, + ) -> None: + del req_id, token_ids, client_id + if stream_ptr is None: + self._force_token(logits) + return + with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)): + self._force_token(logits) + + def _force_token(self, logits: torch.Tensor) -> None: + logits.fill_(float("-inf")) + logits[..., self._forced_token_id] = 0 + + +def _compute_acceptance_length(llm: LLM) -> float: + stats = llm.get_stats(timeout=2) + spec_iterations = [ + stat["specDecodingStats"] + for stat in stats + if stat.get("specDecodingStats") and stat["specDecodingStats"]["numDraftTokens"] > 0 + ] + assert spec_iterations, "No iterations with speculative decoding stats" + accepted = sum(stat["numAcceptedTokens"] for stat in spec_iterations) + requests = sum(stat["numRequestsWithDraftTokens"] for stat in spec_iterations) + return (accepted + requests) / requests + + +@pytest.mark.timeout(10800) +class TestKimiK3(LlmapiAccuracyTestHarness): + MODEL_NAME = "moonshotai/Kimi-K3" + MODEL_PATH = f"{llm_models_root()}/Kimi-K3" + + @skip_pre_blackwell + @pytest.mark.skip_less_mpi_world_size(16) + # The 16-GPU K3 recipes are qualified on GB300 (one NVL72 domain) only: + # on 2-node 180-190 GiB parts (B200/GB200, InfiniBand between nodes) the + # EP16 MoE-comm bring-up hangs and the KV-budget assumptions do not hold, + # so gate on GB300-class device memory. B300 clears this memory gate but + # pairs 8-GPU nodes over InfiniBand (same non-NVL72 topology) -- do not + # schedule these tests on B300; that exclusion is enforced by QA's + # platform selection, not by this marker. + @pytest.mark.skip_less_device_memory(200000) + @pytest.mark.parametrize("mode", ["baseline", "reuse", "sa"]) + def test_w4a16_mxfp4(self, mode: str, monkeypatch: pytest.MonkeyPatch) -> None: + """Verify Kimi K3 accuracy and its key model-feature matrix entries. + + The three modes retain the existing full-GSM8K coverage. The baseline + mode also exercises attention DP, the overlap scheduler, CUDA graphs, + chunked prefill, Torch sampling, a logits processor, and guided + decoding. The reuse mode requires an observed hybrid-cache hit. The SA + mode remains an acceptance-length guard; SA is supported but is not a + column in the model-feature matrix. + """ + if mode == "baseline": + monkeypatch.setenv("TLLM_METRICS_ALL_RANKS", "1") + monkeypatch.setenv("TRTLLM_XGUIDANCE_LENIENT", "1") + self._assert_checkpoint_routing() + + kv_cache_kwargs = dict( + enable_block_reuse=False, + free_gpu_memory_fraction=0.25, + # tokens_per_block=64 keeps the MLA generation path on the + # flashinfer trtllm-gen kernel (K3 has 96 query heads). + tokens_per_block=64, + ) + llm_kwargs = dict( + tensor_parallel_size=16, + moe_expert_parallel_size=16, + enable_attention_dp=True, + max_batch_size=32, + max_num_tokens=8192, + max_seq_len=8192, + trust_remote_code=True, + enable_chunked_prefill=True, + cuda_graph_config=CudaGraphConfig(enable_padding=True, max_batch_size=32), + moe_config=MoeConfig( + max_num_tokens=33024, + use_low_precision_moe_combine=True, + ), + enable_iter_perf_stats=True, + max_stats_len=256, + return_perf_metrics=True, + ) + if mode == "baseline": + llm_kwargs["guided_decoding_backend"] = "xgrammar" + elif mode == "reuse": + kv_cache_kwargs["enable_block_reuse"] = True + # Hybrid models expose reusable prefixes only at KDA state + # snapshot boundaries; without a snapshot cadence, block reuse + # silently never engages. + kv_cache_kwargs["mamba_state_config"] = MambaStateConfig(periodic_snapshot_interval=256) + else: + llm_kwargs.update( + max_batch_size=8, + disable_overlap_scheduler=True, + enable_chunked_prefill=False, + cuda_graph_config=CudaGraphConfig(max_batch_size=8), + speculative_config=SADecodingConfig(max_draft_len=2), + max_stats_len=-1, + ) + # Log corpus-aggregate acceptance length and rate at evaluation + # end. QA records these values from the test log. + monkeypatch.setenv("TLLM_EVAL_SPEC_STATS", "1") + + with LLM( + self.MODEL_PATH, + kv_cache_config=KvCacheConfig(**kv_cache_kwargs), + **llm_kwargs, + ) as llm: + self._assert_resolved_args(llm, mode) + if mode == "baseline": + self._assert_attention_dp(llm) + self._assert_logits_processor(llm) + self._assert_guided_decoding(llm) + elif mode == "reuse": + self._assert_kv_cache_reuse(llm) + + task = GSM8K(self.MODEL_NAME) + task.evaluate(llm) + if mode == "sa": + acceptance_length = _compute_acceptance_length(llm) + print( + "[AL] TestKimiK3::test_w4a16_mxfp4[sa] " + f"acceptance_length = {acceptance_length:.3f}" + ) + assert_acceptance_length( + "TestKimiK3::test_w4a16_mxfp4", + acceptance_length, + ) + + def _assert_checkpoint_routing(self) -> None: + config = load_pretrained_config(self.MODEL_PATH, trust_remote_code=True) + assert isinstance(config, KimiK3Config) + assert config.architectures == ["KimiK3ForConditionalGeneration"] + assert isinstance(config.text_config, KimiLinearConfig) + assert config.text_config.model_type == "kimi_linear" + assert config.text_config.num_nextn_predict_layers == 0 + assert getattr(config.text_config, "sliding_window", None) is None + + @staticmethod + def _assert_resolved_args(llm: LLM, mode: str) -> None: + assert llm.args.enable_attention_dp is True + assert llm.args.kv_cache_config.enable_block_reuse is (mode == "reuse") + # K3's routed-expert quantization is nested in the composite checkpoint + # and is not represented by the modelopt-style args quant_algo field. + assert llm.args.quant_config.quant_algo is None + if mode == "sa": + assert llm.args.disable_overlap_scheduler is True + assert llm.args.enable_chunked_prefill is False + else: + assert llm.args.disable_overlap_scheduler is False + assert llm.args.enable_chunked_prefill is True + assert llm.args.cuda_graph_config is not None + assert llm.args.cuda_graph_config.enable_padding is True + + @staticmethod + def _assert_attention_dp(llm: LLM) -> None: + num_ranks = 16 + output_length = 8 + prompts = [[1, 42, 43] for _ in range(num_ranks)] + scheduling_params = [ + SchedulingParams(attention_dp_rank=rank, attention_dp_relax=False) + for rank in range(num_ranks) + ] + outputs = llm.generate( + prompts, + sampling_params=SamplingParams( + max_tokens=output_length, + temperature=0, + end_id=-1, + ), + scheduling_params=scheduling_params, + use_tqdm=False, + ) + stats_entries = llm.get_stats(timeout=10) + + assert isinstance(outputs, list) + assert len(outputs) == num_ranks + reference_token_ids = outputs[0].outputs[0].token_ids + for rank, output in enumerate(outputs): + token_ids = output.outputs[0].token_ids + assert len(token_ids) == output_length + assert token_ids == reference_token_ids, ( + f"ADP rank {rank} produced tokens that differ from rank 0" + ) + + assert stats_entries + for rank in range(num_ranks): + rank_stats = [ + entry + for entry in stats_entries + if entry.get("attentionDpRank", entry.get("rank")) == rank + ] + assert rank_stats, f"No iteration statistics reported for ADP rank {rank}" + inflight_batching_stats = [ + entry.get("inflightBatchingStats", {}) for entry in rank_stats + ] + assert sum(stats.get("numContextRequests", 0) for stats in inflight_batching_stats) == 1 + assert ( + max( + (stats.get("numGenRequests", 0) for stats in inflight_batching_stats), + default=0, + ) + == 1 + ) + + @staticmethod + def _assert_kv_cache_reuse(llm: LLM) -> None: + prompt_token_ids = [1] + [42] * 510 + [43] + output_length = 8 + sampling_params = SamplingParams( + max_tokens=output_length, + temperature=0, + end_id=-1, + return_perf_metrics=True, + ) + scheduling_params = [SchedulingParams(attention_dp_rank=0, attention_dp_relax=False)] + + cold_output = llm.generate( + [prompt_token_ids], + sampling_params=sampling_params, + scheduling_params=scheduling_params, + use_tqdm=False, + )[0].outputs[0] + warm_output = llm.generate( + [prompt_token_ids], + sampling_params=sampling_params, + scheduling_params=scheduling_params, + use_tqdm=False, + )[0].outputs[0] + + cold_metrics = cold_output.request_perf_metrics + warm_metrics = warm_output.request_perf_metrics + assert cold_metrics is not None + assert warm_metrics is not None + assert cold_metrics.kv_cache_metrics.num_reused_blocks == 0 + assert warm_metrics.kv_cache_metrics.num_reused_blocks > 0 + assert len(cold_output.token_ids) == output_length + assert warm_output.token_ids == cold_output.token_ids + + @staticmethod + def _assert_logits_processor(llm: LLM) -> None: + forced_token_id = 22 + output_length = 4 + outputs = llm.generate( + [[1, 42, 43]], + sampling_params=SamplingParams( + max_tokens=output_length, + temperature=0, + end_id=-1, + logits_processor=_ForceTokenLogitsProcessor(forced_token_id), + ), + use_tqdm=False, + ) + + assert isinstance(outputs, list) + assert len(outputs) == 1 + assert outputs[0].outputs[0].token_ids == [forced_token_id] * output_length + + @staticmethod + def _assert_guided_decoding(llm: LLM) -> None: + prompt_token_ids = llm.tokenizer.encode("Return exactly two decimal digits:") + outputs = llm.generate( + [prompt_token_ids], + sampling_params=SamplingParams( + max_tokens=8, + guided_decoding=GuidedDecodingParams(regex=r"[0-9]{2}"), + ), + use_tqdm=False, + ) + + assert isinstance(outputs, list) + assert len(outputs) == 1 + assert re.fullmatch(r"[0-9]{2}", outputs[0].outputs[0].text) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index a2133a9a0027..f56bdc4c6c94 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -4639,107 +4639,6 @@ def test_nvfp4(self, ep_size, attention_dp): task.evaluate(llm) -@pytest.mark.timeout(10800) -class TestKimiK3(LlmapiAccuracyTestHarness): - MODEL_NAME = "moonshotai/Kimi-K3" - MODEL_PATH = f"{llm_models_root()}/Kimi-K3" - - @skip_pre_blackwell - @pytest.mark.skip_less_mpi_world_size(16) - # The 16-GPU K3 recipes are qualified on GB300 (one NVL72 domain) only: - # on 2-node 180-190 GiB parts (B200/GB200, InfiniBand between nodes) the - # EP16 MoE-comm bring-up hangs and the KV-budget assumptions do not hold, - # so gate on GB300-class device memory. B300 clears this memory gate but - # pairs 8-GPU nodes over InfiniBand (same non-NVL72 topology) — do not - # schedule these tests on B300; that exclusion is enforced by QA's - # platform selection, not by this marker. - @pytest.mark.skip_less_device_memory(200000) - @pytest.mark.parametrize("mode", ["baseline", "reuse", "sa"]) - def test_w4a16_mxfp4(self, mode: str, - monkeypatch: pytest.MonkeyPatch) -> None: - """GSM8K on the bf16 + MXFP4-routed-expert checkpoint (16 GPUs, DEP16). - - No automated L0 stage schedules 16-GPU functional tests; this case is - registered in qa/llm_function_multinode.txt and run by QA's weekly - multinode pipeline (qualified on 4x4 GB300 nodes). Each mode mirrors - the corresponding examples/kimi_k3/eval_extra_llm_options*.yaml config - - keep them in sync when editing either. The `sa` leg additionally - records spec-dec acceptance: AL/AR lines in the eval log (via - TLLM_EVAL_SPEC_STATS) and an iteration-stats AL asserted against - references/acceptance_length.yaml. - """ - kv_cache_kwargs = dict( - enable_block_reuse=False, - free_gpu_memory_fraction=0.25, - # tokens_per_block=64 keeps the MLA generation path on the - # flashinfer trtllm-gen kernel (K3 has 96 query heads). - tokens_per_block=64, - ) - llm_kwargs = dict( - tensor_parallel_size=16, - moe_expert_parallel_size=16, - enable_attention_dp=True, - max_batch_size=32, - max_num_tokens=8192, - max_seq_len=8192, - trust_remote_code=True, - enable_chunked_prefill=True, - cuda_graph_config=CudaGraphConfig(enable_padding=True, - max_batch_size=32), - moe_config=MoeConfig(max_num_tokens=33024, - use_low_precision_moe_combine=True), - ) - if mode == "reuse": - kv_cache_kwargs["enable_block_reuse"] = True - # Hybrid models expose reusable prefixes only at KDA state - # snapshot boundaries; without a snapshot cadence, block reuse - # silently never engages. - kv_cache_kwargs["mamba_state_config"] = MambaStateConfig( - periodic_snapshot_interval=256) - elif mode == "sa": - llm_kwargs.update( - max_batch_size=8, - disable_overlap_scheduler=True, - enable_chunked_prefill=False, - cuda_graph_config=CudaGraphConfig(max_batch_size=8), - speculative_config=SADecodingConfig(max_draft_len=2), - # AL capture needs per-iteration spec-decoding stats. - enable_iter_perf_stats=True, - max_stats_len=-1, - ) - # Log corpus-aggregate AL and AR at eval end ("Spec-dec stats:" - # lines) — QA records acceptance from the test log (TRTLLM-15036). - monkeypatch.setenv("TLLM_EVAL_SPEC_STATS", "1") - - with LLM(self.MODEL_PATH, - kv_cache_config=KvCacheConfig(**kv_cache_kwargs), - **llm_kwargs) as llm: - # Reference-key contract: the K3 checkpoint carries its - # quantization as nested text_config.quantization_config - # (compressed-tensors MXFP4 experts), which the LLM-args layer - # does not surface — unlike modelopt-style hf_quant_config.json - # checkpoints — so the reference matcher sees quant_algo=None - # and the references/gsm8k.yaml entries carry no quant_algo key. - # If this fires, the args-level resolution changed: update the - # yaml keys together with this assert. - assert llm.args.quant_config.quant_algo is None - task = GSM8K(self.MODEL_NAME) - task.evaluate(llm) - if mode == "sa": - acceptance_length = _compute_acceptance_length(llm) - print(f"[AL] TestKimiK3::test_w4a16_mxfp4[sa] " - f"acceptance_length = {acceptance_length:.3f}") - # ref_al/min_al live in references/acceptance_length.yaml. - # The reference was measured on the same workload but - # through the lm-eval-route estimator, whose weighting may - # not match this iteration-stats one exactly — so min_al is - # set as a loose acceptance-collapse tripwire rather than - # the populate-path default. Tighten both once this test's - # own runs establish a baseline. - assert_acceptance_length("TestKimiK3::test_w4a16_mxfp4", - acceptance_length) - - class TestQwen3_4B(LlmapiAccuracyTestHarness): MODEL_NAME = "Qwen3/Qwen3-4B" diff --git a/tests/integration/test_lists/qa/llm_function_multinode.txt b/tests/integration/test_lists/qa/llm_function_multinode.txt index 3823f8a1c843..a20a1007da23 100644 --- a/tests/integration/test_lists/qa/llm_function_multinode.txt +++ b/tests/integration/test_lists/qa/llm_function_multinode.txt @@ -11,7 +11,7 @@ test_e2e.py::test_multi_nodes_eval[MiniMax-M3-tp16-mmlu] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp2pp1-gen_tp2pp1] test_e2e.py::test_openai_disagg_multi_nodes_completion[ctx_tp1pp2-gen_tp1pp2] test_e2e.py::test_openai_disagg_multi_nodes_completion_service_discovery[etcd] -accuracy/test_llm_api_pytorch.py::TestKimiK3::test_w4a16_mxfp4[baseline] -accuracy/test_llm_api_pytorch.py::TestKimiK3::test_w4a16_mxfp4[reuse] -accuracy/test_llm_api_pytorch.py::TestKimiK3::test_w4a16_mxfp4[sa] +accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[baseline] +accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[reuse] +accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[sa] accuracy/test_llm_api_pytorch_multimodal.py::TestKimiK3::test_w4a16_mxfp4 From f629aef8352b675f8c211080dbdd56c598aab26f Mon Sep 17 00:00:00 2001 From: Guiju Zhang Date: Fri, 4 Sep 2026 10:42:51 -0700 Subject: [PATCH 2/4] update Signed-off-by: Guiju Zhang --- tests/integration/defs/.test_durations | 2 + .../defs/accuracy/accuracy_core.py | 43 +++++++++++++++ tests/integration/defs/accuracy/test_glm52.py | 30 +---------- tests/integration/defs/accuracy/test_kimi3.py | 52 ++++--------------- .../defs/accuracy/test_llm_api_pytorch.py | 46 ++++++---------- 5 files changed, 72 insertions(+), 101 deletions(-) diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index ec972db3e704..1f692ebf5bdd 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -56,6 +56,8 @@ "accuracy/test_glm52.py::TestGLM52NVFP4::test_dep[tp_size=8-ep_size=8]": 379.728, "accuracy/test_glm52.py::TestGLM52NVFP4::test_mtp_index_share[tp_size=8-ep_size=8]": 497.186, "accuracy/test_glm52.py::TestGLM52NVFP4::test_runtime": 214.152, + "accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[baseline]": 620.817, + "accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[sa]": 815.941, "accuracy/test_kv_pool_rebalance_accuracy.py::TestKvPoolRebalanceAccuracy::test_rebalance_matches_baseline[no_overlap]": 85.61068378378378, "accuracy/test_kv_pool_rebalance_accuracy.py::TestKvPoolRebalanceAccuracy::test_rebalance_matches_baseline[overlap]": 100.5872467191601, "accuracy/test_llm_api_pytorch.py::TestDeepSeekR1::test_fp8_blockscale[latency]": 1956.791076923077, diff --git a/tests/integration/defs/accuracy/accuracy_core.py b/tests/integration/defs/accuracy/accuracy_core.py index 9a91797ca3e1..cfe8b4d91060 100644 --- a/tests/integration/defs/accuracy/accuracy_core.py +++ b/tests/integration/defs/accuracy/accuracy_core.py @@ -34,6 +34,7 @@ from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization import QuantAlgo +from tensorrt_llm.sampling_params import LogitsProcessor from ..common import venv_check_call, venv_mpi_check_call from ..conftest import llm_models_root @@ -41,6 +42,31 @@ from .video_mme import VideoMME as VideoMMEEvaluator +class ForceTokenLogitsProcessor(LogitsProcessor): + + def __init__(self, forced_token_id: int) -> None: + self._forced_token_id = forced_token_id + + def __call__( + self, + req_id: int, + logits: torch.Tensor, + token_ids: list[list[int]], + stream_ptr: int | None, + client_id: int | None, + ) -> None: + del req_id, token_ids, client_id + if stream_ptr is None: + self._force_token(logits) + return + with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)): + self._force_token(logits) + + def _force_token(self, logits: torch.Tensor) -> None: + logits.fill_(float("-inf")) + logits[..., self._forced_token_id] = 0 + + def compute_theta(num_samples: int, sigma: float, alpha: float = 0.05, @@ -125,6 +151,23 @@ def assert_passing(self, accuracy: float) -> None: assert accuracy <= self.threshold, err_msg +def compute_acceptance_length(llm: PyTorchLLM) -> float: + """Mean acceptance length over speculative iterations. + + Requires enable_iter_perf_stats=True. Used by the AL-regression tests. + """ + stats = llm.get_stats(timeout=2) + spec_iters = [ + stat["specDecodingStats"] for stat in stats + if stat.get("specDecodingStats") + and stat["specDecodingStats"]["numDraftTokens"] > 0 + ] + assert spec_iters, "No iterations with speculative decoding stats" + accepted = sum(stat["numAcceptedTokens"] for stat in spec_iters) + requests = sum(stat["numRequestsWithDraftTokens"] for stat in spec_iters) + return (accepted + requests) / requests + + def assert_acceptance_length(test_key: str, al_value: float) -> None: """Assert acceptance length meets the registered minimum. diff --git a/tests/integration/defs/accuracy/test_glm52.py b/tests/integration/defs/accuracy/test_glm52.py index ac74f8758560..10828b26ed72 100644 --- a/tests/integration/defs/accuracy/test_glm52.py +++ b/tests/integration/defs/accuracy/test_glm52.py @@ -17,7 +17,6 @@ import jsonschema import pytest -import torch from tensorrt_llm import LLM from tensorrt_llm.evaluate import JsonModeEval as JsonModeEvaluator @@ -32,34 +31,9 @@ SchedulingParams, ) from tensorrt_llm.quantization import QuantAlgo -from tensorrt_llm.sampling_params import LogitsProcessor from ..conftest import llm_models_root, parametrize_with_ids, skip_pre_blackwell, skip_ray -from .accuracy_core import GSM8K, JsonModeEval, LlmapiAccuracyTestHarness - - -class _ForceTokenLogitsProcessor(LogitsProcessor): - def __init__(self, forced_token_id: int) -> None: - self._forced_token_id = forced_token_id - - def __call__( - self, - req_id: int, - logits: torch.Tensor, - token_ids: list[list[int]], - stream_ptr: int | None, - client_id: int | None, - ) -> None: - del req_id, token_ids, client_id - if stream_ptr is None: - self._force_token(logits) - return - with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)): - self._force_token(logits) - - def _force_token(self, logits: torch.Tensor) -> None: - logits.fill_(float("-inf")) - logits[..., self._forced_token_id] = 0 +from .accuracy_core import GSM8K, ForceTokenLogitsProcessor, JsonModeEval, LlmapiAccuracyTestHarness class _JsonModeGrammarEval(JsonModeEvaluator): @@ -413,7 +387,7 @@ def _assert_logits_processor(llm: LLM) -> None: max_tokens=output_length, temperature=0, end_id=-1, - logits_processor=_ForceTokenLogitsProcessor(forced_token_id), + logits_processor=ForceTokenLogitsProcessor(forced_token_id), ), use_tqdm=False, ) diff --git a/tests/integration/defs/accuracy/test_kimi3.py b/tests/integration/defs/accuracy/test_kimi3.py index dbc2d01a4b65..d5af4d3b1838 100644 --- a/tests/integration/defs/accuracy/test_kimi3.py +++ b/tests/integration/defs/accuracy/test_kimi3.py @@ -16,7 +16,6 @@ import re import pytest -import torch from tensorrt_llm import LLM from tensorrt_llm._torch.configs import KimiK3Config, KimiLinearConfig @@ -30,47 +29,16 @@ SamplingParams, SchedulingParams, ) -from tensorrt_llm.sampling_params import GuidedDecodingParams, LogitsProcessor +from tensorrt_llm.sampling_params import GuidedDecodingParams from ..conftest import llm_models_root, skip_pre_blackwell -from .accuracy_core import GSM8K, LlmapiAccuracyTestHarness, assert_acceptance_length - - -class _ForceTokenLogitsProcessor(LogitsProcessor): - def __init__(self, forced_token_id: int) -> None: - self._forced_token_id = forced_token_id - - def __call__( - self, - req_id: int, - logits: torch.Tensor, - token_ids: list[list[int]], - stream_ptr: int | None, - client_id: int | None, - ) -> None: - del req_id, token_ids, client_id - if stream_ptr is None: - self._force_token(logits) - return - with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)): - self._force_token(logits) - - def _force_token(self, logits: torch.Tensor) -> None: - logits.fill_(float("-inf")) - logits[..., self._forced_token_id] = 0 - - -def _compute_acceptance_length(llm: LLM) -> float: - stats = llm.get_stats(timeout=2) - spec_iterations = [ - stat["specDecodingStats"] - for stat in stats - if stat.get("specDecodingStats") and stat["specDecodingStats"]["numDraftTokens"] > 0 - ] - assert spec_iterations, "No iterations with speculative decoding stats" - accepted = sum(stat["numAcceptedTokens"] for stat in spec_iterations) - requests = sum(stat["numRequestsWithDraftTokens"] for stat in spec_iterations) - return (accepted + requests) / requests +from .accuracy_core import ( + GSM8K, + ForceTokenLogitsProcessor, + LlmapiAccuracyTestHarness, + assert_acceptance_length, + compute_acceptance_length, +) @pytest.mark.timeout(10800) @@ -166,7 +134,7 @@ def test_w4a16_mxfp4(self, mode: str, monkeypatch: pytest.MonkeyPatch) -> None: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) if mode == "sa": - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print( "[AL] TestKimiK3::test_w4a16_mxfp4[sa] " f"acceptance_length = {acceptance_length:.3f}" @@ -296,7 +264,7 @@ def _assert_logits_processor(llm: LLM) -> None: max_tokens=output_length, temperature=0, end_id=-1, - logits_processor=_ForceTokenLogitsProcessor(forced_token_id), + logits_processor=ForceTokenLogitsProcessor(forced_token_id), ), use_tqdm=False, ) diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index a9271c7738a1..81bb2ed5f610 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -50,7 +50,8 @@ skip_pre_blackwell, skip_pre_hopper, skip_ray, skip_x86) from .accuracy_core import (GSM8K, MMLU, CnnDailymail, GPQADiamond, JsonModeEval, LlmapiAccuracyTestHarness, - LongBenchV1, LongBenchV2, assert_acceptance_length) + LongBenchV1, LongBenchV2, assert_acceptance_length, + compute_acceptance_length) # Keep helper definitions below imports so new imports do not need E402 @@ -104,23 +105,6 @@ def _latest_kv_cache_stats(llm): return entries[-1] -def _compute_acceptance_length(llm) -> float: - """Mean acceptance length over speculative iterations. - - Requires enable_iter_perf_stats=True. Used by the AL-regression tests. - """ - stats = llm.get_stats(timeout=2) - spec_iters = [ - stat['specDecodingStats'] for stat in stats - if stat.get('specDecodingStats') - and stat['specDecodingStats']['numDraftTokens'] > 0 - ] - assert spec_iters, "No iterations with speculative decoding stats" - accepted = sum(s['numAcceptedTokens'] for s in spec_iters) - reqs = sum(s['numRequestsWithDraftTokens'] for s in spec_iters) - return (accepted + reqs) / reqs - - def _assert_non_greedy_cuda_graph_matches_eager(build_llm_kwargs, graph_cuda_graph_config=None): """TRTLLM-14874 regression: non-greedy output must match with and without CUDA graphs. @@ -499,7 +483,7 @@ def test_eagle3(self, overlap_scheduler, eagle3_one_model, enable_iter_perf_stats=True) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print( f"[AL] test_eagle acceptance_length = {acceptance_length:.3f}") assert_acceptance_length( @@ -623,7 +607,7 @@ def test_pard(self, overlap_scheduler): enable_iter_perf_stats=True) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print(f"[AL] test_pard acceptance_length = {acceptance_length:.3f}") assert_acceptance_length("TestLlama3_1_8BInstruct::test_pard", acceptance_length) @@ -757,7 +741,7 @@ def test_dflash(self): enable_iter_perf_stats=True) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print( f"[AL] test_dflash acceptance_length = {acceptance_length:.3f}") assert_acceptance_length("TestLlama3_1_8BInstruct::test_dflash", @@ -824,7 +808,7 @@ def test_ngram(self): enable_iter_perf_stats=True) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print( f"[AL] test_ngram acceptance_length = {acceptance_length:.3f}") assert_acceptance_length("TestLlama3_1_8BInstruct::test_ngram", @@ -859,7 +843,7 @@ def test_suffix_automaton(self, enable_global_pool): enable_iter_perf_stats=True) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print( f"[AL] test_suffix_automaton enable_global_pool={enable_global_pool} " f"acceptance_length = {acceptance_length:.3f}") @@ -965,7 +949,7 @@ def test_draft_target_dynamic_draft_len(self): enable_iter_perf_stats=True) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print(f"[AL] test_draft_target_dynamic_draft_len " f"acceptance_length = {acceptance_length:.3f}") assert_acceptance_length( @@ -1612,7 +1596,7 @@ def test_bfloat16(self, mtp_nextn, attention_dp, cuda_graph, task = GSM8K(self.MODEL_NAME) task.evaluate(llm) if mtp_nextn > 0: - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print(f"[AL] test_bfloat16 mtp_nextn={mtp_nextn} " f"acceptance_length = {acceptance_length:.3f}") assert_acceptance_length( @@ -4253,7 +4237,7 @@ def test_gsm8k_dep8_megamoe_deepgemm(self): assert score >= acc_params.ref_accuracy, ( f"GSM8K accuracy {score:.3f} is below recorded reference " f"{acc_params.ref_accuracy:.3f}") - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print(f"[AL] test_gsm8k_dep8_megamoe_deepgemm " f"acceptance_length = {acceptance_length:.3f}") assert_acceptance_length( @@ -4669,7 +4653,7 @@ def test_eagle3(self): enable_iter_perf_stats=True) as llm: task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print( f"[AL] test_eagle3 acceptance_length = {acceptance_length:.3f}") assert_acceptance_length("TestQwen3_4B::test_eagle3", @@ -4904,7 +4888,7 @@ def test_dspark(self, attention_backend): apply_chat_template=True, chat_template_kwargs={"enable_thinking": False}, )) - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print(f"[AL] test_dspark[{attention_backend}] acceptance_length " f"= {acceptance_length:.3f}") assert_acceptance_length("TestQwen3_8B::test_dspark", @@ -5897,7 +5881,7 @@ def test_dflash(self, mocker): task = GSM8K(model_name) task.evaluate(llm, extra_evaluator_kwargs=self.extra_evaluator_kwargs) - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print( f"[AL] test_dflash acceptance_length = {acceptance_length:.3f}") assert_acceptance_length("TestGPTOSS::test_dflash", @@ -6873,7 +6857,7 @@ def test_dflash(self): task.evaluate(llm, extra_acc_spec=extra_acc_spec, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print( f"[AL] test_dflash acceptance_length = {acceptance_length:.3f}") assert_acceptance_length("TestQwen3_5_4B::test_dflash", @@ -8097,7 +8081,7 @@ def test_nvfp4_4gpus_block_reuse(self, tp_size, ep_size, task.evaluate(llm, extra_evaluator_kwargs=self.EXTRA_EVALUATOR_KWARGS) if use_mtp: - acceptance_length = _compute_acceptance_length(llm) + acceptance_length = compute_acceptance_length(llm) print("[AL] test_nvfp4_4gpus_block_reuse " f"acceptance_length = {acceptance_length:.3f}") assert_acceptance_length( From 8c6b7e5baa4402c9274814a9f8170a8f098fddf3 Mon Sep 17 00:00:00 2001 From: Guiju Zhang Date: Fri, 4 Sep 2026 10:58:17 -0700 Subject: [PATCH 3/4] update doc Signed-off-by: Guiju Zhang --- docs/source/models/supported-models.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index be2b5339b35f..4d1f74baf3e3 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -104,8 +104,8 @@ Note: Support for other models may vary. Features marked "N/A" are not applicabl [^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. -[^16]: Kimi K3 supports suffix-automaton and DSpark speculative decoding. It has no MTP or EAGLE-3 head, and its DSpark checkpoints are not compatible with the plain `DFlash` mode represented by this matrix. -[^17]: Guided decoding for `Glm4MoeForCausalLM` is currently supported only with the `xgrammar` backend; `llguidance` is not working. +[^16]: Guided decoding for `Glm4MoeForCausalLM` is currently supported only with the `xgrammar` backend; `llguidance` is not working. +[^17]: Kimi K3 supports suffix-automaton and DSpark speculative decoding. It has no MTP or EAGLE-3 head, and its DSpark checkpoints are not compatible with the plain `DFlash` mode represented by this matrix. # Encoder-Decoder Feature Support Matrix (PyTorch Backend) From 184fbae507e2b8050b1cd9e8264a02b8b5a04b29 Mon Sep 17 00:00:00 2001 From: Guiju Zhang Date: Fri, 4 Sep 2026 17:26:22 -0700 Subject: [PATCH 4/4] update Signed-off-by: Guiju Zhang --- tests/integration/defs/.test_durations | 1 + tests/integration/defs/accuracy/test_kimi3.py | 5 ++--- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index 1f692ebf5bdd..607a0757d144 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -57,6 +57,7 @@ "accuracy/test_glm52.py::TestGLM52NVFP4::test_mtp_index_share[tp_size=8-ep_size=8]": 497.186, "accuracy/test_glm52.py::TestGLM52NVFP4::test_runtime": 214.152, "accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[baseline]": 620.817, + "accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[reuse]": 684.150, "accuracy/test_kimi3.py::TestKimiK3::test_w4a16_mxfp4[sa]": 815.941, "accuracy/test_kv_pool_rebalance_accuracy.py::TestKvPoolRebalanceAccuracy::test_rebalance_matches_baseline[no_overlap]": 85.61068378378378, "accuracy/test_kv_pool_rebalance_accuracy.py::TestKvPoolRebalanceAccuracy::test_rebalance_matches_baseline[overlap]": 100.5872467191601, diff --git a/tests/integration/defs/accuracy/test_kimi3.py b/tests/integration/defs/accuracy/test_kimi3.py index d5af4d3b1838..01b46064dd44 100644 --- a/tests/integration/defs/accuracy/test_kimi3.py +++ b/tests/integration/defs/accuracy/test_kimi3.py @@ -230,18 +230,17 @@ def _assert_kv_cache_reuse(llm: LLM) -> None: end_id=-1, return_perf_metrics=True, ) - scheduling_params = [SchedulingParams(attention_dp_rank=0, attention_dp_relax=False)] + # Sequential requests on an otherwise idle default router are assigned + # to the same ADP rank, so explicit rank pinning is unnecessary here. cold_output = llm.generate( [prompt_token_ids], sampling_params=sampling_params, - scheduling_params=scheduling_params, use_tqdm=False, )[0].outputs[0] warm_output = llm.generate( [prompt_token_ids], sampling_params=sampling_params, - scheduling_params=scheduling_params, use_tqdm=False, )[0].outputs[0]