[TRTLLM-14813][test] Port Kimi K3 unit tests and wire GB300 L0 stages - #17332
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughKimi K3 test coverage now includes cache-manager behavior, optimized KDA decode, fused activation and MLP parity, MoE routing, CUDA Graph replay, and multi-GPU integration entries. ChangesKimi K3 validation
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 20
🧹 Nitpick comments (24)
tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py (1)
375-376: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse a dedicated mismatch exception in
check_accuracy.check_accuracycurrently raises built-inException, so neither test can use a narrower type without changing the helper. Raise a dedicated mismatch exception and expect it in both mutation tests.Test coverage summary: Changed tests:
test_fc1_swap_mutation_breaks_accuracyandtest_swiglu_act_mutation_breaks_accuracy. No matching entries exist intests/integration/test_lists/test-db/ortests/integration/test_lists/qa/. Coverage verdict: insufficient.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py` around lines 375 - 376, Define a dedicated mismatch exception for check_accuracy and raise it when the accuracy comparison detects a mismatch, preserving the existing message. Update both test_fc1_swap_mutation_breaks_accuracy (tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py:375-376) and test_swiglu_act_mutation_breaks_accuracy (tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py:406-407) to expect that specific exception instead of built-in Exception.Source: Coding guidelines
tensorrt_llm/_torch/modules/kimi_k3_attn_res/kimi_k3_attn_res.py (2)
227-227: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRemove the unused
_optimized_extensionattribute.No code reads or writes
self._optimized_extensionafter__init__. Delete it, or document why the placeholder exists.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/kimi_k3_attn_res/kimi_k3_attn_res.py` at line 227, Remove the unused self._optimized_extension initialization from the __init__ method of the relevant attention module, since no other code accesses it; do not add a replacement unless the attribute has a documented purpose.
79-83: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueFix the stale reference path in the docstring.
The docstring points to
exisiting_optimization_work/Attention_residual/tests/util/attn_res_ref.py. The directory name contains a typo, and the path is outside this repository. Correct the spelling or drop the reference.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/kimi_k3_attn_res/kimi_k3_attn_res.py` around lines 79 - 83, Update the module docstring for the chunked torch reference to remove the stale external path or replace it with the correctly spelled, repository-valid reference path; retain the description of byte-identical reference outputs.tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py (1)
646-659: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMake the mutation of the HF module explicit in the function name or a flag.
copy_hf_moe_block_weightsoverwriteshf.experts[i].w{1,2,3}.weightwith the canonical MXFP4 round-trip values. The name suggests a one-way copy from HF into K3. A caller that reuses the HF module after this call gets silently changed weights.Add a
canonicalize_source: bool = Trueparameter, or rename the function to state the two-way effect. The docstring already documents step 4d, so this is a naming and API-surface concern only.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py` around lines 646 - 659, Make the two-way mutation explicit in copy_hf_moe_block_weights by adding a canonicalize_source: bool = True parameter and documenting its effect. Guard the expert.w1, expert.w2, and expert.w3 weight overwrites so they occur only when canonicalize_source is enabled, while preserving the existing K3 expert-bank storage behavior.tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py (1)
165-190: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueBuild
DeepSeekV3MoeRoutingMethodonce.The
routing_methodproperty constructs a newDeepSeekV3MoeRoutingMethodon every access. If a caller reads the property per forward, this allocates a router object per step. Cache the instance after the validation checks pass.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py` around lines 165 - 190, Cache the DeepSeekV3MoeRoutingMethod created by the routing_method property after its validation checks succeed, returning the cached instance on subsequent accesses instead of constructing a new router each time. Initialize or store the cache on the owning class and preserve all existing validation behavior.tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py (1)
26-44: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate
_meta_safe_cast_dtypeand drop the local torch alias.The function has no type annotations, and the repository guidelines require an annotation on every function. The module already imports
torchat line 15, soimport torch as _torchis redundant.♻️ Proposed change
-def _meta_safe_cast_dtype(module, dtype): +def _meta_safe_cast_dtype(module: nn.Module, dtype: torch.dtype) -> None: @@ - import torch as _torch - - def _cast(t): + def _cast(t: torch.Tensor) -> torch.Tensor: if not t.is_floating_point(): return t if t.is_meta: - return _torch.empty_like(t, dtype=dtype) + return torch.empty_like(t, dtype=dtype) return t.to(dtype=dtype)Based on coding guidelines: "Annotate every function, use
Nonefor procedures".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py` around lines 26 - 44, Update _meta_safe_cast_dtype with annotations for its module, dtype, and None return value, and annotate the nested _cast helper consistently with the project’s typing conventions. Remove the redundant local “import torch as _torch” and use the existing module-level torch import throughout.Source: Coding guidelines
tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py (1)
56-66: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd a case for a non-contiguous last dimension.
test_situ_and_mul_strided_rowscovers a row stride greater than the row width, with the last dimension still contiguous. The kernel also depends on an element stride of 1 on the last dimension, and no test covers that assumption. Add a case that passes a tensor withstride(-1) != 1and asserts the documented behavior. This pairs with the guard requested intensorrt_llm/_torch/modules/kimi_k3_moe/_mlp.py.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py` around lines 56 - 66, Extend test_situ_and_mul_strided_rows with a non-contiguous-last-dimension input whose stride(-1) is not 1, and assert the documented behavior enforced by the _mlp.py guard. Keep the existing row-strided case intact and verify the new case using the expected validation or failure assertion rather than comparing normal output.tensorrt_llm/_torch/modules/kimi_k3_moe/_mxfp4.py (1)
82-84: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winReduce the transient memory of the nearest-code search.
diffsmaterializes a tensor 16 times larger than the group tensor, in fp32. At real K3 routed-expert dimensions (intermediate=2048,hidden=7168), onestore_expertcall allocates roughly 0.9 GB for a single weight.copy_hf_moe_block_weightscalls this three times per expert, for every expert.Use
torch.bucketizeagainst the sorted magnitude midpoints onscaled.abs(), then restore the sign bit. That removes the 16x expansion and keeps the same nearest-magnitude result.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/kimi_k3_moe/_mxfp4.py` around lines 82 - 84, Replace the `diffs`/`argmin` nearest-code search in the `_FP4_VALUES` encoding flow with `torch.bucketize` on `scaled.abs()` using sorted FP4 magnitude midpoints, then restore the sign bit to produce the uint8 `codes`. Preserve the existing nearest-magnitude mapping while avoiding materialization of the 16x expanded tensor.tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py (1)
409-413: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueReuse the new mirror in
_prepare_replay_work_items.
_prepare_replay_work_itemsat Line 322 still runscache_slot.to(torch.long), which is exactly the int32 to int64 cast the mirror now removes for other consumers. The mirror is refreshed at Line 411, before_prepare_replay_work_itemsis called at Line 415, so the slice is already valid.Note that
_prepare_replay_work_itemsreadsself.state_indices[num_contexts:batch_size], so the equivalent mirror slice isself.state_indices_long[num_contexts:batch_size].♻️ Proposed change (outside the selected range, at Lines 321-322)
cache_slot = self.state_indices[num_contexts:batch_size] cache_slot_idx = self.state_indices_long[num_contexts:batch_size]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py` around lines 409 - 413, Update _prepare_replay_work_items to reuse the refreshed int64 mirror by assigning cache_slot_idx from self.state_indices_long[num_contexts:batch_size] instead of converting the int32 state_indices slice with to(torch.long). Preserve the existing cache_slot slice and downstream behavior.tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py (1)
51-69: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueDrop the redundant local
import torch as _torch.Line 42 already imports
torchat module scope. The local alias at Line 60 adds no protection and hides the dependency.♻️ Proposed cleanup
- import torch as _torch - def _cast(t): if not t.is_floating_point(): return t if t.is_meta: - return _torch.empty_like(t, dtype=dtype) + return torch.empty_like(t, dtype=dtype) return t.to(dtype=dtype)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py` around lines 51 - 69, Remove the redundant local import torch as _torch from _meta_safe_cast_dtype and update its empty_like reference to use the existing module-level torch import, while preserving the current casting behavior.tests/unittest/models/test_quant_config_utils.py (1)
289-313: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueTest coverage summary (QA review), plus one optional gap.
- Changed test functions: added
test_update_quant_config_from_compressed_tensors_mxfp4_with_fp8_kv_cache. No test was modified or removed.- Test-list status: this file is a unit test under
tests/unittest/models/, not an integration test id. No entry undertests/integration/test_lists/test-db/orqa/is required for it.- Verdict: sufficient for the branch under test. The assertions match
update_quant_config_from_compressed_tensors: theformat == "mxfp4-pack-quantized"disjunct selectsW4A16_MXFP4, and thekv_cache_schemeblock runs before the early return, sokv_cache_quant_algobecomesFP8.Optional gap:
_compressed_tensors_configstill injects a non-nullinput_activationsblock, so this case only exercises theformatdisjunct. The second disjunct (num_bits == 4,type == "float",strategy == "group",input_activations is None) stays uncovered. Consider one more case that omitsformatand passesinput_activations=Noneexplicitly.As per path instructions: "Always produce a test coverage summary, even if no issues are found."
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/models/test_quant_config_utils.py` around lines 289 - 313, Extend coverage for update_quant_config_from_compressed_tensors by adding a case that omits the format value and explicitly sets input_activations to None, while retaining the existing 4-bit float group quantization fields. Assert the same MXFP4 quantization behavior and relevant configuration outcomes so the second branch condition is exercised.Source: Path instructions
cpp/tensorrt_llm/thop/kdaDecodeOp.cpp (1)
216-222: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMove the
outputvalidation before the tensor is bound, or validate unconditionally.The code assigns
outfirst and validates afterwards. The result is correct, but the reader must track two branches onoutput.has_value(). A single guarded block reads better and keeps the allocation path free of dead checks.♻️ Proposed restructure
int const B = static_cast<int>(x_q.size(1)); int const HV = static_cast<int>(x_v.size(2)); - auto out = output.has_value() ? *output : at::empty({B, 1, HV, kDimV}, x_q.options()); - if (output.has_value()) - { - TORCH_CHECK(out.is_cuda() && out.scalar_type() == at::kBFloat16, "out must be a CUDA bfloat16 tensor"); - TORCH_CHECK(out.is_contiguous(), "out must be contiguous"); - TORCH_CHECK(out.dim() == 4 && out.size(0) == B && out.size(1) == 1 && out.size(2) == HV && out.size(3) == kDimV, - "out must have shape [B, 1, HV, 128]"); - } + if (output.has_value()) + { + at::Tensor const& provided = *output; + TORCH_CHECK( + provided.is_cuda() && provided.scalar_type() == at::kBFloat16, "out must be a CUDA bfloat16 tensor"); + TORCH_CHECK(provided.is_contiguous(), "out must be contiguous"); + TORCH_CHECK(provided.dim() == 4 && provided.size(0) == B && provided.size(1) == 1 && provided.size(2) == HV + && provided.size(3) == kDimV, + "out must have shape [B, 1, HV, 128]"); + } + auto out = output.has_value() ? *output : at::empty({B, 1, HV, kDimV}, x_q.options());🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@cpp/tensorrt_llm/thop/kdaDecodeOp.cpp` around lines 216 - 222, Restructure the output handling around output.has_value() so validation occurs before binding out, or is performed through one unconditional validation path after allocation/binding. Keep the existing CUDA, bfloat16, contiguous, and [B, 1, HV, kDimV] requirements, while avoiding duplicated or dead checks in the allocation path.tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py (2)
225-228: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueInitialize
_lazy_handlesin the constructor.
_lazy_handlesis created in_load_lazy_safetensorsand reassigned incleanup. No constructor declares it. The current call order hides this, but any future caller that reads the attribute before a lazy load raisesAttributeError. Declare it once with the other instance state.As per coding guidelines: "initialize externally visible class members in the constructor".
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py` around lines 225 - 228, Initialize self._lazy_handles in the class constructor alongside the other instance state, using the existing empty collection type. Keep _load_lazy_safetensors and cleanup reusing this attribute rather than relying on lazy creation or first-time assignment.Source: Coding guidelines
244-244: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCorrect the return annotation.
_load_lazy_safetensorsreturnsConsumableWeightsDict, not a plaindict[str, Any]. The sibling loaders_prefetch_and_loadand_load_weights_in_parallelannotate-> ConsumableWeightsDict. Match them so callers see themark_consumedsurface.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py` at line 244, Update the return annotation of _load_lazy_safetensors from dict[str, Any] to ConsumableWeightsDict, matching _prefetch_and_load and _load_weights_in_parallel so callers expose the mark_consumed interface.tensorrt_llm/_torch/modules/mla.py (1)
534-541: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueExtend the fused-layout guard to the DeepSeek-V4 path.
The guard rejects
fuse_qkv_a_proj=Falsefor DSA becauseforward_dsa_projhard-codes the fused[q_lora_rank | kv_lora_rank | qk_rope_head_dim]split.forward_impl_with_deepseek_v4at line 1859 has the same hard-coded assumption:q, kv = self.kv_a_proj_with_mqa(hidden_states).split( [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], -1)
is_deepseek_v4is set fromsparse_algorithm == "deepseek_v4", which is a different flag fromis_dsa. No caller combines DeepSeek-V4 with the separateq_a_projlayout today, so this is not currently reachable. Adding the check keeps the two paths symmetric and converts a future shape error into a clear message.♻️ Proposed change
self.is_deepseek_v4 = sparse_algorithm == "deepseek_v4" + if self.is_deepseek_v4 and not fuse_qkv_a_proj: + # forward_impl_with_deepseek_v4 assumes the fused + # [q_a | kv_a | k_pe] projection layout. + raise NotImplementedError( + "DeepSeek-V4 requires fuse_qkv_a_proj=True; the separate " + "q_a_proj layout is not supported with DeepSeek-V4." + )Note: place this after the
is_deepseek_v4assignment at line 542.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/modules/mla.py` around lines 534 - 541, Extend the fused-layout validation after the is_deepseek_v4 assignment to reject fuse_qkv_a_proj=False when is_deepseek_v4 is true, matching the existing DSA guard. Raise a clear NotImplementedError explaining that DeepSeek-V4 requires the fused qkv_a projection layout.tensorrt_llm/models/quant_config_utils.py (1)
74-77: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winShare the exclusion-list construction with the tail block.
Lines 74-77 duplicate the logic at lines 126-132, and the two copies diverge. The tail block preserves the original
ignoreorder whenmodules_to_not_convertis absent; this copy always routes throughset(...), so the resulting order varies between runs.Extract one helper and call it from both paths.
♻️ Proposed change
def _resolve_exclude_modules(hf_quant_config: Mapping[str, Any]) -> list[str]: """Merge ``modules_to_not_convert`` and ``ignore``, preserving order.""" hf_exclude_modules = hf_quant_config.get("modules_to_not_convert") or [] ignore = hf_quant_config.get("ignore", []) return list(dict.fromkeys(list(hf_exclude_modules) + list(ignore)))Then in the MXFP4 branch:
quant_config.group_size = group_size - hf_exclude_modules = hf_quant_config.get("modules_to_not_convert", None) - quant_config.exclude_modules = list( - set((hf_exclude_modules or []) + hf_quant_config.get("ignore", [])) - ) + quant_config.exclude_modules = _resolve_exclude_modules(hf_quant_config) return🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/models/quant_config_utils.py` around lines 74 - 77, Extract the shared exclusion-list construction into a helper such as _resolve_exclude_modules, merging modules_to_not_convert and ignore while preserving first-seen order and removing duplicates. Replace both the MXFP4 assignment near hf_exclude_modules and the duplicate tail-block logic with calls to this helper, keeping quant_config.exclude_modules behavior consistent across paths.tensorrt_llm/_torch/configs/kimi_linear.py (1)
107-107: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace the validation asserts with explicit exceptions.
Python removes
assertstatements when the interpreter runs with-O. Config validation then passes silently and the failure surfaces later as an obscure runtime error. Line 107 and lines 125-126 both validate user-supplied checkpoint data.Line 125 also raises
KeyErrorinstead of a clear message whenkda_layersis absent fromlinear_attn_config.♻️ Proposed change
self.moe_router_activation_func = moe_router_activation_func - assert self.moe_router_activation_func in ("softmax", "sigmoid") + if self.moe_router_activation_func not in ("softmax", "sigmoid"): + raise ValueError( + "moe_router_activation_func must be 'softmax' or 'sigmoid', " + f"got {self.moe_router_activation_func!r}" + )And for the
linear_attn_configcheck:if linear_attn_config is not None: - assert linear_attn_config["kda_layers"] is not None - assert linear_attn_config["full_attn_layers"] is not None + for key in ("kda_layers", "full_attn_layers"): + if linear_attn_config.get(key) is None: + raise ValueError(f"linear_attn_config must define {key!r}")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/configs/kimi_linear.py` at line 107, Replace the validation asserts in the Kimi configuration initializer, including the moe_router_activation_func check and the kda_layers/linear_attn_config checks, with explicit exceptions that remain active under optimized Python execution. Validate missing linear_attn_config or kda_layers explicitly and raise a clear, descriptive error instead of allowing an implicit KeyError.tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py (2)
221-227: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winCross-check the committed conv pool as well as the SSM pool.
The final check compares
ssm_pool_fusedagainstssm_pool_seq. It does not compareconv_pool_fusedagainstconv_pool_seq. The fused kernel commits the conv windows in place after the golden token, and_promote_sequentialalready copies the promoted conv window intoconv_pool_seqon line 160. A conv-window commit bug would therefore pass this test as long as the round-2 outputs happen to match.Add
ok &= _rep("committed conv", conv_pool_fused, conv_pool_seq)after the SSM check.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py` around lines 221 - 227, Extend the committed pool cross-check after the existing `_rep("committed ssm", ssm_pool_fused, ssm_pool_seq)` assertion to also compare `conv_pool_fused` with `conv_pool_seq` using `_rep("committed conv", ...)`, preserving the accumulated `ok` result.
164-169: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert per check, and tighten the parity tolerance.
Two points:
The test accumulates results into
okand ends with a bareassert ok. On failure the report shows onlyassert False. The reader must find the printed lines in captured stdout to learn which of the three checks failed. Assert inside_repwith the metric values in the message instead.
rel < 3e-2allows a 3% relative L2 error between the fused and sequential paths. Both worlds run the same weights on the same inputs; the only expected divergence is bf16 accumulation order. A systematic 2% bookkeeping error in the replay path would pass. Consider tighteningrel, or document why 3e-2 is the required bound for this shape.♻️ Proposed change
def _rep(name, a, b): a, b = a.float(), b.float() cos = torch.nn.functional.cosine_similarity(a.flatten(), b.flatten(), dim=0).item() rel = ((a - b).norm() / (b.norm() + 1e-12)).item() - print(f" {name}: cos={cos:.6f} rel_l2={rel:.3e}") - return cos > 0.999 and rel < 3e-2 + assert cos > 0.999 and rel < 3e-2, ( + f"{name} parity failed: cos={cos:.6f} rel_l2={rel:.3e}" + )Then drop the
okaccumulator and the finalassert ok.Also applies to: 229-229
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py` around lines 164 - 169, Update _rep to assert each parity check directly, including cosine and relative-L2 values in the failure message, then remove the ok accumulator and final assert. Tighten the relative-L2 threshold below 3e-2 to an appropriate bf16 accumulation tolerance, or document the shape-specific justification if 3e-2 is necessary.tensorrt_llm/_torch/models/modeling_kimi_linear.py (2)
2612-2612: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReuse the already-computed
kda_fp8gate.Line 2574 computes
kda_fp8from_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV. Line 2612 re-reads the same variable and re-implements the same comparison. The two expressions must stay in sync, and the earlier one already drives whetherfinalize_decode_weights()was skipped. A divergence would leave the BF16 fast path unbuilt and the FP8 conversion unapplied.♻️ Proposed change
- if os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV, "1") != "0": + if kda_fp8:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_kimi_linear.py` at line 2612, In the logic surrounding finalize_decode_weights(), replace the repeated environment-variable comparison with the already-computed kda_fp8 gate from the earlier initialization. Reuse kda_fp8 directly so the BF16 fast-path and FP8 conversion decisions remain consistent.
1323-1331: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueMove the path-selection logging out of the per-step decode path.
logger.info_onceruns on every decode call, for every KDA layer. The message is emitted once, but the call, the key lookup, and the branch still execute inside the hot decode loop. Log the selected path once at load time (for example infinalize_decode_weights*()), and keep_forward_decodefree of logging.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_kimi_linear.py` around lines 1323 - 1331, The per-step path-selection logging in _forward_decode must be removed from the hot decode loop. Move the ssm_state_indices branch and its info_once messages to the appropriate load-time finalize_decode_weights*() method, preserving the indexed versus static path selection and logging each path only once.tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py (1)
50-95: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider wrapping the test body in
torch.no_grad().
KimiKDARuntimeparameters require gradients by default. The sequential reference loop and the verify call build autograd graphs fort_stepsforwards and hold activations. The other new KDA tests in this PR (for exampletests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py) use@torch.no_grad(). Add the same decorator for consistency and lower memory use.♻️ Proposed change
+@torch.no_grad() `@pytest.mark.parametrize`("batch", [1, 3]) `@pytest.mark.parametrize`("t_steps", [2, 3]) def test_kda_verify_matches_sequential_decode(batch, t_steps):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py` around lines 50 - 95, Decorate test_kda_verify_matches_sequential_decode with torch.no_grad so both the sequential reference loop and _forward_verify execute without building autograd graphs, matching the other KDA tests and reducing memory use.tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py (1)
28-46: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueConsider a module-scoped fixture for the attention pair.
Five tests call
_make_attention_pair(). Each call builds twoKimiKDALinearAttentionmodules withHIDDEN_SIZE = 7168and 96 heads, then copies the state dict. The tests only read the weights, so a@pytest.fixture(scope="module")would remove the repeated allocation and the repeatedload_state_dict. The sibling filetests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.pyalready uses a module-scopeddispatch_pairfixture.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py` around lines 28 - 46, Convert _make_attention_pair into a module-scoped pytest fixture, preserving its optimized/reference construction, state-dict synchronization, and path assertions. Update the five tests that call _make_attention_pair to receive the fixture through their test parameters and reuse the shared pair.tests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.py (1)
319-352: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
data["M"]instead of the module-levelM.
_fla_sequential_referencereads the module globalMat Line 352, butcpu_referencereadsdata["M"].make_conv_dataacceptsMas a parameter. If any caller passes a differentM, the FLA reference silently loops over the wrong token count while the CPU golden stays correct.♻️ Proposed fix
- B, H, K, V, W = data["B"], data["H"], data["K"], data["V"], data["W"] + B, H, K, V, W = data["B"], data["H"], data["K"], data["V"], data["W"] + num_spec = data["M"] T = data["T"]- for i_t in range(a + 1 + M): + for i_t in range(a + 1 + num_spec):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.py` around lines 319 - 352, Update _fla_sequential_reference to read the token-count value from data["M"] and use that local value in its processing loop, matching cpu_reference and the M supplied by make_conv_data; remove its dependency on the module-level M.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py`:
- Around line 230-242: Remove the os.path.isfile guard and its early return from
_is_kimi_k3_checkpoint. Always open config.json and let lookup, read, or parse
failures propagate before branching, while preserving the existing model_type
checks for valid configurations.
In `@tensorrt_llm/_torch/modules/fused_moe/quantization.py`:
- Around line 6295-6303: Update the class-level cache key used by the
scale-loading logic around load_expert_w3_w1_weight_scale_mxfp4 and
load_expert_w2_weight_scale_mxfp4 to include the CUDA device alongside the
existing key fields. Ensure each device maintains separate cached indices before
calling torch.ops.trtllm.shuffle_matrix, preventing indices from one device
being reused on another.
In `@tensorrt_llm/_torch/modules/kimi_k3_attn_res/_attn_res_kernels.py`:
- Around line 70-84: Add a return annotation to intree_attn_res_fwd declaring
its four-tensor tuple result, preserving the existing
torch.ops.trtllm.attn_res_fwd behavior and operator contract.
In `@tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py`:
- Around line 89-96: Remove the torch.cuda.synchronize call from
_write_identity_rope_values while preserving the identity-value writes. Rely on
current-stream ordering for subsequent kernels; do not add another device-wide
synchronization, and only introduce event-based cross-stream coordination if the
implementation demonstrates a genuine cross-stream hazard.
In `@tensorrt_llm/_torch/modules/kimi_k3_moe/_mlp.py`:
- Around line 127-147: Update the `situ_and_mul` implementation to explicitly
validate that `x` has unit stride on its last dimension before launching
`situ_and_mul_kernel`, matching the kernel’s contiguous-last-dimension contract.
Keep the existing shape and stride handling unchanged for valid inputs, and
reject non-unit last-dimension strides with the established validation
mechanism.
In `@tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py`:
- Line 499: In the Kimi K3 MoE block’s flattening and output-restoration logic,
replace both `hidden_states.view(-1, self.hidden_size)` and
`y.view(*orig_shape)` with `reshape` calls so non-contiguous tensors are handled
safely while preserving the existing shapes.
In `@tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py`:
- Around line 89-92: Initialize self.weight and self.e_score_correction_bias
deterministically during KimiK3MoEGate construction, preferably by adding and
invoking reset_parameters; ensure both parameters receive valid values before
forward can run without a weight load.
- Around line 108-129: The fused and eager routing paths use different sigmoid
arithmetic, so do not claim byte-for-byte parity. Add a CUDA parity test for the
Kimi K3 routing implementation that compares both paths with an appropriate
numerical tolerance, and remove or revise any exact-parity
assertion/documentation associated with `_use_fused_routing`.
In `@tensorrt_llm/_torch/modules/kimi_kda/_kda_decode.py`:
- Line 21: Bound the module-level _DUMMY_CACHE with a precise key/value type and
prevent batch-size-dependent onorm_g tensors from accumulating: update the
onorm_g fallback in the decode path to use one reusable maximum-batch buffer and
slice it for the requested batch, or remove onorm_g from the cache. Preserve the
existing contiguous behavior before kernel launch.
In `@tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py`:
- Around line 95-101: Update is_intree_prefill_available and
is_intree_mtp_available to catch ImportError only for unavailable modules. Let
other loader failures propagate, or if the probes must remain non-throwing, log
those exceptions at warning level before returning False; remove the broad
BLE001 handlers while preserving the expected ImportError fallback.
In `@tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py`:
- Around line 239-244: Initialize dt_bias in the KDA mixer constructor instead
of leaving it uninitialized from torch.empty. Use an intentional default
compatible with MetaInitMode, such as the existing supported initialization
patterns, before KDA execution; keep A_log, f_a_proj, and f_b_proj unchanged.
In `@tensorrt_llm/_torch/modules/mla.py`:
- Around line 604-610: Update MLA construction call sites in test_mla_helix.py
and test_mla_registry.py to provide rms_norm_eps explicitly or populate
config.pretrained_config.rms_norm_eps, including cases that omit config, so
initialization no longer reaches the ValueError in the MLA constructor.
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 2250-2254: Update the Kimi manager kwargs construction to select
conv_state_layout conditionally for MambaHybridCacheManagerV2, using the KDA
q_k_v layout, while retaining model_type="qwen3_next" only for the legacy
manager path. Pass the resulting conditional kwargs at the call site, ensuring
the V2 route includes conv_state_layout and does not include model_type.
In `@tensorrt_llm/_torch/pyexecutor/config_utils.py`:
- Around line 389-397: Update the Kimi linear dtype override block guarded by
is_kimi_linear so it logs at warning level when quant_config supplied an
explicit mamba_ssm_cache_dtype, while retaining the existing info-level log for
non-explicit defaults. Preserve the fp32 override behavior and identify the
user-supplied value through kv_cache_config.mamba_ssm_cache_dtype.
In `@tensorrt_llm/models/quant_config_utils.py`:
- Around line 63-71: Update the MXFP4 selection condition in the surrounding
quantization configuration logic to require the expected 4-bit float weight
settings even when format is mxfp4-pack-quantized, rejecting incompatible
declarations instead of configuring W4A16_MXFP4. In the same branch, replace the
direct weights_quant_config["group_size"] access with the established validation
path so a missing or unsupported group size raises the actionable ValueError
rather than KeyError.
In `@tests/unittest/_torch/executor/test_mamba_cache_manager.py`:
- Around line 661-678: Correct test_kimi_defaults_to_v2 to match
KimiLinearForCausalLM.get_model_defaults: assert use_kv_cache_manager_v2 is
False unless the model defaults are intentionally changed, and update the
expected manager accordingly. Replace _hybrid_model_config() in the final
get_kv_cache_manager_cls assertion with the existing Kimi linear model
configuration helper, preserving the test’s intended Kimi coverage.
- Around line 197-227: Align the Kimi tests and production contracts: update
test_kimi_explicit_v2_manager_uses_qkv_convolution_layout and its setup to use
model_type="qwen3_next", and ensure KimiLinearForCausalLM.get_model_defaults()
sets use_kv_cache_manager_v2=True so test_kimi_defaults_to_v2 resolves V2
automatically. Preserve the already-correct
test_kimi_kda_cache_params_preserve_qkv_and_fp32_state_geometry behavior.
In `@tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py`:
- Around line 58-64: In dispatch_pair, replace the optimized prefill path
assertion in
tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py:58-64 with
pytest.skip("optimized KDA prefill op unavailable") when
optimized.prefill_kernel_path is not "optimized"; apply the same conditional
skip to the corresponding fixture in
tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.py:44-50,
while preserving the reference-path assertion.
In `@tests/unittest/_torch/modules/kimi_kda/test_kda_decode_op.py`:
- Around line 302-305: Update the test fixture’s A_log tensor to use the
per-head shape [num_heads] while leaving dt_bias at [projection_size]. Keep the
existing dtype, device, and initialization unchanged so the fixture matches the
fused kernel contract.
In `@tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py`:
- Around line 19-24: Add both Kimi K3 MoE test modules to the appropriate GPU
test target list, such as the GB300 test-db YAML. Extend coverage for
non-contiguous last dimensions, the KimiK3RMSNorm fused and
KIMI_K3_FUSED_RMSNORM=0 paths, and NonSituActivation, while preserving the
existing test parametrizations.
---
Nitpick comments:
In `@cpp/tensorrt_llm/thop/kdaDecodeOp.cpp`:
- Around line 216-222: Restructure the output handling around output.has_value()
so validation occurs before binding out, or is performed through one
unconditional validation path after allocation/binding. Keep the existing CUDA,
bfloat16, contiguous, and [B, 1, HV, kDimV] requirements, while avoiding
duplicated or dead checks in the allocation path.
In `@tensorrt_llm/_torch/configs/kimi_linear.py`:
- Line 107: Replace the validation asserts in the Kimi configuration
initializer, including the moe_router_activation_func check and the
kda_layers/linear_attn_config checks, with explicit exceptions that remain
active under optimized Python execution. Validate missing linear_attn_config or
kda_layers explicitly and raise a clear, descriptive error instead of allowing
an implicit KeyError.
In `@tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py`:
- Around line 225-228: Initialize self._lazy_handles in the class constructor
alongside the other instance state, using the existing empty collection type.
Keep _load_lazy_safetensors and cleanup reusing this attribute rather than
relying on lazy creation or first-time assignment.
- Line 244: Update the return annotation of _load_lazy_safetensors from
dict[str, Any] to ConsumableWeightsDict, matching _prefetch_and_load and
_load_weights_in_parallel so callers expose the mark_consumed interface.
In `@tensorrt_llm/_torch/models/modeling_kimi_linear.py`:
- Line 2612: In the logic surrounding finalize_decode_weights(), replace the
repeated environment-variable comparison with the already-computed kda_fp8 gate
from the earlier initialization. Reuse kda_fp8 directly so the BF16 fast-path
and FP8 conversion decisions remain consistent.
- Around line 1323-1331: The per-step path-selection logging in _forward_decode
must be removed from the hot decode loop. Move the ssm_state_indices branch and
its info_once messages to the appropriate load-time finalize_decode_weights*()
method, preserving the indexed versus static path selection and logging each
path only once.
In `@tensorrt_llm/_torch/modules/kimi_k3_attn_res/kimi_k3_attn_res.py`:
- Line 227: Remove the unused self._optimized_extension initialization from the
__init__ method of the relevant attention module, since no other code accesses
it; do not add a replacement unless the attribute has a documented purpose.
- Around line 79-83: Update the module docstring for the chunked torch reference
to remove the stale external path or replace it with the correctly spelled,
repository-valid reference path; retain the description of byte-identical
reference outputs.
In `@tensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.py`:
- Around line 26-44: Update _meta_safe_cast_dtype with annotations for its
module, dtype, and None return value, and annotate the nested _cast helper
consistently with the project’s typing conventions. Remove the redundant local
“import torch as _torch” and use the existing module-level torch import
throughout.
In `@tensorrt_llm/_torch/modules/kimi_k3_moe/_mxfp4.py`:
- Around line 82-84: Replace the `diffs`/`argmin` nearest-code search in the
`_FP4_VALUES` encoding flow with `torch.bucketize` on `scaled.abs()` using
sorted FP4 magnitude midpoints, then restore the sign bit to produce the uint8
`codes`. Preserve the existing nearest-magnitude mapping while avoiding
materialization of the 16x expanded tensor.
In `@tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py`:
- Around line 646-659: Make the two-way mutation explicit in
copy_hf_moe_block_weights by adding a canonicalize_source: bool = True parameter
and documenting its effect. Guard the expert.w1, expert.w2, and expert.w3 weight
overwrites so they occur only when canonicalize_source is enabled, while
preserving the existing K3 expert-bank storage behavior.
In `@tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py`:
- Around line 165-190: Cache the DeepSeekV3MoeRoutingMethod created by the
routing_method property after its validation checks succeed, returning the
cached instance on subsequent accesses instead of constructing a new router each
time. Initialize or store the cache on the owning class and preserve all
existing validation behavior.
In `@tensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.py`:
- Around line 51-69: Remove the redundant local import torch as _torch from
_meta_safe_cast_dtype and update its empty_like reference to use the existing
module-level torch import, while preserving the current casting behavior.
In `@tensorrt_llm/_torch/modules/mamba/mamba2_metadata.py`:
- Around line 409-413: Update _prepare_replay_work_items to reuse the refreshed
int64 mirror by assigning cache_slot_idx from
self.state_indices_long[num_contexts:batch_size] instead of converting the int32
state_indices slice with to(torch.long). Preserve the existing cache_slot slice
and downstream behavior.
In `@tensorrt_llm/_torch/modules/mla.py`:
- Around line 534-541: Extend the fused-layout validation after the
is_deepseek_v4 assignment to reject fuse_qkv_a_proj=False when is_deepseek_v4 is
true, matching the existing DSA guard. Raise a clear NotImplementedError
explaining that DeepSeek-V4 requires the fused qkv_a projection layout.
In `@tensorrt_llm/models/quant_config_utils.py`:
- Around line 74-77: Extract the shared exclusion-list construction into a
helper such as _resolve_exclude_modules, merging modules_to_not_convert and
ignore while preserving first-seen order and removing duplicates. Replace both
the MXFP4 assignment near hf_exclude_modules and the duplicate tail-block logic
with calls to this helper, keeping quant_config.exclude_modules behavior
consistent across paths.
In `@tests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.py`:
- Around line 319-352: Update _fla_sequential_reference to read the token-count
value from data["M"] and use that local value in its processing loop, matching
cpu_reference and the M supplied by make_conv_data; remove its dependency on the
module-level M.
In `@tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py`:
- Around line 221-227: Extend the committed pool cross-check after the existing
`_rep("committed ssm", ssm_pool_fused, ssm_pool_seq)` assertion to also compare
`conv_pool_fused` with `conv_pool_seq` using `_rep("committed conv", ...)`,
preserving the accumulated `ok` result.
- Around line 164-169: Update _rep to assert each parity check directly,
including cosine and relative-L2 values in the failure message, then remove the
ok accumulator and final assert. Tighten the relative-L2 threshold below 3e-2 to
an appropriate bf16 accumulation tolerance, or document the shape-specific
justification if 3e-2 is necessary.
In `@tests/unittest/_torch/modeling/test_kimi_kda_verify_parity.py`:
- Around line 50-95: Decorate test_kda_verify_matches_sequential_decode with
torch.no_grad so both the sequential reference loop and _forward_verify execute
without building autograd graphs, matching the other KDA tests and reducing
memory use.
In `@tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py`:
- Around line 28-46: Convert _make_attention_pair into a module-scoped pytest
fixture, preserving its optimized/reference construction, state-dict
synchronization, and path assertions. Update the five tests that call
_make_attention_pair to receive the fixture through their test parameters and
reuse the shared pair.
In `@tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py`:
- Around line 56-66: Extend test_situ_and_mul_strided_rows with a
non-contiguous-last-dimension input whose stride(-1) is not 1, and assert the
documented behavior enforced by the _mlp.py guard. Keep the existing row-strided
case intact and verify the new case using the expected validation or failure
assertion rather than comparing normal output.
In `@tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py`:
- Around line 375-376: Define a dedicated mismatch exception for check_accuracy
and raise it when the accuracy comparison detects a mismatch, preserving the
existing message. Update both test_fc1_swap_mutation_breaks_accuracy
(tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py:375-376) and
test_swiglu_act_mutation_breaks_accuracy
(tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py:406-407) to expect
that specific exception instead of built-in Exception.
In `@tests/unittest/models/test_quant_config_utils.py`:
- Around line 289-313: Extend coverage for
update_quant_config_from_compressed_tensors by adding a case that omits the
format value and explicitly sets input_activations to None, while retaining the
existing 4-bit float group quantization fields. Assert the same MXFP4
quantization behavior and relevant configuration outcomes so the second branch
condition is exercised.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5a46443c-07d4-49a9-8005-468fcf8c965a
📒 Files selected for processing (58)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cppcpp/tensorrt_llm/thop/kdaDecodeOp.cpptensorrt_llm/_torch/configs/__init__.pytensorrt_llm/_torch/configs/kimi_linear.pytensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/checkpoints/hf/weight_loader.pytensorrt_llm/_torch/models/modeling_kimi_linear.pytensorrt_llm/_torch/modules/fused_moe/communication/communication_factory.pytensorrt_llm/_torch/modules/fused_moe/communication/deep_ep_low_latency.pytensorrt_llm/_torch/modules/fused_moe/configurable_moe.pytensorrt_llm/_torch/modules/fused_moe/create_moe.pytensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.pytensorrt_llm/_torch/modules/fused_moe/moe_op_backend.pytensorrt_llm/_torch/modules/fused_moe/moe_scheduler.pytensorrt_llm/_torch/modules/fused_moe/quantization.pytensorrt_llm/_torch/modules/kimi_k3_attn_res/__init__.pytensorrt_llm/_torch/modules/kimi_k3_attn_res/_attn_res_kernels.pytensorrt_llm/_torch/modules/kimi_k3_attn_res/kimi_k3_attn_res.pytensorrt_llm/_torch/modules/kimi_k3_mla/__init__.pytensorrt_llm/_torch/modules/kimi_k3_mla/kimi_k3_mla_attention.pytensorrt_llm/_torch/modules/kimi_k3_moe/__init__.pytensorrt_llm/_torch/modules/kimi_k3_moe/_mlp.pytensorrt_llm/_torch/modules/kimi_k3_moe/_moe_kernels.pytensorrt_llm/_torch/modules/kimi_k3_moe/_mxfp4.pytensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.pytensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.pytensorrt_llm/_torch/modules/kimi_kda/__init__.pytensorrt_llm/_torch/modules/kimi_kda/_kda_decode.pytensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.pytensorrt_llm/_torch/modules/kimi_kda/kimi_kda_mixer.pytensorrt_llm/_torch/modules/mamba/mamba2_metadata.pytensorrt_llm/_torch/modules/mla.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/config_utils.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/_torch/utils.pytensorrt_llm/mapping.pytensorrt_llm/models/quant_config_utils.pytests/integration/test_lists/test-db/l0_gb300_multi_gpus.ymltests/unittest/_torch/executor/test_mamba_cache_manager.pytests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.pytests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.pytests/unittest/_torch/modeling/test_kimi_kda_verify_parity.pytests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.pytests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.pytests/unittest/_torch/modules/kimi_kda/test_kda_decode_op.pytests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.pytests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.pytests/unittest/_torch/modules/moe/test_kimi_k3_mlp.pytests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.pytests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.pytests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.pytests/unittest/_torch/modules/test_kimi_k3_mla_backend.pytests/unittest/_torch/thop/serial/test_moe.pytests/unittest/models/test_quant_config_utils.py
dcd9e49 to
b5b1db2
Compare
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
|
/bot run --disable-fail-fast |
|
PR_Github #64438 [ run ] triggered by Bot. Commit: |
|
PR_Github #64438 [ run ] completed with state
|
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/executor/test_mamba_cache_manager.py`:
- Line 231: Add the return annotation -> SimpleNamespace to the
_kimi_model_config function definition, preserving its existing behavior and
ensuring the referenced type is available in the module.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: aed144fe-b5e3-4e0c-a431-889642210e98
📒 Files selected for processing (3)
tests/unittest/_torch/executor/test_mamba_cache_manager.pytests/unittest/_torch/modules/kimi_kda/test_kda_decode_op.pytests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py
🚧 Files skipped from review as they are similar to previous changes (2)
- tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py
- tests/unittest/_torch/modules/kimi_kda/test_kda_decode_op.py
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
Test-only + test-list wiring, no product code, so the risk is bounded to CI. Verified the three new -k selectors resolve on main (test_kimi_kda_verify_parity.py, test_kda_mtp_decode_cute_parity.py, test_kimi_kda_fused_verify_parity.py) and that the two .py paths added to l0_gb300_multi_gpus.yml (test_kda_prefill_op.py, kimi_k3_attn_res/test_attn_res_op.py) already exist — no empty selection, no missing path.
Approving over the 2 unresolved threads: @xxi-nv's maintainability point was answered and acted on (12 cases -> 4 at 6938d9c12), and the CodeRabbit one is a return annotation. One real finding left inline on test_moe.py — not blocking.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/unittest/_torch/modules/moe/test_moe_backend.py`:
- Around line 1269-1284: Add test_trtllm_bf16_dsv3_routing_kimi_k3_shape to
every applicable test-db CI list so the new standalone test is selected, while
leaving l0_gb300_multi_gpus.yml and qa/ unchanged because they have no matching
entry.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 38ce620b-6f52-4565-86e0-b321c0e46102
📒 Files selected for processing (1)
tests/unittest/_torch/modules/moe/test_moe_backend.py
|
/bot run |
Validation summary (out-of-CI, GB300 node)
|
Port the K3 kernel/module unit tests from the feature branch onto the KimiLinear model base, and wire the K3 op-parity and verify-parity suites into l0_gb300_multi_gpus (post_merge stage). Ported as-is: - tests/unittest/_torch/modules/kimi_kda/test_kda_decode_op.py - tests/unittest/_torch/modules/moe/test_kimi_k3_mlp.py - tests/unittest/_torch/modules/moe/test_kimi_k3_moe_gate.py - tests/unittest/_torch/modules/moe/test_kimi_k3_situ_and_mul.py - tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py Ported additive deltas: - test_mamba_cache_manager.py: KimiLinear cache-geometry and V2-manager wiring tests (BlockReuseConfig usage kept on the main-side API) - test_moe.py: Kimi-K3 896-expert/topK-16 no-groups routing shape Deliberately not ported (depend on source changes that ride in later PRs): the mxe2m1 do_finalize signature change and the separated-routing no-finalize host-scales test, the fused finalize+AllReduce+RMSNorm tests, the cute-dsl MLA backend fallback tests, and the fused-A-GEMM MXFP8 epilogue tests. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
… backend PR tests/unittest/_torch/modules/test_kimi_k3_mla_backend.py imports _KIMI_K3_MLA_GEN_BACKEND_ENV / _select_mla_generation_backend from kimi_k3_mla_attention, but that helper ships with the later cute-dsl MLA generation-backend PR and does not exist on main, so the file fails collection (l0_cpu sweeps unittest/_torch/modules). Drop it here; the backend-selection PR carries its own test. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…racts, A_log shape, trim MLP parametrization - test_mamba_cache_manager.py: the Kimi default is the Mixed manager, not V2 (get_model_defaults declares no use_kv_cache_manager_v2), so rewrite test_kimi_defaults_to_v2 as test_kimi_defaults_to_mixed_manager and check the manager class against a Kimi config instead of the Qwen3.5 hybrid config. Split the explicit-V2 test: constructor geometry stays green; the [q|k|v] conv_state_layout contract is a strict xfail until the runtime route stops passing model_type unconditionally. - test_kda_decode_op.py: A_log dummy is per-head ([num_heads]), matching the mixer parameter, instead of [projection_size]. - test_kimi_k3_mlp.py: trim the fused-vs-unfused parity matrix from 12 to 4 cases (decode/prefill token shapes x linear_beta set/None); the half-swap mutation control is unchanged. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
The RoutingDSv3_noGroups_896x16 param added to thop/serial/test_moe.py::TestMoeFp4 never ran: the whole class carries an unconditional skip pointing new coverage at modules/moe/test_moe_backend.py. Drop the dead param and cover the shape there instead: test_trtllm_bf16_dsv3_routing_kimi_k3_shape reuses the existing TRTLLM-Gen BF16 DeepSeekV3 routing test with 896 experts / top_k 16 (single group), exercising the (1024, 32) routing-kernel tier including the cooperative small-batch kernel at num_tokens 1. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
… test helper return type The moved routing case test_trtllm_bf16_dsv3_routing_kimi_k3_shape is now verified green on GB300 (both parametrizations), so wire it into l0_gb300_multi_gpus.yml alongside the other K3 unit suites. Also add the missing return annotation on _kimi_model_config flagged in review. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…pending runtime support The prefix-reuse debug log and the aggregated recurrent-cache status counters these tests assert are not implemented in MambaHybridCacheManagerV2 on this branch; premerge runs fail them on every machine type. Mark the affected cases strict-xfail (matching the existing conv_state_layout xfail) until the runtime-side logging and accounting land. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
/bot run |
d661a8e to
9111f44
Compare
|
PR_Github #64660 [ run ] triggered by Bot. Commit: |
|
PR_Github #64660 [ run ] completed with state |
10 participants
Description
Ports the Kimi K3 unit-test suites that accompany the model PR, and wires the
K3 kernel/op parity tests into the GB300 multi-GPU L0 list.
Base: #17269 (KimiLinear model) has merged; this PR is rebased onto
main. The wired selections have verified green runs on a GB300 node.
New test files:
Additive deltas:
the current BlockReuseConfig API)
Deliberately excluded (they depend on source changes that ride in later PRs):
the moe-backend separated-routing test, mnnvl fused-finalize tests, fmha
cute-dsl MLA page-index tests, and the dsv3 mxfp8 variants.
Test-list wiring
l0_gb300_multi_gpus.yml gains the K3 verify-parity
-kselections and theKDA/attention-residual op test files. That block is
stage: post_merge, sothese do not run in pre-merge CI; per test-list policy this PR stays draft
until every wired selection has a green run on representative hardware
(Blackwell sm_100 family, 4 GPUs). Results will be posted here before the PR
is marked ready. No l0_a10/QA list changes in this PR.
Test Coverage
The PR content is itself test coverage; the wired selections are
random-weight tests requiring no model checkpoint.
PR Checklist
[TRTLLM-14813][test]conventionDev Engineer Review
situ_and_mul, and MoE backend behavior.top_k=16.QA Engineer Review
Test changes
Added:
test_optimized_decode_matches_fla_referencetest_optimized_decode_updates_indexed_recurrent_state_pool_in_placetest_sm103_selector_dispatches_each_supported_head_at_boundarytest_selector_preserves_legacy_compact_heads_off_sm103test_sm103_selector_is_cuda_graph_safetest_fused_gate_up_matches_unfused_referencetest_gate_up_half_swap_mutation_breaks_accuracytest_situ_and_mul_matches_eager_referencetest_situ_and_mul_strided_rowstest_kimi_k3_mlp_fused_activation_matches_eagertest_kimi_k3_mlp_rejects_fused_flag_with_custom_activationtest_situ_and_mul_fake_registrationtest_situ_and_mul_cuda_graph_capturetest_trtllm_bf16_dsv3_routing_kimi_k3_shapeExpanded
test_trtllm_bf16_unquantized_moewith configurable expert-count andtop_kparameters.Added cache-manager regression and rank-aware logging coverage in
test_mamba_cache_manager.py.Test-list coverage
Modified
tests/integration/test_lists/test-db/l0_gb300_multi_gpus.yml.The file covers K3 parity, speculative decoding, KDA decode, attention-residual, and MoE routing integration selections. The unit tests are not individually listed in the provided test-list summary.
Verdict
Needs follow-up. CI failed, and CBTS coverage data is unavailable.