From 22f96c5fd503e1c9c4000c816100d41ede5b8ca4 Mon Sep 17 00:00:00 2001 From: Zhenhuan Chen Date: Sun, 16 Aug 2026 21:34:23 -0700 Subject: [PATCH] [None][feat] Add Kimi K3 to layer-wise benchmarks Hybrid KDA+MLA KV cache, K3 layer-call convention with the attn-residual snapshot stack, MoE discovery, and NVTX ranges. Also skips weight-stripped layers in the post-load walks, which DUMMY and truncated slices now reach. Signed-off-by: Zhenhuan Chen --- examples/layer_wise_benchmarks/README.md | 44 ++++ examples/layer_wise_benchmarks/run.py | 40 ++++ .../_torch/models/modeling_kimi_linear.py | 38 ++- .../tools/layer_wise_benchmarks/mark_utils.py | 13 ++ .../tools/layer_wise_benchmarks/runner.py | 216 +++++++++++++++++- .../test_lists/test-db/l0_b200.yml | 1 + .../tools/test_layer_wise_benchmarks.py | 48 ++++ 7 files changed, 383 insertions(+), 17 deletions(-) diff --git a/examples/layer_wise_benchmarks/README.md b/examples/layer_wise_benchmarks/README.md index 77ff7d15b6f6..9fecec99a595 100644 --- a/examples/layer_wise_benchmarks/README.md +++ b/examples/layer_wise_benchmarks/README.md @@ -100,6 +100,28 @@ NP=1 ./mpi_launch.sh ./run.sh config_gen.yaml --model nvidia/NVIDIA-Nemotron-3-N NP=2 ./mpi_launch.sh ./run.sh config_ctx.yaml --model Qwen/Qwen3-Next-80B-A3B-Instruct --layer-indices 6,7 --no-enable-attention-dp --mamba-ssm-cache-dtype float16 --batch-size 4 NP=2 ./mpi_launch.sh ./run.sh config_gen.yaml --model Qwen/Qwen3-Next-80B-A3B-Instruct --layer-indices 6,7 --no-enable-attention-dp --mamba-ssm-cache-dtype float16 --batch-size 512 +# Run Kimi K3 (KimiLinear: KDA + MLA hybrid), $K3 being a local checkpoint directory. +# Layer indices are 0-based, so 4,5,6,7 is three KDA layers followed by one full-attention +# (MLA) layer -- one full period of the 3:1 pattern. Layer 0 is the only dense layer. +# --moe-backend must be TRTLLM: the SiTU routed experts support only TRTLLM and +# MEGAMOE_DEEPGEMM, the config files default to CUTLASS, and of those two only TRTLLM is in +# the balance-method whitelist -- MEGAMOE_DEEPGEMM additionally needs +# --balance-method NotModified. GEN also builds a prefill runner, so +# --moe-backend-for-prefill must be set to match. +# The attn-residual snapshot stack grows one entry every attn_res_block_size (12) layers and +# is seeded from layer_indices[0], so a slice at 0,1,2,3 measures a cheaper attn_res than one +# at 48,49,50,51. +NP=4 ./mpi_launch.sh ./run.sh config_ctx.yaml --model $K3 --layer-indices 4,5,6,7 --tokens-per-block 64 --moe-backend TRTLLM +NP=4 ./mpi_launch.sh ./run.sh config_gen.yaml --model $K3 --layer-indices 4,5,6,7 --tokens-per-block 64 --moe-backend TRTLLM --moe-backend-for-prefill TRTLLM --seq-len-q 1 + +# Run Kimi K3 with 3 draft tokens per step (seq_len_q = 1 + 3). --spec-max-draft-len makes +# the cache manager allocate the KDA multi-token verify buffers; without it the run stops +# with an error. No sampler exists here, so no draft is ever accepted. +NP=4 ./mpi_launch.sh ./run.sh config_gen.yaml --model $K3 --layer-indices 4,5,6,7 --tokens-per-block 64 --moe-backend TRTLLM --moe-backend-for-prefill TRTLLM --batch-size 32 --seq-len-q 4 --spec-max-draft-len 3 + +# Run Kimi K3 at the serving parallelism (DEP16, see examples/kimi_k3/eval_extra_llm_options.yaml) +NP=16 ./mpi_launch.sh ./run.sh config_gen.yaml --model $K3 --layer-indices 4,5,6,7 --tokens-per-block 64 --moe-backend TRTLLM --moe-backend-for-prefill TRTLLM --seq-len-q 1 --moe-max-num-tokens 33024 --use-low-precision-moe-combine + # Run with DeepEP A2A NP=4 ./mpi_launch.sh -x TRTLLM_FORCE_COMM_METHOD=DEEPEP ./run.sh config_ctx.yaml --moe-backend CUTEDSL NP=4 ./mpi_launch.sh -x TRTLLM_FORCE_COMM_METHOD=DEEPEP ./run.sh config_gen.yaml --moe-backend CUTEDSL @@ -377,3 +399,25 @@ Limitations: 2. Error `huggingface_hub.errors.HfHubHTTPError: 429 Client Error: Too Many Requests for url: https://huggingface.co/nvidia/DeepSeek-R1-0528-FP4-v2/resolve/main/config.json`. Please use a local model through the `--model` option, or follow Hugging Face's instructions: "We had to rate limit your IP. To continue using our service, create a HF account or login to your existing account, and make sure you pass a HF_TOKEN if you're using the API." + +3. Error `Routing results are not replaced` with a MoE backend that routes inside the kernel. + + `--balance-method` overrides routing in Python, so it needs the top-k to be computed outside the MoE kernel. Under attention DP that is always the case; with `--no-enable-attention-dp` the top-k can be fused into the kernel instead. Either add `-x TLLM_TRTLLMGEN_FORCE_SEPARATED_ROUTING=1` or set `--balance-method NotModified`. + +4. Error `Kimi K3 needs --spec-max-draft-len N for seq_len_q N+1`. + + Kimi K3's multi-token verify path reads speculative KDA state buffers that the cache manager only allocates when it is built with a speculative config. Add `--spec-max-draft-len` equal to `seq_len_q - 1`. Note that the harness has no sampler, so no draft is ever accepted and the result is the zero-acceptance corner of the verify kernel, not a steady-state number. + +5. Error `AttributeError: 'KimiLinearConfig' object has no attribute 'n_routed_experts'` with `--scaled-from`. + + Weak scaling rewrites `n_routed_experts`, which Kimi K3's config does not have (it uses `num_experts`). `--scaled-from` is not supported for this model. + +6. Numbers look plausible but a factor of ~10 too slow on a hybrid model. + + Several model fast paths fall back silently. Check these lines in the rank-0 log: + + - `KDA kernel dispatch: prefill=optimized decode=optimized verify=optimized` + - `fused prefill/decode/verify projections on KDA layers` — `N` must equal the number of KDA layers in `--layer-indices`; `0` means every KDA layer is running the reference path (~70 us/layer of glue around a ~5 us kernel) + - `Mamba Cache (kda-replay) is allocated` — only with `--spec-max-draft-len`; without it the legacy per-step verify buffers are used instead + + For reference, Kimi K3 on 4 GPUs (DEP4, one GB200 node) with `--layer-indices 4,5,6,7` and `--load-format DUMMY` runs ~60 ms per CTX iteration at `--batch-size 1 --seq-len-q 8193`, and ~4.2 ms per GEN iteration at `--batch-size 32 --seq-len-q 4 --seq-len-kv-cache 8193 --spec-max-draft-len 3`. diff --git a/examples/layer_wise_benchmarks/run.py b/examples/layer_wise_benchmarks/run.py index 66f1ceffaa6c..6a0004dc599d 100644 --- a/examples/layer_wise_benchmarks/run.py +++ b/examples/layer_wise_benchmarks/run.py @@ -52,6 +52,13 @@ def comma_separated_floats(s): parser.add_argument("--kv-pool-headroom", type=int, default=1) parser.add_argument("--enable-swa-scratch-reuse", action="store_true") # Model init args +parser.add_argument( + "--vision-config", + type=str, + choices=["none", "checkpoint"], + help="Which vision tower to build from a composite multimodal checkpoint. Default" + ' "none": drop it and profile the text decoder layers, which is all this tool runs.', +) parser.add_argument("--load-format", type=str, choices=["AUTO", "DUMMY"]) parser.add_argument("--max-num-tokens", type=int) parser.add_argument("--moe-backend", type=str) @@ -59,6 +66,13 @@ def comma_separated_floats(s): # whose Literal is the authoritative list. A second list here can only drift. parser.add_argument("--moe-backend-for-prefill", type=str) parser.add_argument("--moe-max-num-tokens", type=int) +parser.add_argument( + "--spec-max-draft-len", + type=int, + help="Draft tokens per generation step. Required by models whose multi-token verify" + " path needs cache-manager-allocated speculative buffers (Kimi K3 KDA). Must equal" + " seq_len_q - 1.", +) group = parser.add_mutually_exclusive_group() group.add_argument( "--use-low-precision-moe-combine", action="store_true", dest="use_low_precision_moe_combine" @@ -137,6 +151,8 @@ def comma_separated_floats(s): args.kv_cache_dtype = "auto" if args.mamba_ssm_cache_dtype is None: args.mamba_ssm_cache_dtype = "auto" +if args.vision_config is None: + args.vision_config = "none" if args.load_format is None: args.load_format = "DUMMY" if args.max_num_tokens is None: @@ -153,8 +169,27 @@ def comma_separated_floats(s): parser.error("Both --replay-start-iter and --replay-stop-iter must be provided or none") if args.replay_verify_metadata is None: args.replay_verify_metadata = True +if args.spec_max_draft_len is not None: + if args.run_type != "GEN": + parser.error("--spec-max-draft-len only applies to --run-type GEN") + if args.spec_max_draft_len < 1: + parser.error("--spec-max-draft-len must be >= 1") + if args.seq_len_q_list != [args.spec_max_draft_len + 1]: + parser.error( + f"--spec-max-draft-len {args.spec_max_draft_len} implies a single seq_len_q" + f" {args.spec_max_draft_len + 1}, got {args.seq_len_q_list}" + ) print(args) +# Speculative decoding shape only: the harness has no sampler, so no drafts are +# ever accepted. This exists so the cache manager allocates the multi-token +# verify buffers (Kimi K3 KDA replay caches) that seq_len_q > 1 needs. +spec_config = None +if args.spec_max_draft_len is not None: + from tensorrt_llm.llmapi.llm_args import SADecodingConfig + + spec_config = SADecodingConfig(max_draft_len=args.spec_max_draft_len) + # MPI args rank = mpi_rank() world_size = mpi_world_size() @@ -175,6 +210,8 @@ def comma_separated_floats(s): layer_indices=args.layer_indices, kv_pool_headroom=args.kv_pool_headroom, enable_swa_scratch_reuse=args.enable_swa_scratch_reuse, + spec_config=spec_config, + vision_config=args.vision_config, ) attn_workspace = torch.empty((0,), device="cuda", dtype=torch.int8) logger.info("Layer-wise benchmarks: Create KV cache manager ... Done") @@ -200,6 +237,8 @@ def comma_separated_floats(s): mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype, use_low_precision_moe_combine=args.use_low_precision_moe_combine, use_cuda_graph=args.use_cuda_graph, + spec_config=spec_config, + vision_config=args.vision_config, ) logger.info("Layer-wise benchmarks: Create runner ... Done") @@ -270,6 +309,7 @@ def comma_separated_floats(s): mamba_ssm_cache_dtype=args.mamba_ssm_cache_dtype, use_low_precision_moe_combine=args.use_low_precision_moe_combine, use_cuda_graph=False, + vision_config=args.vision_config, ) logger.info("Layer-wise benchmarks: Create runner for prefill ... Done") diff --git a/tensorrt_llm/_torch/models/modeling_kimi_linear.py b/tensorrt_llm/_torch/models/modeling_kimi_linear.py index 7f15e37fe0c7..71bd1e90b22e 100644 --- a/tensorrt_llm/_torch/models/modeling_kimi_linear.py +++ b/tensorrt_llm/_torch/models/modeling_kimi_linear.py @@ -549,6 +549,16 @@ def _swap_linear_to_fp8_weight_read( return 1 +def _has_weights(module: nn.Module) -> bool: + """False once ``modeling_utils.remove_weights()`` has stripped a module. + + Post-load finalization walks every decoder layer, so it must skip layers + whose parameters were dropped — the layer-wise benchmarks keep only the + profiled slice resident. + """ + return not getattr(module, "_weights_removed", False) + + def _convert_moe_mlps_to_fp8_weight_read( model: nn.Module, include_fused_gate_up: bool = True ) -> int: @@ -563,6 +573,8 @@ def _convert_moe_mlps_to_fp8_weight_read( count = 0 for layer in model.layers: + if not _has_weights(layer): + continue moe = getattr(layer, "block_sparse_moe", None) if moe is None: continue @@ -623,7 +635,7 @@ def _convert_kda_projections_to_fp8_weight_read(model: nn.Module) -> int: count = 0 for layer in model.layers: - if not getattr(layer, "is_kda", False): + if not getattr(layer, "is_kda", False) or not _has_weights(layer): continue mixer = getattr(getattr(layer, "self_attn", None), "mixer", None) if mixer is None: @@ -693,7 +705,7 @@ def _convert_mla_projections_to_fp8_weight_read(model: nn.Module) -> int: for layer in model.layers: # MLA layers are the non-KDA layers (each layer is exactly one of the # two); their projections live on the KimiK3MLAAttention mixer. - if getattr(layer, "is_kda", False): + if getattr(layer, "is_kda", False) or not _has_weights(layer): continue mixer = getattr(getattr(layer, "self_attn", None), "mixer", None) if mixer is None: @@ -2235,6 +2247,20 @@ def forward( prefix_sum = prefix_sum + hidden_states return prefix_sum, num_snapshots + def skip_forward( + self, + hidden_states: torch.Tensor, + block_residual: torch.Tensor, + attn_metadata: AttentionMetadata, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """No-op stand-in for ``forward``, matching ``DecoderLayer.skip_forward``. + + ``modeling_utils.skip_forward()`` only drops a module's weights when it + finds this attribute, so without it the layer-wise benchmarks would + allocate all 93 layers instead of the profiled slice. + """ + return hidden_states, block_residual + # --------------------------------------------------------------------------- # Model. @@ -2472,7 +2498,7 @@ def checkpoint_name_plan( # when the routed MoE is TP-sharded (moe_ep=1 -> ids 0..num_experts)). expert_jobs = [] for layer_idx, layer in enumerate(self.model.layers): - if not getattr(layer, "is_moe", False): + if not getattr(layer, "is_moe", False) or not _has_weights(layer): continue moe = layer.block_sparse_moe base = f"{prefix}model.layers.{layer_idx}.block_sparse_moe.experts" @@ -2541,7 +2567,7 @@ def _load_trunk_params( mla_mixers = [ layer.self_attn.mixer for layer in self.model.layers - if not getattr(layer, "is_kda", True) + if not getattr(layer, "is_kda", True) and _has_weights(layer) ] mla_kv_b_mixers = {id(mixer.kv_b_proj.weight): mixer for mixer in mla_mixers} mla_head_shard_linears = {} @@ -2864,7 +2890,7 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: # This must run after every KDA parameter is loaded and sharded. num_kda_fused = 0 for layer in self.model.layers: - if getattr(layer, "is_kda", False): + if getattr(layer, "is_kda", False) and _has_weights(layer): if not kda_fp8: layer.self_attn.finalize_decode_weights() num_kda_fused += int( @@ -2919,7 +2945,7 @@ def _finalize_weight_load(self, num_params: int, num_moe_layers: int) -> None: # in BF16. n_glue = 0 for layer in self.model.layers: - if getattr(layer, "is_kda", False): + if getattr(layer, "is_kda", False) and _has_weights(layer): layer.self_attn.finalize_decode_weights_fp8() n_glue += int(layer.self_attn._bfa_proj_weight is not None) logger.info( diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py b/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py index 97fcf7f7dabf..53237af8992a 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/mark_utils.py @@ -2,6 +2,12 @@ from tensorrt_llm._torch.models.modeling_deepseekv3 import DeepseekV3Gate, Deepseekv3MoE from tensorrt_llm._torch.models.modeling_deepseekv4 import DeepseekV4Gate, DeepseekV4MoE +from tensorrt_llm._torch.models.modeling_kimi_linear import ( + KimiK3MoEGate, + KimiK3MoERuntime, + KimiKDARuntime, + KimiMLARuntime, +) from tensorrt_llm._torch.models.modeling_nemotron_h import MLPLayer, NemotronHMOE from tensorrt_llm._torch.models.modeling_qwen3_next import ( Qwen3NextGatedDeltaNet, @@ -26,6 +32,13 @@ def mark_ranges(): Qwen3NextSparseMoeBlock.forward = nvtx.annotate("Qwen3NextSparseMoeBlock")( Qwen3NextSparseMoeBlock.forward ) + # Kimi K3. `KimiK3MLAAttention` overrides `MLA.forward`, so the range below + # is on its `KimiMLARuntime` wrapper. The gate is entered through + # `compute_logits`, not `forward`. Its MLPs are the shared `GatedMLP`. + KimiKDARuntime.forward = nvtx.annotate("KimiKDARuntime")(KimiKDARuntime.forward) + KimiMLARuntime.forward = nvtx.annotate("KimiMLARuntime")(KimiMLARuntime.forward) + KimiK3MoERuntime.forward = nvtx.annotate("KimiK3MoERuntime")(KimiK3MoERuntime.forward) + KimiK3MoEGate.compute_logits = nvtx.annotate("KimiK3MoEGate")(KimiK3MoEGate.compute_logits) MLA.forward = nvtx.annotate("MLA")(MLA.forward) Attention.forward = nvtx.annotate("Attention")(Attention.forward) MoE.forward = nvtx.annotate("MoE")(MoE.forward) diff --git a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py index adb1d0c17338..e2fa43776a11 100644 --- a/tensorrt_llm/tools/layer_wise_benchmarks/runner.py +++ b/tensorrt_llm/tools/layer_wise_benchmarks/runner.py @@ -11,22 +11,29 @@ import torch import tensorrt_llm._torch.model_config +import tensorrt_llm._torch.pyexecutor.config_utils import tensorrt_llm.bindings from tensorrt_llm._torch.attention_backend.utils import get_attention_backend from tensorrt_llm._torch.custom_ops.cute_dsl_custom_ops import GroupedGemmInputsHelper from tensorrt_llm._torch.metadata import KVCacheParams from tensorrt_llm._torch.model_config import ModelConfig -from tensorrt_llm._torch.models.modeling_utils import PostInitCaller, skip_forward +from tensorrt_llm._torch.models.modeling_utils import PostInitCaller, remove_weights, skip_forward from tensorrt_llm._torch.modules.fused_moe.fused_moe_trtllm_gen import TRTLLMGenFusedMoE from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata -from tensorrt_llm._torch.pyexecutor._util import get_kv_cache_manager_cls +from tensorrt_llm._torch.pyexecutor._util import _mamba_conv_layout_kwargs, get_kv_cache_manager_cls from tensorrt_llm._torch.pyexecutor.config_utils import ( + extract_mamba_kv_cache_params, + get_kimi_linear_layer_masks, get_qwen3_hybrid_layer_masks, + is_hybrid_linear, + is_kimi_linear, is_mla, is_nemotron_hybrid, is_qwen3_hybrid, load_pretrained_config, + unwrap_kimi_text_config, ) +from tensorrt_llm._torch.pyexecutor.mamba_cache_manager import MixedMambaHybridCacheManager from tensorrt_llm._torch.pyexecutor.model_loader import ( ModelLoader, _construct_checkpoint_loader, @@ -36,7 +43,7 @@ from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._torch.utils import get_model_extra_attrs, model_extra_attrs from tensorrt_llm._utils import local_mpi_size, mpi_rank, mpi_world_size, torch_dtype_to_binding -from tensorrt_llm.llmapi.llm_args import KvCacheConfig, MoeConfig, TorchLlmArgs +from tensorrt_llm.llmapi.llm_args import DecodingBaseConfig, KvCacheConfig, MoeConfig, TorchLlmArgs from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping @@ -398,7 +405,9 @@ def __init__( mamba_ssm_cache_dtype: str, use_low_precision_moe_combine: bool, use_cuda_graph: bool, - ): + spec_config: Optional[DecodingBaseConfig] = None, + vision_config: Optional[str] = None, + ) -> None: super().__init__() checkpoint_loader = _construct_checkpoint_loader("pytorch", None, "HF") @@ -406,6 +415,9 @@ def __init__( llm_args = TorchLlmArgs( model=pretrained_model_name_or_path, load_format=load_format, + # `ModelLoader(spec_config=...)` below is what reaches + # `model_config.spec_config`; this keeps `llm_args` consistent with it. + **({"speculative_config": spec_config} if spec_config is not None else {}), **{} if use_cuda_graph else {"cuda_graph_config": None}, moe_config=MoeConfig( backend=moe_backend, @@ -421,17 +433,30 @@ def __init__( model_loader = ModelLoader( llm_args=llm_args, mapping=mapping, - spec_config=None, + spec_config=spec_config, sparse_attention_config=None, max_num_tokens=max_num_tokens, max_seq_len=max_seq_len, ) - with self.scaled_from_ctx(scaled_from, mapping), self.skip_unused_layers_ctx(layer_indices): + with ( + self.scaled_from_ctx(scaled_from, mapping), + self.vision_config_ctx(vision_config), + self.skip_unused_layers_ctx(layer_indices), + ): model, _ = model_loader.load( checkpoint_dir=pretrained_model_name_or_path, checkpoint_loader=checkpoint_loader ) + finalize_weight_load = getattr(model, "_finalize_weight_load", None) + if load_format == "DUMMY" and finalize_weight_load is not None: + # Models that build decode fast-path constants at the end of + # `load_weights` never get them under DUMMY, and then silently run a + # reference path instead (Kimi K3's KDA decode: ~70 us/layer of glue + # around a ~5 us kernel). Run the hook so DUMMY measures the same + # kernels as a real load. The arguments only feed a log line. + finalize_weight_load(0, 0) + def forward(position_ids, hidden_states, attn_metadata, residual, **kwargs): # TODO: to be more general, we should call DecoderModel.forward for layer_idx in layer_indices: @@ -445,12 +470,71 @@ def forward(position_ids, hidden_states, attn_metadata, residual, **kwargs): hidden_states = layer(position_ids, hidden_states, attn_metadata, **kwargs) return hidden_states, residual - model.forward = forward + def forward_block_residual( + position_ids: torch.Tensor, + hidden_states: torch.Tensor, + attn_metadata, + residual: torch.Tensor, + num_snapshots: int = 0, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + # Kimi K3 layers take a preallocated snapshot bank plus the count of + # valid rows, and no `position_ids` -- MLA derives RoPE positions from + # `attn_metadata`. `create_run_pack` passes the bank through the + # `residual` slot and seeds the count for a mid-model slice. + block_residual = residual + for layer_idx in layer_indices: + layer = model.model.layers[layer_idx] + hidden_states, num_snapshots = layer( + hidden_states, block_residual, num_snapshots, attn_metadata + ) + return hidden_states, block_residual + + # Layers carrying a snapshot stack take `block_residual` where the generic + # ones take `residual`, and derive positions from `attn_metadata`. + first_layer = model.model.layers[layer_indices[0]] + if "block_residual" in inspect.signature(first_layer.forward).parameters: + model.forward = forward_block_residual + else: + model.forward = forward self.model_config = model.model_config self.model = model self.layer_indices = layer_indices + @staticmethod + @contextlib.contextmanager + def vision_config_ctx(vision_config: Optional[str]): + """Pick which config a composite multimodal checkpoint resolves to. + + A checkpoint shipping `vision_config` beside `text_config` is kept composite + by config loading and resolved to a `*ForConditionalGeneration` wrapper so + the vision tower stays available; the text-only path is the fallback for + checkpoints without one. This harness profiles text decoder layers, so the + default drops the tower and uses the inner text config, which already names + its own architecture. Keyed on `vision_config` rather than a model check + because the same condition gates every such route. + """ + if vision_config != "none": + yield + return + + model_config_module = tensorrt_llm._torch.model_config + load_pretrained_config_orig = model_config_module.load_pretrained_config + + def load_pretrained_config_text_only(*args, **kwargs): + config = load_pretrained_config_orig(*args, **kwargs) + text_config = getattr(config, "text_config", None) + if getattr(config, "vision_config", None) is not None and text_config is not None: + return text_config + return config + + model_config_module.load_pretrained_config = load_pretrained_config_text_only + try: + yield + finally: + model_config_module.load_pretrained_config = load_pretrained_config_orig + @staticmethod @contextlib.contextmanager def scaled_from_ctx(scaled_from, mapping): @@ -503,7 +587,14 @@ def call_new(cls, *args, **kwargs): skip_forward(module) num_hidden_layers = model.model_config.pretrained_config.num_hidden_layers if hasattr(model.model, "embed_tokens"): - skip_forward(model.model.embed_tokens) + embed_tokens = model.model.embed_tokens + if hasattr(embed_tokens, "skip_forward"): + skip_forward(embed_tokens) + else: + # Plain `nn.Embedding` (Kimi K3): no `skip_forward` to swap in, + # but `model.forward` never reaches it, so dropping the + # weights is enough and saves vocab_size * hidden_size bytes. + remove_weights(embed_tokens) for layer_idx in range(num_hidden_layers): layer = model.model.layers[layer_idx] if layer_idx not in layer_indices: @@ -619,6 +710,40 @@ def create_run_pack( ) kwargs = {} + # Fail here rather than deep inside the model's verify path, which reads + # buffers the cache manager only allocates for a speculative config. + if ( + run_type == "GEN" + and seq_len_q > 1 + and getattr(kv_cache_manager, "is_speculative", None) is not None + and not kv_cache_manager.is_speculative() + ): + raise ValueError( + f"seq_len_q {seq_len_q} needs --spec-max-draft-len {seq_len_q - 1}:" + " multi-token verify reads speculative recurrent-state buffers" + ) + + # An attn-residual model (Kimi K3) carries a + # [num_snapshots, num_tokens, hidden_size] stack instead of a residual + # tensor, pushing one snapshot every `attn_res_block_size` layers. A slice + # starting at `layer_indices[0]` inherits `ceil(that / block_size)` of them; + # since the mixing cost scales with the depth, a slice started mid-model + # must not begin from an empty stack. + attn_res_block_size = getattr( + unwrap_kimi_text_config(pretrained_config), "attn_res_block_size", None + ) + if attn_res_block_size is not None: + # The bank is preallocated at full-model capacity and `num_snapshots` + # counts the valid rows, so a slice starting at `layer_indices[0]` + # inherits `ceil(that / block_size)` of them. The mixing cost scales + # with the count, so a mid-model slice must not start from zero. + kwargs["num_snapshots"] = ceil_div(self.layer_indices[0], attn_res_block_size) + residual = torch.rand( + (self.model.model.num_attn_res_snapshots, batch_size * seq_len_q, hidden_size), + dtype=torch.bfloat16, + device="cuda", + ) + # DeepSeek-V4 (multi-head hyper-connection) decoder layers take the initial residual # as ``hc_state`` shaped ``[num_tokens, hc_mult, hidden_size]`` (not a 2D hidden-states # tensor), and their MoE routing requires ``input_ids``. Both are absent from the @@ -691,6 +816,10 @@ def replace_routing_method_ctx(self, balance_method: BalanceMethod, balance_rati moe_modules.append(layer.mixer.experts) elif layer.__class__.__name__ in ["GatedMLP"]: pass + elif (block_sparse_moe := getattr(layer, "block_sparse_moe", None)) is not None: + # Latent-MoE layout (Kimi K3): the routed experts sit behind the + # block, and layers below `first_k_dense_replace` are dense. + moe_modules.append(block_sparse_moe.routed_experts) else: moe_modules.append(layer.mlp.experts) @@ -753,9 +882,12 @@ def create_kv_cache_manager( layer_indices, kv_pool_headroom=1, enable_swa_scratch_reuse=False, - ): + spec_config: Optional[DecodingBaseConfig] = None, + vision_config: Optional[str] = None, + ) -> KVCacheManager: # Please refer to `tensorrt_llm/_torch/pyexecutor/py_executor_creator.py` for `tokens_per_block` - model_config = ModelConfig.from_pretrained(pretrained_model_name_or_path) + with Runner.vision_config_ctx(vision_config): + model_config = ModelConfig.from_pretrained(pretrained_model_name_or_path) validate_and_set_kv_cache_quant(model_config, kv_cache_dtype) validate_and_set_mamba_ssm_cache_dtype(model_config, mamba_ssm_cache_dtype) if model_config.enable_flash_mla: @@ -781,7 +913,9 @@ def create_kv_cache_manager( "NVFP4": tensorrt_llm.bindings.DataType.NVFP4, None: torch_dtype_to_binding(config.torch_dtype), }[model_config.quant_config.kv_cache_quant_algo] - if is_mla(config): + # Hybrids below also carry MLA fields, but only some of their layers + # are MLA, so the pure-MLA route must exclude them. + if is_mla(config) and not is_hybrid_linear(config): layer_mask = [i in layer_indices for i in range(config.num_hidden_layers)] num_layers = sum(layer_mask) kv_cache_manager = kv_cache_manager_cls( @@ -802,6 +936,66 @@ def create_kv_cache_manager( sparse_attention_config=model_config.sparse_attention_config, pretrained_config=model_config.pretrained_config, ) + elif is_kimi_linear(config): + # Needs its own branch: neither pure-MLA nor pure-mamba fits. KDA + # recurrent/conv state goes on the mamba side and the MLA latent cache + # on the paged-KV side. Mirrors `_util._create_kv_cache_manager`. + # spec_config=None: it only feeds `num_draft_layers`, and the masks + # below come from `layer_indices` instead. + mamba_params = extract_mamba_kv_cache_params( + config, spec_config=None, quant_config=model_config.quant_config + ) + # Dimensions live on the inner config for a composite checkpoint. + text_config = unwrap_kimi_text_config(config) + full_layer_mask, full_mamba_layer_mask = get_kimi_linear_layer_masks(config) + layer_mask = [ + full_layer_mask[i] and i in layer_indices + for i in range(text_config.num_hidden_layers) + ] + mamba_layer_mask = [ + full_mamba_layer_mask[i] and i in layer_indices + for i in range(text_config.num_hidden_layers) + ] + kimi_extra_kwargs = {} + if spec_config is not None and issubclass( + kv_cache_manager_cls, MixedMambaHybridCacheManager + ): + # Multi-token verify reads per-slot replay caches when the fused + # kernel is available, else the legacy per-step buffers. + from tensorrt_llm._torch.modules.kimi_kda._kda_kernels import ( + is_kda_mtp_verify_available, + ) + + if is_kda_mtp_verify_available(): + kimi_extra_kwargs["kda_replay_num_spec"] = spec_config.tokens_per_gen_step - 1 + kv_cache_manager = kv_cache_manager_cls( + # mamba (KDA) cache parameters + mamba_params.state_size, + mamba_params.conv_kernel, + mamba_params.num_heads, + mamba_params.n_groups, + mamba_params.head_dim, + sum(mamba_layer_mask), + mamba_layer_mask, + mamba_params.dtype, + mamba_params.mamba_ssm_cache_dtype, + # kv cache parameters (MLA latent cache) + kv_cache_config, + tensorrt_llm.bindings.internal.batch_manager.CacheType.SELFKONLY, + num_layers=sum(layer_mask), + layer_mask=layer_mask, + num_kv_heads=1, + head_dim=text_config.kv_lora_rank + text_config.qk_rope_head_dim, + tokens_per_block=tokens_per_block, + max_seq_len=max_seq_len, + max_batch_size=max_batch_size, + mapping=mapping, + dtype=kv_cache_dtype, + spec_config=spec_config, + # KDA conv state is [Q | K | V]: the qwen3_next section layout. + **_mamba_conv_layout_kwargs(kv_cache_manager_cls, "qwen3_next"), + **kimi_extra_kwargs, + ) elif is_nemotron_hybrid(config): mamba_layer_mask = [ i in layer_indices and char == "M" diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index daa214082f0b..e1acb231f60a 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -188,6 +188,7 @@ l0_b200: - unittest/tools/test_layer_wise_benchmarks.py::test_deepseek_r1_ctx_dep[1] - unittest/tools/test_layer_wise_benchmarks.py::test_nemotron_gen_dep[1] - unittest/tools/test_layer_wise_benchmarks.py::test_qwen3_next_gen_tep[1] + - unittest/tools/test_layer_wise_benchmarks.py::test_kimi_k3_gen_dep[1] - unittest/tools/test_layer_wise_benchmarks.py::test_performance_alignment[1] - unittest/kv_cache_manager_v2_tests # ------------- KV Cache V2 Scheduler IT --------------- diff --git a/tests/unittest/tools/test_layer_wise_benchmarks.py b/tests/unittest/tools/test_layer_wise_benchmarks.py index edb68422956c..cec58396bb2b 100644 --- a/tests/unittest/tools/test_layer_wise_benchmarks.py +++ b/tests/unittest/tools/test_layer_wise_benchmarks.py @@ -340,6 +340,54 @@ def test_qwen3_next_gen_tep(llm_root, world_size): ) +# Kimi K3's MXFP4 routed experts and KDA kernels require SM100+. +@skip_pre_blackwell +@pytest.mark.parametrize("world_size", [1, 4]) +def test_kimi_k3_gen_dep(llm_root, world_size): + if torch.cuda.device_count() < world_size: + pytest.skip(f"needs {world_size:d} GPUs to run this test") + model_root = llm_models_root(check=True) + profile_dir = f"profiles/test_kimi_k3_gen_dep_{world_size}" + if world_size == 1: + # EP1 puts all 896 experts on one GPU, and GEN builds a second (prefill) + # model, so halve the slice: layer 6 is KDA and 7 is MLA, still covering + # both attention paths. Balanced routing needs the top-k computed outside + # the MoE kernel, which only happens once there is expert parallelism. + layer_args = ["--layer-indices=6,7", "--balance-method=NotModified"] + else: + # 0-based: three KDA layers then one full-attention (MLA) layer. + layer_args = ["--layer-indices=4,5,6,7"] + check_call( + [ + "./mpi_launch.sh", + "./run.sh", + "config_gen.yaml", + "--model", + model_root / "Kimi-K3", + *layer_args, + "--tokens-per-block=64", + # SiTU routed experts support no other backend, and GEN also builds + # a prefill runner. + "--moe-backend=TRTLLM", + "--moe-backend-for-prefill=TRTLLM", + # 1 golden + 3 draft tokens per generation step. + "--batch-size=32", + "--seq-len-q=4", + "--spec-max-draft-len=3", + ], + cwd=llm_root / "examples" / "layer_wise_benchmarks", + env={ + **os.environ, + "NP": f"{world_size:d}", + "PROFILE_DIR": profile_dir, + }, + ) + check_call( + ["python3", "parse.py", "--profile-dir", profile_dir, f"--world-size={world_size}"], + cwd=llm_root / "examples" / "layer_wise_benchmarks", + ) + + # The pinned DeepSeek-V3-Lite NVFP4 checkpoint requires SM100+; on older # architectures the benchmark crashes the test process (seen on A10, where # this module runs as part of the unittest/tools directory).