[TRTLLM-14815][feat] Enable disaggregated serving for Kimi K3 - #17334
[TRTLLM-14815][feat] Enable disaggregated serving for Kimi K3#17334brnguyen2 wants to merge 11 commits into
Conversation
|
The new peer check wants both sides to have the same layer set, but each rank only holds its own pipeline stage, and the transfer code below it takes the overlap on purpose. So this rejects hybrid models with pipeline parallelism, including a Qwen3-Next test already in pre-merge. Also, the function that seeds the replay caches on the generation side is never called. |
|
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:
WalkthroughThe change adds byte-aware bounce-buffer reservations, Mamba/KDA recurrent-state transfer validation, Kimi K3 disaggregated serving configurations, container launcher fallbacks, cache-transceiver handling, and an endpoint parity harness. ChangesKimi K3 disaggregated transfer and serving
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 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: 15
Note
Due to the large number of review comments, Critical, Major severity comments were prioritized as inline comments.
🟡 Minor comments (11)
tests/integration/defs/kimi_k3_disagg_parity.py-536-544 (1)
536-544: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAlign the message with the actual behavior.
The message states "using each server's own" model name. The code uses a single
modelvalue for bothEndpointobjects (Line 562). If the two servers report different model ids, the candidate request carries the reference model name.Either pass each side its own model name, or correct the message.
🔧 Proposed fix: per-endpoint model names
- model = model if model is not None else cand_model - if model is None: + ref_model = model + model = model if model is not None else cand_model + if model is None: print("[parity] ERROR: neither endpoint exposes /v1/models; pass --model explicitly") return 1 if cand_model is not None and cand_model != model: print( f"[parity] NOTE: endpoints serve different model names " f"({model!r} vs {cand_model!r}); using each server's own" ) + models_by_role = { + "reference": ref_model or model, + "candidate": cand_model or model, + }Then use
models_by_role[role]when constructing eachEndpoint.🤖 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/integration/defs/kimi_k3_disagg_parity.py` around lines 536 - 544, Correct the model-selection behavior around the endpoint construction flow: preserve separate model names for reference and candidate servers when their reported IDs differ, store them by role (for example via models_by_role), and pass the corresponding role-specific value when constructing each Endpoint instead of reusing the single model variable. Keep the existing fallback and mismatch handling intact.tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py-375-376 (1)
375-376: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winBoth mutation tests expect the broadest exception type.
check_accuracysignals a mismatch through an assertion, but both tests accept anyException. A kernel, loader, or packing failure therefore satisfies the mutation test for the wrong reason, and the mutation stops proving what the docstrings claim.
tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py#L375-L376: changepytest.raises(Exception, match="Mismatch percentage")topytest.raises(AssertionError, match="Mismatch percentage")intest_fc1_swap_mutation_breaks_accuracy.tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py#L406-L407: apply the same change intest_swiglu_act_mutation_breaks_accuracy.Based on coding guidelines: "Catch the narrowest exception possible ... prefer built-in exception types."
🤖 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, In tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py:375-376 and :406-407, update the pytest.raises calls in test_fc1_swap_mutation_breaks_accuracy and test_swiglu_act_mutation_breaks_accuracy to expect AssertionError instead of Exception, preserving the existing “Mismatch percentage” match.Source: Coding guidelines
examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml-67-71 (1)
67-71: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winIncomplete removal of the block-count bounce-gate documentation. The bounce gate moved from a block count to a byte threshold (
TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES, default 2 MiB). Both configs still carry residue of the old wording, which contradictsexamples/kimi_k3/disagg/ctx_config.yamlandexamples/kimi_k3/disagg/README.md.
examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml#L67-L71: restore the missing sentence subject on Line 67 and replace the "default gate of 96 blocks / under 6144 tokens at tokens_per_block=64" text with the byte-gate description.examples/kimi_k3/disagg/gen_config_no_sa.yaml#L40-L45: delete the orphaned fragment# at tokens_per_block=64). See ctx_config.yaml.on Line 42 and keep the byte-gate sentence above it.🤖 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 `@examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml` around lines 67 - 71, The bounce-gate documentation still uses obsolete block-count wording. In examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml lines 67-71, restore the sentence subject and describe the byte threshold using TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES and its 2 MiB default instead of the 96-block/6144-token text. In examples/kimi_k3/disagg/gen_config_no_sa.yaml lines 40-45, remove the orphaned “at tokens_per_block=64). See ctx_config.yaml.” fragment and retain the preceding byte-gate sentence.tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml-6-9 (1)
6-9: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winSet
TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES=1intest_disaggregated_overlap_transceiver_runtime_python_bounce. The test currently sets only the legacy block gate, so the default 2 MiB byte gate keeps short prompts on the per-block path and the coalesced-bounce assertion fails.Test coverage: The test is listed in
tests/integration/test_lists/test-db/l0_gb200_multi_gpus.yml. Coverage is insufficient until the byte gate is lowered.🤖 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/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml` around lines 6 - 9, Update the test setup for test_disaggregated_overlap_transceiver_runtime_python_bounce to set TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES=1 alongside the existing legacy block gate. Ensure short prompts use the coalesced-bounce WRITE path so the assertion exercises the intended behavior.Source: Path instructions
tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py-389-394 (1)
389-394: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winReject non-finite SiTu activation values.
float("nan")and infinity pass the current positive-value check. The native SiTu calculation can then produce invalid activations. Require finitetrtllm_gen_activation_alphaandtrtllm_gen_activation_betabefore weight creation.Proposed fix
+import math + - if (self.trtllm_gen_activation_alpha <= 0.0 - or self.trtllm_gen_activation_beta <= 0.0): + if ( + not math.isfinite(self.trtllm_gen_activation_alpha) + or not math.isfinite(self.trtllm_gen_activation_beta) + or self.trtllm_gen_activation_alpha <= 0.0 + or self.trtllm_gen_activation_beta <= 0.0 + ):🤖 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/fused_moe/fused_moe_trtllm_gen.py` around lines 389 - 394, Update the validation around trtllm_gen_activation_alpha and trtllm_gen_activation_beta to reject NaN and infinite values in addition to non-positive values. Require both parameters to be finite and strictly positive before proceeding to weight creation, while preserving the existing ValueError and reported values.tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py-396-428 (1)
396-428: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick winConfirm that
build_fused_weightstargets the device that holds the expert bank.
devicecomes fromtorch.cuda.current_device(). The bank buffers may live on a different device when the module was constructed with an explicitdeviceargument, or when the ambient CUDA device changed between construction and this call.pack_routed_expert_weightscopies each source tensor with.to(device), so the call succeeds but the fused buffers can land on a device other than the one the forward pass runs on.invoke_native_situ_moethen mixeshidden_states.devicewith the weight device.Derive the device from
self.expert_bank.w1_packed.deviceinstead, or assert that the two match.🐛 Proposed change
- device = torch.device(f"cuda:{torch.cuda.current_device()}") + device = self.expert_bank.w1_packed.device + if device.type != "cuda": + raise RuntimeError( + f"build_fused_weights requires CUDA expert-bank buffers, got {device}" + )🤖 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 396 - 428, Update build_fused_weights to derive device from self.expert_bank.w1_packed.device instead of torch.cuda.current_device(), ensuring packed weights and alpha/beta buffers remain on the expert bank’s device for invoke_native_situ_moe. Use the bank tensor’s device consistently for pack_routed_expert_weights and make_situ_alpha_beta.tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py-560-571 (1)
560-571: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
_moe_infercrashes on an empty token batch.If
xhas zero rows,outputsstays empty.sorted_tokens.new_empty(0)then produces a 1-D tensor of shape(0,). Line 565 calls.view(*topk_ids.shape, -1), and PyTorch cannot infer the-1dimension from a zero-element tensor. The call raisesRuntimeErrorinstead of returning an empty result.Return an empty result with the correct trailing dimension before the dispatch loop.
🐛 Proposed guard
def _moe_infer( self, x: torch.Tensor, topk_ids: torch.Tensor, topk_weight: torch.Tensor, ) -> torch.Tensor: @@ + if x.shape[0] == 0: + return x.new_empty((0, self.moe_hidden_size)) cnts = topk_ids.new_zeros((topk_ids.shape[0], self.num_experts))🤖 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 560 - 571, Update _moe_infer to handle an empty token batch before the dispatch loop, returning an empty tensor with the expected output trailing dimension instead of constructing sorted_tokens.new_empty(0) and reshaping it. Preserve the existing dispatch and aggregation path for non-empty inputs.tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py-192-251 (1)
192-251: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAlign fused and eager sigmoid semantics
noaux_tc_opuses0.5 * tanhf(0.5 * logits) + 0.5, while the eager path usestorch.sigmoid. With bias near a selection boundary, these formulas can select different experts. Align the implementations or add a CUDA parity test with near-tie logits. The raw-score gather and1e-20normalization match.🤖 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 192 - 251, Align the fused routing path in forward with the eager _score semantics by ensuring noaux_tc receives or processes scores using the same sigmoid formula as the CPU/reference path; do not change the raw-score gather or 1e-20 renormalization behavior. If the custom op cannot be changed, add CUDA parity coverage using near-tie logits and update the implementation so fused and eager expert selection matches.tensorrt_llm/_torch/modules/kimi_k3_attn_res/__init__.py-5-15 (1)
5-15: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winFix the stale and misspelled path in the module docstring.
The docstring points at
exisiting_optimization_work/Attention_residual(misspelled "exisiting")._attn_res_kernels.pystates the kernel is now source-integrated atcpp/tensorrt_llm/kernels/kimiK3AttnRespluscpp/tensorrt_llm/thop/attnResOp.cpp. Update this docstring to the in-tree locations so readers do not look for a directory that is not in the repository.🤖 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/__init__.py` around lines 5 - 15, Update the module docstring near the fused `attn_res_fwd` description to replace the stale misspelled `exisiting_optimization_work/Attention_residual` references with the current in-tree locations `cpp/tensorrt_llm/kernels/kimiK3AttnRes` and `cpp/tensorrt_llm/thop/attnResOp.cpp`; preserve the existing behavioral and fallback descriptions.tensorrt_llm/_torch/modules/kimi_k3_attn_res/_attn_res_kernels.py-49-51 (1)
49-51: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winCorrect the docstring and add the missing return annotation.
Line 50 states the kernel is "Blackwell sm_100 only", but Line 51 also accepts
103. The module docstring at Lines 5-6 already says sm_100/sm_103. Align Line 50 with the code.
intree_attn_res_fwdhas no return annotation. It returns a 4-tuple of tensors per its docstring; annotate it.Based on the coding guideline "Annotate every function".📝 Proposed fix
def is_attn_res_optimized_supported() -> bool: - """The optimized ``attn_res_fwd`` kernel is Blackwell sm_100 only.""" + """The optimized ``attn_res_fwd`` kernel is Blackwell sm_100/sm_103 only.""" return get_attn_res_sm_version() in (100, 103)def intree_attn_res_fwd( layer_residual: torch.Tensor, block_residual: torch.Tensor, res_weight: torch.Tensor, rms_weight: torch.Tensor, rms_eps: float, -): +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:Also applies to: 70-76
🤖 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/_attn_res_kernels.py` around lines 49 - 51, Update the docstring of is_attn_res_optimized_supported to state that the optimized attn_res_fwd kernel supports Blackwell sm_100 and sm_103, matching its existing return condition and module documentation. Add an explicit return annotation to intree_attn_res_fwd describing its documented 4-tuple of tensors, and preserve its current behavior.Source: Coding guidelines
tensorrt_llm/_torch/pyexecutor/config_utils.py-600-608 (1)
600-608: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winPreserve Kimi configuration overrides and dtype metadata.
Pass
**kwargstoKimiLinearConfig.from_dict. Copy top-leveldtypeandtorch_dtypeintotext_dictwhen the text config does not define them. The current Kimi K3 dtype isbfloat16, so the fallback masks this loss today, but other composite configs can select the wrong 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/pyexecutor/config_utils.py` around lines 600 - 608, Update the Kimi configuration branch around KimiLinearConfig.from_dict to pass the existing **kwargs overrides. Before constructing text_dict’s KimiLinearConfig, propagate top-level dtype and torch_dtype into text_dict only when those keys are absent there, preserving text-config values and dtype metadata for composite configurations.
🧹 Nitpick comments (24)
tests/unittest/_torch/modeling/test_kimi_kda_fused_verify_parity.py (1)
192-229: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAssert per stage instead of one aggregated boolean.
The test accumulates
okacross three comparisons and asserts once at Line 229. If the test fails, the report shows onlyassert False. Thecos/rel_l2values print to stdout, which pytest captures by default. A developer cannot tell which round failed from the failure output.Assert after each comparison with a message. Consider also comparing
conv_pool_fusedagainstconv_pool_seqin the committed-state cross-check; only the SSM pool is compared today, so a conv-window commit bug stays undetected.♻️ Proposed refactor
-def _rep(name, a, b): +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}: cos={cos:.6f} rel_l2={rel:.3e}"Then drop the
okaccumulator and the trailingassert ok.🤖 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 192 - 229, Update the parity test to assert each comparison immediately with a stage-specific failure message, removing the ok accumulator and trailing assert. In the committed-state cross-check after _promote_sequential, compare both ssm_pool_fused and conv_pool_fused against their sequential counterparts so convolution state mismatches are detected.tests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.py (2)
60-61: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the optional fields as
int | None.
num_shared_expertsandrouted_expert_hidden_sizedefault toNonebut are annotatedint. Python 3.10+ is the project target, so useint | None.Based on coding guidelines: "use precise types instead of
dict/object/Any" and "prefer built-in generic types and|".♻️ Proposed change
- num_shared_experts: int = None - routed_expert_hidden_size: int = None + num_shared_experts: int | None = None + routed_expert_hidden_size: int | None = None🤖 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 60 - 61, Update the annotations for num_shared_experts and routed_expert_hidden_size to int | None, preserving their existing defaults of None.Source: Coding guidelines
386-403: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse
monkeypatch.setattrfor theinvoke_native_situ_moeswap.The manual assign/
try/finallyworks, butmonkeypatch.setattrrestores the attribute even when the test process is interrupted between the assignment and thetryblock, and it removes the need for theorigbookkeeping.🤖 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 386 - 403, Update the `invoke_native_situ_moe` replacement in the test to use the pytest `monkeypatch.setattr` fixture instead of manual assignment with `orig` and `try`/`finally`. Preserve the `swiglu_invoke` wrapper behavior while letting `monkeypatch` handle restoration.tests/unittest/_torch/modeling/test_kda_mtp_decode_cute_parity.py (1)
319-352: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueRead
Mfromdatainside_fla_sequential_reference.
cpu_referenceandcute_runreaddata["M"], but this helper reads the module-level constantM. If a caller builds data with a differentnum_spec, the reference silently processes the wrong token count. BindMlocally fromdata.♻️ Proposed change
- 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 bind M from data["M"] locally before using it in the token-processing loop, matching cpu_reference and cute_run and honoring callers with different num_spec values.tests/unittest/_torch/modules/kimi_k3_attn_res/test_attn_res_op.py (1)
36-52: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd a negative case for the fused-path fallback contract.
_apply_attn_res_fusedreturnsNoneoutside its contract (non-bfloat16 dtype,K + 1 > 12,M > 16384, or a hidden size that is not a multiple of 1024 in [4096, 8192]). The current parameters only cover accepted shapes, so a regression that widens or narrows the guard stays undetected. Add one case withnum_snapshots=12and one with a non-bfloat16prefix_sum, and assert the helper returnsNone.🤖 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_k3_attn_res/test_attn_res_op.py` around lines 36 - 52, Extend test_fused_attn_res_matches_torch_reference coverage with negative cases for the fused-path contract: include num_snapshots=12 and a non-bfloat16 prefix_sum, then call _apply_attn_res_fused directly and assert it returns None for both cases.tests/unittest/_torch/modules/kimi_kda/test_kda_prefill_op.py (1)
28-46: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winConvert
_make_attention_pairinto a module-scoped fixture.Four tests call
_make_attention_pair()independently. Each call builds twoKimiKDALinearAttentionmodules withhidden_size=7168and 96 heads on the GPU, plus a fullload_state_dictcopy. The modules are read-only in every test, so the allocations and the state-dict copy repeat with no added coverage. A module-scoped fixture reduces the runtime and the peak device memory, which matters for the 8191-token cases intest_kda_prefill_op_partial_final_chunk_large_batch.Note: the sibling files
test_kda_prefill_state_parity.pyandtest_kda_cache_soundness.pyalready use module-scopeddispatch_pairfixtures, so this also aligns the three files.♻️ Proposed change
-def _make_attention_pair() -> tuple[KimiKDALinearAttention, KimiKDALinearAttention]: +@pytest.fixture(scope="module") +def attention_pair() -> tuple[KimiKDALinearAttention, KimiKDALinearAttention]: common = {Then take
attention_pairas a test argument instead of calling the helper.🤖 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 that constructs and returns the unchanged optimized/reference pair once per module. Update all tests that call _make_attention_pair() to accept the attention_pair fixture argument and use it instead, preserving the existing assertions and pair behavior.tests/unittest/_torch/modules/kimi_kda/test_kda_cache_soundness.py (1)
122-143: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSkip when FLA cache helpers are unavailable.
pytest.importorskip("fla")does not check these symbols. CatchImportErroraround the FLA imports and callpytest.skip(...)when either helper is unavailable.🤖 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_cache_soundness.py` around lines 122 - 143, Update _flush_tensor_cache_pins to catch ImportError while importing the FLA cache helpers fla_pci and fla_pco, and call pytest.skip(...) if either import is unavailable. Keep the TensorRT-LLM helper import and cache-churning behavior unchanged when the FLA helpers are present.examples/disaggregated/slurm/benchmark/start_server.sh (1)
25-25: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueQuote
${config_file}only; keep${trtllm_serve_cmd}unquoted.Shellcheck reports SC2086 for this line.
${trtllm_serve_cmd}must stay unquoted because it can holdpython3 -m tensorrt_llm.commands.serveand depends on word splitting.${config_file}is a single value and should be quoted so a path with spaces still works.♻️ Proposed fix
-${trtllm_serve_cmd} disaggregated -c ${config_file} -t 7200 -r 7200 +# shellcheck disable=SC2086 # trtllm_serve_cmd may hold multiple words +${trtllm_serve_cmd} disaggregated -c "${config_file}" -t 7200 -r 7200🤖 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 `@examples/disaggregated/slurm/benchmark/start_server.sh` at line 25, Update the disaggregated server command using trtllm_serve_cmd so config_file is enclosed in quotes while trtllm_serve_cmd remains unquoted for intentional word splitting.Source: Linters/SAST tools
tensorrt_llm/_torch/disaggregation/native/bounce/impl.py (1)
581-598: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUpdate the type annotations now that the list holds
Noneplaceholders.
block_bytes_per_groupreturnsOptional[int]entries, but the return type is a barelist, and the consumers still declareblock_bytes_per_group: List[int](line 60 and line 97). Precise types here document the placeholder contract and let type checkers catch a missingNoneguard.♻️ Proposed annotation fix
-def block_bytes_per_group(page_table) -> list: +def block_bytes_per_group(page_table) -> List[Optional[int]]:Then update both declarations on
VmmBounceTransport:block_bytes_per_group: List[Optional[int]]As per coding guidelines: "Annotate every function ... avoid unnecessary
Any... use preciseCallablearguments".🤖 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/disaggregation/native/bounce/impl.py` around lines 581 - 598, Update block_bytes_per_group to return a precise list type with Optional[int] entries, reflecting its None placeholders. In VmmBounceTransport, change both block_bytes_per_group declarations to List[Optional[int]] and ensure the required typing imports are present so type checkers enforce None handling.Source: Coding guidelines
tensorrt_llm/_torch/disaggregation/native/auxiliary.py (1)
62-70: 🗄️ Data Integrity & Integration | 🔵 Trivial | 💤 Low valueConsider validating equal buffer counts before the elementwise arithmetic.
compute_aux_transfer_descspairs buffers positionally between two independently builtAuxBufferMetainstances. If the two sides ever expose a different number of aux buffers, numpy raises a broadcast error with no context, and equal-length-but-reordered lists would pair the wrong buffers silently. A short guard makes the cross-process contract explicit.♻️ Proposed guard
+ if src_meta.item_sizes.size != dst_meta.item_sizes.size: + raise ValueError( + "compute_aux_transfer_descs: aux buffer count differs " + f"(local={src_meta.item_sizes.size}, peer={dst_meta.item_sizes.size})" + ) src_ptrs = src_meta.ptrs + src_meta.item_sizes * src_slot🤖 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/disaggregation/native/auxiliary.py` around lines 62 - 70, Update compute_aux_transfer_descs to validate that src_meta and dst_meta expose equal buffer counts before computing pointer offsets or elementwise sizes. Raise a clear, contextual error when the counts differ, while preserving positional pairing for matching counts.tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py (2)
1678-1734: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate
_maybe_promote_drafter_states.The helper has no parameter types and no return type. Add
scheduled_batch: ScheduledRequests,attn_metadata: Optional["AttentionMetadata"], and-> None.As per coding guidelines: "Annotate every function, use
Nonefor procedures".♻️ Proposed annotation
- def _maybe_promote_drafter_states(self, scheduled_batch, attn_metadata): + def _maybe_promote_drafter_states( + self, + scheduled_batch: ScheduledRequests, + attn_metadata: Optional["AttentionMetadata"]) -> None:🤖 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/pyexecutor/mamba_cache_manager.py` around lines 1678 - 1734, Annotate _maybe_promote_drafter_states with scheduled_batch: ScheduledRequests, attn_metadata: Optional["AttentionMetadata"], and a -> None return type. Reuse the existing ScheduledRequests and AttentionMetadata imports or forward-reference conventions without changing the helper’s behavior.Source: Coding guidelines
417-437: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valuePrefer
ValueErroroverassertfor constructor validation.The three
assertstatements guard a user-reachable configuration contract. Python removes them under-O. RaiseValueErrorso the misconfiguration is always rejected.As per coding guidelines: "raise
ValueErrorrather than assertions".🤖 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/pyexecutor/mamba_cache_manager.py` around lines 417 - 437, In the constructor validation for _use_kda_replay_update, replace all three assert checks with ValueError raises so these configuration constraints remain enforced under optimized Python execution. Preserve the existing conditions and descriptive messages for mutual exclusivity, required speculative decoding, and matching replay width.Source: Coding guidelines
examples/kimi_k3/disagg/ctx_config.yaml (1)
47-61: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueClarify the "512 MiB" throughput reference.
Line 50-51 states "512 MiB measured at ~455 GB/s/GPU". Lines 57-60 then state that 512 is undersized and silently degrades every 8k request. A reader can take the first sentence as a recommended value. Report the measured bandwidth without the 512 MiB figure, or label it explicitly as a historical measurement.
The same sentence appears in
examples/kimi_k3/disagg/gen_config_no_sa.yamlat Line 39.🤖 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 `@examples/kimi_k3/disagg/ctx_config.yaml` around lines 47 - 61, Clarify the Fabric-VMM bounce-buffer comment in both ctx_config.yaml and gen_config_no_sa.yaml by removing the ambiguous “512 MiB measured” wording or explicitly labeling it as a historical measurement. Ensure the documentation does not imply 512 MiB is a recommended size, consistent with kv_cache_bounce_size_mb: 1024 and the sizing explanation.tensorrt_llm/_torch/modules/kimi_k3_moe/__init__.py (1)
30-37: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueSort
__all__to clear the Ruff RUF022 hint.Only the last two entries are out of order. Swapping them satisfies the isort-style ordering.
♻️ Proposed change
__all__ = [ "KimiK3MoEGate", "KimiK3RoutedExpertBank", "KimiK3SparseMoeBlock", "MoEBlockProvenance", - "copy_hf_moe_gate_weights", "copy_hf_moe_block_weights", + "copy_hf_moe_gate_weights", ]Note that prior repository learnings record Ruff
selectas{D, E, F, I, PLE, W}, which excludesRUF. If that is still the configuration, this hint does not fail CI and the change is cosmetic.🤖 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/__init__.py` around lines 30 - 37, Reorder the final two entries in __all__ so copy_hf_moe_block_weights precedes copy_hf_moe_gate_weights, preserving all other exports unchanged.Source: Linters/SAST tools
tensorrt_llm/_torch/modules/kimi_k3_moe/_mlp.py (1)
111-149: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low valueEnforce the contiguity precondition of
situ_and_mul.The kernel addresses elements as
x_row_ptr + offsetsandx_row_ptr + offsets + d. This is correct only when the last dimension ofxhas stride 1. The docstring at Line 117 states the requirement, but the op does not check it. A caller that passes a transposed or sliced view produces silently wrong output rather than an error.The current caller at Line 243 passes a contiguous
nn.Linearoutput, so no defect exists today. Add the check to protect the op as a publictorch.ops.trtllmentry point.🛡️ Proposed guard
b, n = x.shape assert n % 2 == 0 + assert x.stride(-1) == 1, "situ_and_mul requires a contiguous last dimension" d = n // 2🤖 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/_mlp.py` around lines 111 - 149, Enforce the documented contiguous-last-dimension precondition in situ_and_mul before launching situ_and_mul_kernel. Validate that x.stride(-1) equals 1 and reject invalid strided views with a clear error; leave the existing contiguous-input execution path unchanged.tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py (1)
327-338: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueCheck the incompatible-flag combination before the hardware check.
assert_native_situ_supportedraises on non-Blackwell hardware. A caller that sets bothuse_fused_cubin=Trueandnon_situ_activation_mutation=Truetherefore sees a hardware error on most machines, not the configuration error that actually applies. Move the flag check first so the message stays accurate on every platform.♻️ Proposed reorder
if use_fused_cubin: + if non_situ_activation_mutation: + raise RuntimeError( + "non_situ_activation_mutation is a Python-reference mutation " + "control; it cannot be combined with use_fused_cubin=True" + ) # Fail before any weight processing when the platform cannot run # the fused path at all. assert_native_situ_supported( hidden_size=self.moe_hidden_size, intermediate_size=config.moe_intermediate_size, ) - if non_situ_activation_mutation: - raise RuntimeError( - "non_situ_activation_mutation is a Python-reference mutation " - "control; it cannot be combined with use_fused_cubin=True" - )🤖 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 327 - 338, In the use_fused_cubin branch of the Kimi K3 MoE block, move the non_situ_activation_mutation incompatibility check before assert_native_situ_supported. This ensures the configuration RuntimeError is raised first for the incompatible flag combination, regardless of platform, while preserving the hardware validation for other fused-cubin requests.tensorrt_llm/_torch/modules/kimi_k3_moe/_moe_kernels.py (1)
41-43: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse built-in generic types across the new
kimi_k3_moepackage. The coding guidelines require built-in generic types and|unions. The repository targets Python 3.10+, and every one of these modules already importsfrom __future__ import annotations, so the deprecatedtypingaliases are not needed.kimi_k3_moe_gate.pyalready usestorch.dtype | Nonein its signatures, so the package is currently inconsistent with itself.
tensorrt_llm/_torch/modules/kimi_k3_moe/_moe_kernels.py#L41-L43: drop thetypingimport and replaceDict[str, torch.Tensor]withdict[str, torch.Tensor],Tuple[int, int, int]withtuple[int, int, int], andOptional[int]withint | None. Note that_CACHE_PERMUTE_INDICESat Line 57 is a runtime annotation;dict[tuple, torch.Tensor]is valid at runtime on 3.10+.tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_block.py#L43-L46: keepAny, and replaceList[str]withlist[str],Tuple[int, int]withtuple[int, int], andOptional[nn.Module]/Optional[torch.device]/Optional[torch.Tensor]with| Noneunions.tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py#L25-L27: keepAny, and replaceTuple[torch.Tensor, torch.Tensor]withtuple[torch.Tensor, torch.Tensor]to match the| Nonestyle already used in__init__.Based on the guideline "prefer built-in generic types and
|" and the learning that TensorRT-LLM requires Python >=3.10, sofrom __future__ import annotationsis not needed for these forms.🤖 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/_moe_kernels.py` around lines 41 - 43, Replace deprecated typing aliases with Python 3.10 built-in generics and | None unions across the kimi_k3_moe package. In tensorrt_llm/_torch/modules/kimi_k3_moe/_moe_kernels.py lines 41-43, remove the typing import and update Dict, Tuple, and Optional usages; in kimi_k3_moe_block.py lines 43-46, retain Any while updating List, Tuple, and Optional annotations; in kimi_k3_moe_gate.py lines 25-27, retain Any and replace Tuple with tuple. Keep the existing from __future__ import annotations and runtime cache annotation valid.Sources: Coding guidelines, Learnings
tensorrt_llm/_torch/modules/kimi_k3_moe/kimi_k3_moe_gate.py (1)
165-190: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low value
routing_methodbuilds a new object on every access.Each read constructs a
DeepSeekV3MoeRoutingMethod, which in turn constructs aDeepseekv3RoutingImpl. Two consecutive reads therefore return objects that are not identical, and any consumer that caches identity or compares instances sees a difference. If a caller reads this property inside a forward pass, the allocation lands on the hot path.Build the object once and memoize it.
♻️ Proposed memoization
+ self._routing_method: DeepSeekV3MoeRoutingMethod | None = None`@property` def routing_method(self) -> DeepSeekV3MoeRoutingMethod: """Return the shared DeepSeekV3 router used by ConfigurableMoE.""" + if self._routing_method is not None: + return self._routing_method if self.moe_router_activation_func != "sigmoid": @@ - return DeepSeekV3MoeRoutingMethod( + self._routing_method = DeepSeekV3MoeRoutingMethod( top_k=self.top_k, n_group=self.num_expert_group, topk_group=self.topk_group, routed_scaling_factor=self.routed_scaling_factor, callable_e_score_correction_bias=lambda: self.e_score_correction_bias, is_fused=True, ) + return self._routing_method🤖 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, Memoize the routing object returned by Kimi K3’s routing_method property so repeated accesses reuse one DeepSeekV3MoeRoutingMethod instance instead of rebuilding it. Preserve the existing validation checks, initialize or cache the object after validation, and return the cached instance on subsequent reads.tensorrt_llm/_torch/pyexecutor/config_utils.py (1)
389-397: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueLog the dtype override at warning level when the user set it explicitly.
This branch also overrides an explicit
quant_config.mamba_ssm_cache_dtype. An info-level message is easy to miss. Uselogger.warningwhen the value came fromquant_config, so a user who requested a different SSM cache dtype learns that the request was ignored.🤖 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/pyexecutor/config_utils.py` around lines 389 - 397, Update the Kimi K3 dtype override branch around is_kimi_linear and mamba_ssm_cache_dtype to track whether the value originated from quant_config.mamba_ssm_cache_dtype, and use logger.warning for that explicit-user override while retaining logger.info for implicit/default values. Preserve the existing fp32 assignment and message context.tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py (1)
225-228: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueInitialize
_lazy_handlesin the class rather than only on first use.
self._lazy_handlesis created incleanupand in_load_lazy_safetensors. Any other reader that runs before either method raisesAttributeError. Declare it once as a class attribute or 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 _lazy_handles in the class definition or constructor of the relevant weight loader so it always exists before any reader accesses it; retain cleanup and _load_lazy_safetensors behavior while ensuring the value starts as an empty collection.tensorrt_llm/_torch/configs/kimi_linear.py (1)
107-107: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winReplace the config assertions with explicit
ValueErrorchecks.These asserts validate values that come from a checkpoint
config.json. Python removesassertstatements under-O, and the current failures are hard to diagnose: line 107 raises a bareAssertionError, and lines 125-126 raiseKeyErrorwhenkda_layersorfull_attn_layersis missing rather than reporting the missing key.Based on the coding guideline "use validators ... raise `ValueError` rather than assertions".♻️ Proposed validation change
- 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}" + )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 a non-null {key!r}" + )Also applies to: 124-127
🤖 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 assertions in the Kimi configuration validation, including the checks around moe_router_activation_func, kda_layers, and full_attn_layers, with explicit validation that raises clear ValueError exceptions. Validate missing keys before indexing so absent kda_layers or full_attn_layers report the configuration problem instead of raising KeyError, and preserve the existing accepted-value constraints.Source: Coding guidelines
tensorrt_llm/_torch/modules/mla.py (1)
534-541: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winExtend the unsupported-layout guard to the DeepSeek-V4 path.
forward_impl_with_deepseek_v4at Line 1859 splitskv_a_proj_with_mqa(hidden_states)into[q_lora_rank, kv_lora_rank + qk_rope_head_dim]. That split assumes the fused layout, exactly likeforward_dsa_proj. If a future model setsfuse_qkv_a_proj=Falsetogether withdeepseek_v4, the split silently reads wrong columns instead of failing. Add the same explicit rejection so the failure is loud.🛡️ Proposed guard extension
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." + )🤖 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 unsupported-layout guard in the MLA initialization or forward setup around `forward_impl_with_deepseek_v4` so `deepseek_v4` with `fuse_qkv_a_proj=False` raises the same explicit `NotImplementedError` as the existing `is_dsa` guard. Reuse the existing fused-layout requirement and message semantics, while preserving supported DeepSeek-V4 and DSA configurations.tensorrt_llm/_torch/modules/kimi_kda/_kda_kernels.py (1)
95-101: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winLog the swallowed import failure in the availability probes.
Both probes catch
Exceptionand discard it. When the CuTe DSL import fails, the dispatch silently falls back to FLA and no record of the cause remains. The dispatch already logs the selected paths at Lines 206-210, so the reason belongs next to it. Keep the broad catch if CuTe DSL can raise non-ImportErrortypes, but record the exception.The repository guideline requires catching the narrowest exception possible; state the justification in a comment when a broad catch is intentional.
♻️ Proposed change
def is_intree_prefill_available() -> bool: """True when the in-tree CuTe DSL prefill op can be imported.""" try: _load_prefill_module() return True - except Exception: + except Exception as exc: # CuTe DSL import can raise non-ImportError types + _log_probe_failure("kda_prefill", exc) return FalseAlso applies to: 134-140
🤖 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/_kda_kernels.py` around lines 95 - 101, Update both availability probes, is_intree_prefill_available and the corresponding probe around the second catch block, to log the caught exception instead of silently discarding it. Retain the broad Exception catch only if required for non-ImportError CuTe DSL failures, and add a concise comment documenting that justification; otherwise narrow the exception type. Place the failure logging alongside the existing dispatch-path logging.Sources: Coding guidelines, Linters/SAST tools
tensorrt_llm/_torch/modules/kimi_k3_attn_res/kimi_k3_attn_res.py (1)
180-347: 📐 Maintainability & Code Quality | 🔵 Trivial | 🏗️ Heavy liftUse one implementation for
attn_res.KimiK3AttnResidualOphas no production or test call sites, whilemodeling_kimi_linear.pymaintains separate fused dispatch and fp32 fallback logic. Route the model through this module, or remove the unused class and state the test-only scope of the remaining helpers.🤖 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 180 - 347, The attn_res implementation is duplicated and KimiK3AttnResidualOp is currently unused. Route the model’s residual-selection path in modeling_kimi_linear.py through KimiK3AttnResidualOp, including optimized and reference dispatch, and remove the duplicate fused/fp32 logic; alternatively remove the unused class and explicitly limit its helpers to tests.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ffb4453b-2cb9-42c3-87f5-04566039a37d
📒 Files selected for processing (73)
cpp/tensorrt_llm/batch_manager/kv_cache_manager_v2/blockRadixTree.cppcpp/tensorrt_llm/thop/kdaDecodeOp.cppexamples/disaggregated/slurm/benchmark/run_benchmark.shexamples/disaggregated/slurm/benchmark/start_server.shexamples/disaggregated/slurm/benchmark/start_worker.shexamples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yamlexamples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.pyexamples/kimi_k3/disagg/README.mdexamples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yamlexamples/kimi_k3/disagg/ctx_config.yamlexamples/kimi_k3/disagg/disagg_proxy_config.yamlexamples/kimi_k3/disagg/gen_config_no_sa.yamltensorrt_llm/_torch/configs/__init__.pytensorrt_llm/_torch/configs/kimi_linear.pytensorrt_llm/_torch/custom_ops/cpp_custom_ops.pytensorrt_llm/_torch/disaggregation/native/auxiliary.pytensorrt_llm/_torch/disaggregation/native/bounce/config.pytensorrt_llm/_torch/disaggregation/native/bounce/core.pytensorrt_llm/_torch/disaggregation/native/bounce/impl.pytensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.pytensorrt_llm/_torch/disaggregation/native/peer.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/transceiver.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/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yamltests/integration/defs/kimi_k3_disagg_parity.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_prefill_op.pytests/unittest/_torch/modules/kimi_kda/test_kda_prefill_state_parity.pytests/unittest/_torch/modules/moe/test_kimi_k3_situ_moe.pytests/unittest/disaggregated/region/test_aux.pytests/unittest/disaggregated/test_bounce.pytests/unittest/disaggregated/test_kda_mamba_transfer.pytests/unittest/models/test_quant_config_utils.py
…ative disaggregation Extend the Python-native disaggregation framework to transfer Kimi K3 KDA (Kimi Delta Attention) recurrent and conv states between context and generation instances: SSM mixer peer descriptors, per-request auxiliary state payloads, bounce-buffer staging for non-fabric pools, and transceiver routing for hybrid linear-attention models. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Replace the Kimi K3 disaggregated-serving fail-fast in the cache manager routing with the shared hybrid transceiver validation: the Python NIXL transceiver selects MixedMambaHybridCacheManager, whose KDA recurrent/conv states transfer through the bounce buffer. Also log the selected hybrid cache manager class once at routing time. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…i K3 parity harness - test_kda_mamba_transfer.py: KDA recurrent/conv state transfer through the native disaggregation path. - test_bounce.py / region/test_aux.py: cover bounce-buffer staging and auxiliary-state payloads for hybrid models. - kimi_k3_disagg_parity.py: two-endpoint aggregated-vs-disaggregated parity harness (multi-node; not wired into any test list here). - Update the overlap transceiver-runtime python bounce test config. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
… and benchmark wiring - examples/kimi_k3/disagg/: ctx/gen/proxy configs, SLURM benchmark harness config, and a README covering K3 disagg constraints (matched DEP16, Python NIXL transceiver, bounce-buffer sizing, UCX transport pins). Spec-decode (SA) variants land with K3 SA support. - slurm/benchmark harness: worker/server env plumbing (TRTLLM_WORKER_UCX_TLS, PATH/PYTHONPATH prepends) used by the configs. - cache_transceiver_test: K3-shaped KDA payload config and harness support. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…oss peers The peer-compatibility gate rejected any pair of ranks whose mamba/KDA layer sets differ. With pipeline parallelism each rank publishes only its own stage's layers, so the sets legitimately differ (or are disjoint, or one stage holds no recurrent layers at all) while the transfer path intersects the two sets on purpose. Drop the set-equality requirement and treat a missing recurrent layer group on either side as nothing to validate; keep the per-slot size invariants, which are layer-agnostic. Add a regression test covering partial overlap, disjoint stages, a recurrent-layer-free stage, and a size mismatch on the overlap. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
6f7aac5 to
58f0da3
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 |
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (1)
tests/unittest/disaggregated/test_kda_mamba_transfer.py (1)
575-583: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd an explicit hardware marker to
test_kda_transfer.This test constructs
MixedMambaHybridCacheManagerinstances and runs a NIXL loopback transfer. Mark it with the repository’s GPU/NIXL marker so CPU collection runs skip it instead of failing.🤖 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/disaggregated/test_kda_mamba_transfer.py` around lines 575 - 583, Add the repository’s existing GPU/NIXL pytest marker to test_kda_transfer, alongside its timeout and parameterization decorators. Keep the test body and parameter sets unchanged so CPU-only collection skips this MixedMambaHybridCacheManager and NIXL loopback test.
🤖 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 `@examples/disaggregated/slurm/benchmark/start_server.sh`:
- Around line 11-16: Strip one leading and trailing literal single quote from
TRTLLM_PATH_PREPEND and TRTLLM_PYTHONPATH_PREPEND before using them in the
prepend export logic. Apply this change in
examples/disaggregated/slurm/benchmark/start_server.sh:11-16 and
examples/disaggregated/slurm/benchmark/start_worker.sh:44-49, preserving the
existing conditional exports and PATH/PYTHONPATH composition.
In `@examples/disaggregated/slurm/benchmark/start_worker.sh`:
- Around line 85-91: Add the standard NVIDIA copyright header for 2026 at the
top of start_worker.sh, preserving the existing trtllm_serve_cmd selection and
launch commands unchanged.
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 166-170: Update SpeculativeConfig validation in
TorchLlmArgs.validate_speculative_config to reject speculative decoding for Kimi
K3 when disaggregated serving is enabled. Add the guard before selecting the
Kimi K3 disaggregated manager, while preserving existing behavior for
non-disaggregated configurations.
In `@tests/integration/defs/kimi_k3_disagg_parity.py`:
- Around line 123-128: Update _served_model to catch connection-level failures
alongside HTTPError, including urllib.error.URLError and relevant
OSError/timeout cases, so unavailable endpoints return None instead of
propagating exceptions. Report these failures with the existing parity note
format while preserving the current HTTP status reporting.
- Around line 540-544: Update the model-name handling around the parity endpoint
construction and the NOTE: either pass the resolved candidate model name to the
candidate Endpoint while retaining the reference name for the reference
Endpoint, or change the message to accurately state that the reference model
name is used for both and that --model overrides it. Ensure the log and actual
request behavior remain consistent.
- Around line 89-93: Resolve the missing kimi_k3_sa_harness dependency used by
kimi_k3_disagg_parity.py: either add that module with PROMPTS_AND_CHECKS,
_compare_logits_parity, and _parity_prompts, or update the import to reference
their actual existing definitions. Ensure the parity test imports successfully
without changing its required symbols.
In `@tests/unittest/disaggregated/region/test_aux.py`:
- Around line 138-187: Add tests/unittest/disaggregated/region/test_aux.py,
including test_aux_buffer_zero_max_draft_len_round_trip and
test_aux_transfer_layout_ctx_no_spec_gen_sa, to the appropriate CI test lists
under tests/integration/test_lists/test-db/ and
tests/integration/test_lists/qa/. Ensure the configured test commands discover
and execute these tests.
---
Nitpick comments:
In `@tests/unittest/disaggregated/test_kda_mamba_transfer.py`:
- Around line 575-583: Add the repository’s existing GPU/NIXL pytest marker to
test_kda_transfer, alongside its timeout and parameterization decorators. Keep
the test body and parameter sets unchanged so CPU-only collection skips this
MixedMambaHybridCacheManager and NIXL loopback test.
🪄 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: dc28e7de-971c-4e71-9024-972508653da5
📒 Files selected for processing (23)
examples/disaggregated/slurm/benchmark/run_benchmark.shexamples/disaggregated/slurm/benchmark/start_server.shexamples/disaggregated/slurm/benchmark/start_worker.shexamples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yamlexamples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.pyexamples/kimi_k3/disagg/README.mdexamples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yamlexamples/kimi_k3/disagg/ctx_config.yamlexamples/kimi_k3/disagg/disagg_proxy_config.yamlexamples/kimi_k3/disagg/gen_config_no_sa.yamltensorrt_llm/_torch/disaggregation/native/bounce/config.pytensorrt_llm/_torch/disaggregation/native/bounce/core.pytensorrt_llm/_torch/disaggregation/native/bounce/impl.pytensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.pytensorrt_llm/_torch/disaggregation/native/peer.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/transceiver.pytensorrt_llm/_torch/pyexecutor/_util.pytests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yamltests/integration/defs/kimi_k3_disagg_parity.pytests/unittest/disaggregated/region/test_aux.pytests/unittest/disaggregated/test_bounce.pytests/unittest/disaggregated/test_kda_mamba_transfer.py
🚧 Files skipped from review as they are similar to previous changes (16)
- examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py
- examples/kimi_k3/disagg/README.md
- tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml
- tensorrt_llm/_torch/disaggregation/native/bounce/core.py
- tensorrt_llm/_torch/disaggregation/transceiver.py
- examples/kimi_k3/disagg/ctx_config.yaml
- examples/disaggregated/slurm/cache_transceiver_test/configs/kda_payload_kimi_k3.yaml
- examples/disaggregated/slurm/benchmark/run_benchmark.sh
- tensorrt_llm/_torch/disaggregation/native/peer.py
- examples/kimi_k3/disagg/disagg_proxy_config.yaml
- examples/kimi_k3/disagg/gen_config_no_sa.yaml
- examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml
- tensorrt_llm/_torch/disaggregation/native/bounce/config.py
- tests/unittest/disaggregated/test_bounce.py
- tensorrt_llm/_torch/disaggregation/native/transfer.py
- tensorrt_llm/_torch/disaggregation/native/bounce/impl.py
| trtllm_serve_cmd="trtllm-serve" | ||
| if ! command -v trtllm-serve >/dev/null 2>&1; then | ||
| trtllm_serve_cmd="python3 -m tensorrt_llm.commands.serve" | ||
| fi | ||
|
|
||
| ${nsys_prefix} trtllm-llmapi-launch ${numa_bind_cmd} \ | ||
| trtllm-serve ${model_path} \ | ||
| ${trtllm_serve_cmd} ${model_path} \ |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Locate the console-script declarations for both commands.
fd -H -t f -a 'setup.py|pyproject.toml' | xargs rg -n -C6 'trtllm-serve|trtllm-llmapi-launch|console_scripts|\[project.scripts\]'Repository: NVIDIA/TensorRT-LLM
Length of output: 2361
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- start_worker.sh header and relevant range ---'
sed -n '1,25p;75,100p' examples/disaggregated/slurm/benchmark/start_worker.sh
printf '%s\n' '--- setup.py declaration range ---'
sed -n '465,492p' setup.py
printf '%s\n' '--- launch script metadata and first lines ---'
stat -c '%A %n' tensorrt_llm/llmapi/trtllm-llmapi-launch
sed -n '1,35p' tensorrt_llm/llmapi/trtllm-llmapi-launch
printf '%s\n' '--- declaration parser ---'
python3 - <<'PY'
from pathlib import Path
text = Path("setup.py").read_text()
start = text.index("entry_points={")
end = text.index("extras_require=", start)
block = text[start:end]
print("console_scripts contains trtllm-serve:", "trtllm-serve=" in block)
print("console_scripts contains trtllm-llmapi-launch:", "trtllm-llmapi-launch=" in block)
print("scripts contains trtllm-llmapi-launch:", "scripts=['tensorrt_llm/llmapi/trtllm-llmapi-launch']" in block)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 2998
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- launch script metadata and first lines ---'
python3 - <<'PY'
from pathlib import Path
p = Path("tensorrt_llm/llmapi/trtllm-llmapi-launch")
st = p.stat()
print("mode:", oct(st.st_mode & 0o777))
print("size:", st.st_size)
print(p.read_text().splitlines()[:35])
PY
printf '%s\n' '--- declaration parser ---'
python3 - <<'PY'
from pathlib import Path
text = Path("setup.py").read_text()
start = text.index("entry_points={")
end = text.index("extras_require=", start)
block = text[start:end]
print("console_scripts contains trtllm-serve:", "trtllm-serve=" in block)
print("console_scripts contains trtllm-llmapi-launch:", "trtllm-llmapi-launch=" in block)
print("scripts contains trtllm-llmapi-launch:", "scripts=['tensorrt_llm/llmapi/trtllm-llmapi-launch']" in block)
PY
printf '%s\n' '--- references to the launch command and packaging mode ---'
rg -n -C3 'trtllm-llmapi-launch|in-place|\\.pth|develop|editable' \
setup.py pyproject.toml README.md docs examples scripts 2>/dev/null | head -250 || true
printf '%s\n' '--- tracked-file status ---'
git status --short -- examples/disaggregated/slurm/benchmark/start_worker.shRepository: NVIDIA/TensorRT-LLM
Length of output: 23088
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- start_worker.sh setup and launch context ---'
sed -n '25,95p' examples/disaggregated/slurm/benchmark/start_worker.sh
printf '%s\n' '--- documented in-place path setup ---'
sed -n '95,125p' examples/kimi_k3/disagg/README.md
rg -n -C5 'TRTLLM_PATH_PREPEND|TRTLLM_PYTHONPATH_PREPEND' \
examples/disaggregated examples/kimi_k3 scripts 2>/dev/null | head -220
printf '%s\n' '--- launch command references outside generated/build paths ---'
rg -n -C4 'trtllm-llmapi-launch' \
--glob '!build/**' --glob '!dist/**' --glob '!node_modules/**' . | head -220
printf '%s\n' '--- relevant setup.py context ---'
sed -n '1,100p;430,492p' setup.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 40015
🏁 Script executed:
#!/bin/bash
set -eu
printf '%s\n' '--- neighboring shell headers ---'
for f in examples/disaggregated/slurm/benchmark/start_server.sh \
examples/disaggregated/slurm/benchmark/run_benchmark.sh \
tensorrt_llm/llmapi/trtllm-llmapi-launch; do
printf '%s\n' "--- $f ---"
sed -n '1,8p' "$f"
done
printf '%s\n' '--- start_worker.sh tracked metadata ---'
git ls-files --stage examples/disaggregated/slurm/benchmark/start_worker.sh
git log -1 --format='%h %ad %s' --date=short -- examples/disaggregated/slurm/benchmark/start_worker.shRepository: NVIDIA/TensorRT-LLM
Length of output: 1154
Add the NVIDIA copyright header. examples/disaggregated/slurm/benchmark/start_worker.sh is modified but has no NVIDIA copyright header for 2026.
🧰 Tools
🪛 Shellcheck (0.11.0)
[info] 90-90: Double quote to prevent globbing and word splitting.
(SC2086)
[info] 91-91: Double quote to prevent globbing and word splitting.
(SC2086)
[info] 91-91: Double quote to prevent globbing and word splitting.
(SC2086)
🤖 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 `@examples/disaggregated/slurm/benchmark/start_worker.sh` around lines 85 - 91,
Add the standard NVIDIA copyright header for 2026 at the top of start_worker.sh,
preserving the existing trtllm_serve_cmd selection and launch commands
unchanged.
There was a problem hiding this comment.
None of the launcher scripts in examples/disaggregated/slurm/benchmark (run_benchmark.sh, start_server.sh, submit-side scripts) carry a copyright header, so adding one only to start_worker.sh would be inconsistent. Keeping this file consistent with the directory; a header sweep over these scripts belongs in a separate cleanup.
There was a problem hiding this comment.
@brnguyen2 Understood. I will not require a one-file header change in this PR.
The repository guideline still identifies the missing header, but a consistent update should cover the launcher-script set in a separate cleanup. Would you like me to create a follow-up GitHub issue for that header sweep?
✏️ Learnings added
Learnt from: brnguyen2
Repo: NVIDIA/TensorRT-LLM PR: 17334
File: examples/disaggregated/slurm/benchmark/start_worker.sh:90-96
Timestamp: 2026-08-07T03:18:05.235Z
Learning: For `examples/disaggregated/slurm/benchmark` launcher scripts, maintainers prefer a consistent directory-wide NVIDIA copyright-header sweep in a separate cleanup change. Do not request a header addition only in a single modified launcher script when peer launcher scripts also lack headers.
Note: Learnings are effective only in the context of similar code segments. To apply general coding guidelines, please configure review instructions. You can manage existing learnings in the app.
You are interacting with an AI system.
| from kimi_k3_sa_harness import ( # noqa: E402 | ||
| PROMPTS_AND_CHECKS, | ||
| _compare_logits_parity, | ||
| _parity_prompts, | ||
| ) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Resolve the imported symbols in the sibling harness module.
fd -t f 'kimi_k3_sa_harness.py' | while IFS= read -r f; do
ast-grep outline "$f" --items all
rg -n -C4 'PROMPTS_AND_CHECKS|def _parity_prompts|def _compare_logits_parity' "$f"
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 157
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' 'Candidate files:'
git ls-files '*kimi_k3*' | sed -n '1,120p'
printf '%s\n' 'Harness symbol definitions and references:'
rg -n -C8 'PROMPTS_AND_CHECKS|def _parity_prompts|def _compare_logits_parity|_parity_prompts\(|_compare_logits_parity\(' \
tests/integration/defs tests 2>/dev/null | sed -n '1,260p'Repository: NVIDIA/TensorRT-LLM
Length of output: 11170
🏁 Script executed:
#!/bin/bash
set -e
printf '%s\n' 'Definition directory contents:'
find tests/integration/defs -maxdepth 1 -type f -printf '%f\n' | sort | sed -n '1,240p'
printf '%s\n' 'All harness filename matches, including ignored and hidden paths:'
find . -type f -name 'kimi_k3_sa_harness.py' -print
printf '%s\n' 'References to the harness module:'
rg -n -C3 'kimi_k3_sa_harness|PROMPTS_AND_CHECKS|_parity_prompts|_compare_logits_parity' . \
-g '*.py' -g '!build/**' -g '!dist/**' -g '!node_modules/**' 2>/dev/null | sed -n '1,320p'Repository: NVIDIA/TensorRT-LLM
Length of output: 6986
Add tests/integration/defs/kimi_k3_sa_harness.py or update the import. The module is absent from this repository, so kimi_k3_disagg_parity.py cannot import its required symbols.
🤖 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/integration/defs/kimi_k3_disagg_parity.py` around lines 89 - 93,
Resolve the missing kimi_k3_sa_harness dependency used by
kimi_k3_disagg_parity.py: either add that module with PROMPTS_AND_CHECKS,
_compare_logits_parity, and _parity_prompts, or update the import to reference
their actual existing definitions. Ensure the parity test imports successfully
without changing its required symbols.
There was a problem hiding this comment.
kimi_k3_sa_harness.py is added by PR #17327 (Kimi K3 serving parsers and speculative decoding); this parity script is a manually invoked harness, not a pytest-collected test, so the import is only exercised when the harness is run. It resolves once #17327 lands; we are keeping the merge ordering of the K3 PR series in mind.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
repo="NVIDIA/TensorRT-LLM"
echo "Pull request metadata:"
gh pr view 17327 --repo "$repo" --json number,state,mergedAt,baseRefName,headRefName,url \
--jq '{number, state, mergedAt, baseRefName, headRefName, url}'
echo
echo "Matching changed files:"
gh pr view 17327 --repo "$repo" --json files \
--jq '.files[] | select(.path == "tests/integration/defs/kimi_k3_sa_harness.py") | {path, additions, deletions}'Length of output: 890
@brnguyen2, confirmed. #17327 is open and adds tests/integration/defs/kimi_k3_sa_harness.py.
The import works after #17327 merges and this PR includes that commit. Until then, an independent checkout of this PR cannot run the manual harness. Please preserve that merge order or rebase this PR after #17327 lands. I will keep the finding open because the dependency is not yet present in the current branch.
Test coverage summary: tests/integration/defs/kimi_k3_disagg_parity.py is a manually invoked harness. It is not listed in a CI or QA test list. Coverage needs follow-up after the K3 PR series is merged.
You are interacting with an AI system.
|
PR_Github #64435 [ run ] triggered by Bot. Commit: |
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - CONCERNS
Verdict: The code paths shown look internally consistent, but this feature ships with no CI-runnable coverage and two author-acknowledged unresolved correctness follow-ups, so it should not merge to main as-is (mergeable_state is also blocked).
Concerns
-
[MAJOR]
tests/integration/defs/kimi_k3_disagg_parity.py:1(and missing test-list wiring) - feature ships with no CI coverage- What is wrong: Substantial production code is added (bounce byte-gate in
bounce/config.py/impl.py,MambaPolicy.validate_peer_compatible/payload_bytes, transceiver hybrid routing,_util.pydisagg enablement) plus new unit tests and a parity harness, but per the description there are no changes undertests/integration/test_lists/,test-db/, orqa/, and the parity harness is not registered. - How it fails: None of the new tests run in CI, so a future regression in the recurrent-state transfer path (extra_bytes sizing overrun, peer-validation false-reject) lands undetected.
- Suggested fix: Register the applicable unit tests in an L0 list and wire (or explicitly document the gating job for) the integration harness before un-drafting.
- What is wrong: Substantial production code is added (bounce byte-gate in
-
[MAJOR]
tensorrt_llm/_torch/pyexecutor/_util.py:163- self-flagged unresolved correctness items- What is wrong: The PR's own Dev/QA review states the generation-side replay-cache seeding function is not invoked, and that peer validation may reject valid hybrid models when ranks hold different pipeline-stage layer subsets. The replay-cache seeding code is not in this diff.
- How it fails: A gen request whose recurrent state depends on the seeded replay cache would decode from uninitialized state after the ctx->gen handoff, producing wrong tokens with no error.
- Suggested fix: Confirm (with a test) that replay-cache seeding is invoked on the gen path, and resolve the PP peer-validation follow-up before approval.
Minor notes (non-blocking)
tests/unittest/disaggregated/test_kda_mamba_transfer.py:575-test_kda_transferruns a real NIXL loopback but has no GPU/NIXL marker; CPU-only collection will fail instead of skipping.tensorrt_llm/_torch/pyexecutor/_util.py:170- theand not is_disaggchange makes the kimi_linear branch fall through to unseen manager-selection code; add a test asserting the selected manager class for kimi_linear+disagg (theinfo_oncelog is not a check).
QA view
- Test coverage: partial - unit tests exist for the byte-gate, hybrid bounce reserve, aux zero-draft round-trip, and KDA peer validation, but none are registered in a test list so CI does not run them; the gen-side replay-cache path and the
_util.pydisagg fall-through have no test here. - SM coverage: code touches Blackwell (GB200/GB300, sm100) fabric-VMM bounce and the NIXL transfer path; tests run only on CPU (logic) and single-node loopback (unmarked). No CI run on the target arch within this PR - a real coverage gap.
- Test code:
test_kda_transfermissing GPU/NIXL marker; nothing registered intest_lists//test-db//qa/; parity harness is a standalone multi-node script (self-test only). - Test time: unknown - tests are not wired into CI, so no measurable impact now; wiring the loopback/Blackwell runs later would add non-trivial time.
- Needs
/qa-verify: yes - Blackwell-specific transfer path with no CI test on that arch, tests not list-wired, and unresolved replay-cache/PP follow-ups; a human should re-run token-parity/GSM8K on GB200/GB300 disagg before trusting this.
Possible new issues
extra_bytesregion sizing intransfer.pydispatch_task assumes sendercollect_fragssizes are slot-independent; a future slot-dependent sender would overrun the bounce slot into the neighbor. Guarded only by the K3-geometry unit test._check_peer_compatiblenow raisesValueErrorinstead of returningbool; on the sender listener thread exceptions are only logged, so an incompatible peer fails silently sender-side and relies on the receiver-side gate - a fragile mixed contract.
What I could not verify
- The manager-selection fall-through in
get_kv_cache_manager_clsbeyond the shown lines (which manager is actually chosen for kimi_linear+disagg). - The generation-side replay-cache seeding code (not in this diff) - cannot confirm it is invoked.
- Because the PR is stacked on #17269, the shown diff may include unrelated content that will shrink after rebase; runtime behavior on hardware is unverified.
Automated review by NVCortex Lite, run by @fredricz-20070104.
fredricz-20070104
left a comment
fredricz-20070104
left a comment
There was a problem hiding this comment.
Review summary - Approve (non-blocking)
Approving so this is not blocked on me. The points raised in my review comment above are non-blocking — please read them and address what you agree with before merging.
Worth doing before this is relied on: Arch-specific (Blackwell) transfer path with no CI test on that arch, tests not registered in any list, and author-acknowledged unresolved follow-ups (gen-side replay-cache seeding not invoked, PP peer validation). A human QA should re-run token-parity/GSM8K on GB200/GB300 disagg before this is trusted.
Automated review by NVCortex Lite, run by @fredricz-20070104.
…rness error handling, launcher quote stripping - test_kda_mamba_transfer.py: shut down every GPU-backed manager created by _create_kda_managers (not just the ranks under test) and wrap run_kda_transfer_test in try/finally so managers and transceivers are released on early assertion failures. - kimi_k3_disagg_parity.py: _served_model now returns None on connection-level failures (URLError/OSError) and malformed responses, matching its documented contract; fix the mismatched-model NOTE to state the actual behavior (one model name sent to both endpoints). - start_server.sh / start_worker.sh: strip literal single quotes from TRTLLM_PATH_PREPEND / TRTLLM_PYTHONPATH_PREPEND before use, matching run_benchmark.sh. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
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/integration/defs/kimi_k3_disagg_parity.py`:
- Around line 129-131: The _served_model() response handling must validate
malformed /v1/models payloads before indexing and catch json.JSONDecodeError
alongside existing connection failures, returning None consistently for invalid
JSON, missing or null data, and null entries. Add regression tests covering
these malformed payloads and connection failures, and register the tests in the
appropriate test-db/ or qa/ test list.
🪄 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: 89bcf8c7-5597-4d03-b61f-cef603462277
📒 Files selected for processing (4)
examples/disaggregated/slurm/benchmark/start_server.shexamples/disaggregated/slurm/benchmark/start_worker.shtests/integration/defs/kimi_k3_disagg_parity.pytests/unittest/disaggregated/test_kda_mamba_transfer.py
🚧 Files skipped from review as they are similar to previous changes (3)
- examples/disaggregated/slurm/benchmark/start_server.sh
- examples/disaggregated/slurm/benchmark/start_worker.sh
- tests/unittest/disaggregated/test_kda_mamba_transfer.py
| except (urllib.error.URLError, OSError, KeyError, IndexError) as e: | ||
| print(f"[parity] NOTE: {base_url}/v1/models unavailable ({e})") | ||
| return None |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tests/integration/defs/kimi_k3_disagg_parity.py"
ast-grep outline "$file" --items all --type function
rg -n -C6 'def _http_json|def _served_model|JSONDecodeError|models\["data"\]|\["id"\]' "$file"
rg -n -C4 'kimi_k3|parity|self_test' tests/integration/test_lists 2>/dev/null || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 8755
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tests/integration/defs/kimi_k3_disagg_parity.py"
printf '%s\n' '--- source around _http_json and _served_model ---'
sed -n '90,140p' "$file"
printf '%s\n' '--- source around self-test and entry point ---'
sed -n '600,700p' "$file"
printf '%s\n' '--- changed files ---'
git status --short
git diff --stat -- "$file" tests/integration/test_lists
git diff -- "$file" tests/integration/test_lists | sed -n '1,260p'
printf '%s\n' '--- exact test-list references ---'
rg -n -i -C3 'kimi[_-]?k3|disagg[_-]?parity|kimi' tests/integration/test_lists/test-db tests/integration/test_lists/qa 2>/dev/null || true
printf '%s\n' '--- all references to the definition and self-test ---'
rg -n -C3 'kimi_k3_disagg_parity|_served_model|_self_test' tests tests/integration/test_lists 2>/dev/null || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 42095
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
import json
payloads = [
[],
{"data": None},
{"data": [None]},
{"data": [{}]},
{"data": [{"id": "model"}]},
]
print("--- _served_model extraction behavior ---")
for payload in payloads:
try:
result = payload["data"][0]["id"]
except Exception as exc:
print(f"{payload!r} -> {type(exc).__name__}: {exc}")
else:
print(f"{payload!r} -> {result!r}")
print("--- _http_json decode behavior ---")
try:
json.loads(b"not-json")
except Exception as exc:
print(f"invalid JSON -> {type(exc).__name__}")
print(f"JSONDecodeError is ValueError: {isinstance(exc, ValueError)}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 578
Handle malformed /v1/models responses before indexing.
_served_model() allows TypeError to escape for malformed payloads such as [], {"data": null}, and {"data": [null]}. Invalid JSON raises uncaught json.JSONDecodeError. Validate the response shape before indexing and handle json.JSONDecodeError.
Add regression coverage for malformed payloads and connection failures. The current self-test does not cover _served_model(), and no matching test-db/ or qa/ test-list entry 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 `@tests/integration/defs/kimi_k3_disagg_parity.py` around lines 129 - 131, The
_served_model() response handling must validate malformed /v1/models payloads
before indexing and catch json.JSONDecodeError alongside existing connection
failures, returning None consistently for invalid JSON, missing or null data,
and null entries. Add regression tests covering these malformed payloads and
connection failures, and register the tests in the appropriate test-db/ or qa/
test list.
… synthetic rank info Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
|
PR_Github #64435 [ run ] completed with state
|
| # skip bounce below this many bytes (the operative gate; see DEFAULT_MIN_BYTES for the rationale) | ||
| min_bytes: int = DEFAULT_MIN_BYTES | ||
| # legacy block-count gate, kept for back-compat (TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS): both gates | ||
| # must pass, and the default of 1 makes this one vacuous so the byte gate decides | ||
| min_blocks: int = 1 |
There was a problem hiding this comment.
This flips the bounce gate for every existing bounce user, not just Kimi K3: min_blocks goes 96 -> 1 (vacuous) and the operative gate becomes 2 MiB. On a production-sized model a single block is usually already well past 2 MiB, so transfers that take the per-block path today will start going through the bounce arena. That is a behaviour change to an already-shipped, separately-opted-into feature (kv_cache_bounce_size_mb > 0) and it is not mentioned in the title or the description.
The rationale in the comment above is convincing and I am not arguing the direction — the 96-block number was clearly calibrated for 128-token blocks. Please call the gate change out in the PR description, and ideally in the kv_cache_bounce_size_mb field doc in llm_args.py, so operators who sized their arena around the old 96-block gate are not surprised by the new arena pressure.
There was a problem hiding this comment.
Good catch, and agreed this should not silently change behavior for existing bounce users. Rather than only documenting it, 90dbf68 guards the change: plain-KV transfers keep the original 96-block gate (TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS, default 96 restored), and the byte gate (TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES, 2 MiB) now applies only when the payload carries recurrent (mamba/KDA) state, i.e. extra_bytes > 0 at the reserve() gate. Existing deployments that opted into kv_cache_bounce_size_mb see no change in which transfers use the arena, and Kimi K3 is unaffected since every K3 request carries the fixed ~433 MiB KDA payload, which clears any sane byte gate. There is an inline TODO at the gate to investigate whether the byte-only gate is safe (or better) for plain-KV payloads too, so the special case can be removed; a follow-up ticket will be filed for that.
| # The byte gate below which a transfer keeps the per-block path is lowered to 1 via the | ||
| # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES env (set by the test) so the ordinary short test prompts still | ||
| # take the coalesced-bounce WRITE path (the production default of 2 MiB may exceed a short prompt's | ||
| # KV footprint on a tiny model). |
There was a problem hiding this comment.
Heads-up on a cross-PR collision: #17392 ([TRTLLM-15078][test] Prune non-Llama-3.1-8B Llama tests) deletes this file, while this PR edits it. Both are open against main, so whichever lands second either conflicts or silently drops the other's intent. Worth syncing with @xinhe-nv before either merges.
There was a problem hiding this comment.
Thanks for the heads-up. I checked #17392: it deletes this config file (and prunes the associated test) as part of the TinyLlama/Llama test cleanup. I will watch the merge order; if #17392 lands first, I will rebase and move this config's content to wherever that reorganization expects it (or drop the edit if the bounce integration test is relocated). If this PR lands first, the conflict will show up in #17392 and can be resolved there with @xinhe-nv.
…or recurrent-state payloads Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
…tent MambaPolicy.validate_peer_compatible now checks the global (per-rank bytes x mamba_tp) recurrent-state size, so the fixed-size synthetic mamba group in make_page_table() reads as a replicated state under heterogeneous TP and fails registration in the tp2-vs-tp1 registrar tests. Shard the fixture's mamba pools from a fixed global size by a mamba_tp parameter (default 2, matching make_rankinfo's default tp_size and preserving the previous byte values) and pass mamba_tp=1 for the tp=1 peers. Also restores the intended failure mode of test_peer_registrar_rejects_misaligned_subbyte_head_mismatch, which had been passing on the mamba mismatch instead of the alignment check. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
BowenFu
left a comment
BowenFu
left a comment
There was a problem hiding this comment.
Bounce-gate finding addressed: plain-KV payloads keep the original 96-block gate (min_blocks default restored to 96) and the 2 MiB byte gate now applies only when extra_bytes > 0, so an existing bounce deployment is unchanged. Also checked the narrower case Codex raised -- a non-Kimi hybrid model with bounce explicitly enabled -- and it cannot regress either: on main impl.py:230 bails any transfer whose layer group has no known slot size ("e.g. mamba"), so hybrids never reach the arena today; this PR is what first makes them eligible.
Approving over the 3 open CodeRabbit threads: the start_worker.sh copyright nit (whole directory is header-less), the kimi_k3_disagg_parity.py import of a module from #17327, and the _http_json error-handling nit -- the latter two are in a manually-invoked harness that pytest does not collect. The #17392 file collision still needs merge-order care.
6 participants
Description
Adds disaggregated-serving support for Kimi K3 (KimiLinear), in four commits:
MambaPolicy state-region mapping for the KDA mixer, peer registration,
bounce-buffer config/impl updates sized for the KDA state payload, and
transceiver handling for hybrid (attention + recurrent-state) models.
pyexecutor/_util.pypreviouslyraised NotImplementedError for disaggregated serving (referencing this
ticket); it now routes through the shared hybrid transceiver validation.
updates, plus a Kimi K3 disagg logits-parity integration harness.
Stacked on #17269 — draft until that merges; the diff then shrinks to
these four commits after a rebase onto main.
Notes
the SA disagg test depend on [TRTLLM-14814][feat] Kimi K3 serving parsers, chat template, and speculative decoding (suffix automaton + DFlash scaffold) #17327 and follow once both PRs are in.
wired into any test list here; unit-suite results on Blackwell hardware
will be posted before un-drafting. Prior validation of this code on the
feature bring-up branch: token-level parity between disaggregated and
aggregated serving on the target model.
Test Coverage
PR Checklist
[TRTLLM-14815][feat]conventionDev Engineer Review
trtllm-servefallback handling.#17327.test-db/, orqa/changes are included.QA Engineer Review
kimi_k3_disagg_parity.pyparity harness with endpoint, token, logprob, GSM8K, reporting, CLI, and self-test coverage.tests/integration/test_lists/,test-db/, orqa/.