[TRTLLM-14704][feat] Support multi-modal part of K3 - #17050
Conversation
|
In order to test it, I renamed |
brnguyen2
left a comment
brnguyen2
left a comment
There was a problem hiding this comment.
Approving — the comments below are optional touch-ups, not blockers.
Solid bring-up — the delta table in the module docstring and the head_dim / replication comments make this much easier to follow than the usual VLM port, and the MMMU channel extractor is a real scoring bug worth fixing. Given this targets feat/kimi_k3 and the branch is racing to functional completeness, I've split the items below into before merge (cheap, seconds each) and follow-up PR (everything else — none of it should hold this up).
Before merge (cheap)
Ticket. Title is [None][feat] on a ~1100-line model bring-up. This should carry TRTLLM-14704 per the repo convention — [TRTLLM-14704][feat] Support multi-modal part of K3. [None] is for chores/docs/waives.
Stale path reference. run_gsm8k_kimi_k3.sbatch → run_eval_kimi_k3.sbatch is a user-facing rename; docs are updated, but tests/integration/defs/kimi_k3_disagg_parity.py:42 still points at the old name. One-line fix.
Description vs. diff. Two things a future reader (and the merge-back reviewer) will need, and they're just description lines:
- The
modeling_kimi_k25.pychanges aren't wire-up — they change behavior for the shipped K2.5 model:_vision_requires_replicationnow silently replicates the vision tower for anytp_sizethat doesn't divide the head count, andload_weightsnow materializes every checkpoint slice up front. Both look right, but please say so explicitly and note whether K2.5 was re-validated. - TEP16 support in the sbatch isn't mentioned at all.
Follow-up PR (file under TRTLLM-14674, not blocking)
Test coverage. tests/unittest/others/test_lm_eval.py isn't referenced by any file under tests/integration/test_lists/, so the tests listed as this PR's coverage don't run in pre-merge CI. The 650 lines of new model + config code have no tests at all — is_kimi_k3_multimodal_config is a pure dict predicate and KimiK3Config.from_dict round-trip is a few lines; both are cheap to pin and both are exactly the kind of routing logic that breaks silently when a checkpoint field is renamed. This should land before the branch merges back, but it doesn't need to be in this PR.
Inheritance structure. KimiK3VisionModel.__init__ and K3MoonViT3dEncoder.__init__ deliberately skip their parents and re-implement them — ~90 lines duplicated from the K2.5 versions, already diverged (K3 adds qkv_hidden_size, drops layer_norm_eps). A shared _parse_vision_cfg(vision_cfg) helper plus a _build_blocks() hook would let both share the config parsing and keep future K2.5 fixes reaching K3. Worth doing while the K2.5/K3 pairing is still fresh in someone's head, but it's a refactor of working code — follow-up.
183b018 to
a97b9c2
Compare
moraxu
changed the title
Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Signed-off-by: Fred Wei <20514172+WeiHaocheng@users.noreply.github.com>
Signed-off-by: Michal Guzek <mguzek@nvidia.com>
- Pass the vision tower head count from both KimiK*VisionModel inits and the projector MLPs into _vision_requires_replication / _get_vision_tp_mapping instead of re-deriving it there with a hardcoded K3 default (a K2.5 config omitting both head-count keys could silently force a shardable 16-head tower to tp=1). - Document the runtime-dtype-dependent eps of the reference-matching default-eps RMSNorms in the K3 vision tower (norm0/norm1 and final_layernorm). - Clean up the per-job TEP runtime yaml: remove it via trap on batch script exit and gitignore the .eval_tep_runtime/.eval_dflash_runtime patterns. - Unit-test is_kimi_k3_multimodal_config (composite, language_model_only opt-out, missing/empty/non-dict sub-configs, text-only kimi_linear, wrong model_type) so released-config field renames fail loudly. Signed-off-by: Michal Guzek <mguzek@nvidia.com>
- Align the K2.5-vs-K3 delta table in the module docstring. - Stop propagating the text model's kv_cache_quant_algo into the vision tower's QuantConfig: the vision encoder layers keep no KV cache, so the inherited value could only steer attention kernel selection for cache-less attention (see NVIDIA#12851 for a past instance). - Replace the hand-rolled torch_dtype/dtype/string normalization with the shared resolve_hf_torch_dtype helper and document the write-back contract that makes model_config.torch_dtype safe everywhere after __init__. - Document why the wrapper declares _supports_sdpa and why a second __init__ on a built instance must early-return. Signed-off-by: Michal Guzek <mguzek@nvidia.com>
- run_eval_kimi_k3.sbatch: collect every per-job runtime config (the --dflash rewrite included) in a CLEANUP_FILES array removed by a single EXIT trap, since a second bare 'trap ... EXIT' would replace the first rather than add to it. - Deduplicate the VLM load_weights override: the K2.5 wrapper now builds its tower through a _VISION_MODEL_CLS class attribute, Kimi K3 overrides only that attribute and inherits load_weights (and the MetaInitMode deferral/recreation logic) unchanged. - extract_kimi_k3_mmmu_answer: guard on the extraction succeeding rather than the channel span being non-empty, so channel content that reduces to nothing (e.g. markdown-bold whitespace after the cascade's bold-stripping) keeps falling back instead of returning ''; add a regression test. - Unit-test _vision_requires_replication (12 heads/tp16 -> replicate, 16 heads/tp8 -> shard, attention-DP -> always replicate). - Fix a D205 docstring in the K3 lm_eval tests flagged by main's ruff coverage. Signed-off-by: Michal Guzek <mguzek@nvidia.com>
- KimiK3VisionConfig: route the flash_attention_2 default through the attn_implementation kwarg — transformers v5 PretrainedConfig.__init__ overwrites a directly-assigned _attn_implementation with the kwarg's None default; explicit caller overrides still win. - KimiK3VisionModel: validate the vision-config switches the tower hardcodes (norm_type/mlp_type/activation_func/pos_emb_type and the three bias flags) so an unsupported checkpoint variant fails loudly at construction instead of silently building a different architecture. - run_eval_kimi_k3.sbatch: escape sed replacement metacharacters in the --dflash drafter path so paths containing \, & or | substitute literally into the per-job YAML. - run_dspark_acceptance.sbatch: validate that value-taking options have a value before reading $2 (set -u aborted with an unbound-variable error instead of a usage message). - README: note that .venv-3.12 should be substituted with .venv-<major>.<minor> when the container ships a different Python. Signed-off-by: Michal Guzek <mguzek@nvidia.com>
03e58c9 to
f3e5868
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. |
|
@CodeRabbit fullreview |
|
✅ Action performedFull review finished. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (4)
tensorrt_llm/_torch/configs/kimi_k3.py (1)
34-68: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winComplete the new Python type annotations.
Add
-> Noneto the configuration constructors. Use parameterized Python 3.10 container types for configuration inputs andmerge_kernel_size. AnnotatePatchMergerMLPV2.forwardwith its tensor-or-sequence input and return types.
tensorrt_llm/_torch/configs/kimi_k3.py#L34-L68: add a constructor return type and replace the baretupleannotation with a precise built-in generic.tensorrt_llm/_torch/configs/kimi_k3.py#L121-L129: add a constructor return type and replace bare nesteddictannotations with precise types.tensorrt_llm/_torch/pyexecutor/config_utils.py#L516-L516: replace the baredictparameter type with a parameterized configuration mapping type.tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py#L265-L272: annotatexand the tensor-or-list return value.As per coding guidelines, “Annotate every function” and “prefer built-in generic types and
|.”🤖 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_k3.py` around lines 34 - 68, Complete the annotations across all affected sites: in tensorrt_llm/_torch/configs/kimi_k3.py:34-68, add -> None to the configuration constructor and change merge_kernel_size from bare tuple to a precise built-in generic; in tensorrt_llm/_torch/configs/kimi_k3.py:121-129, add -> None to the other constructor and parameterize its nested dict annotations; in tensorrt_llm/_torch/pyexecutor/config_utils.py:516, replace the bare dict parameter with a parameterized configuration mapping type; and in tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py:265-272, annotate PatchMergerMLPV2.forward so x and its tensor-or-list return value use precise types, preferring Python 3.10 built-in generics and | unions.Source: Coding guidelines
tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py (1)
29-100: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winTest the loader result in addition to the predicate.
These tests validate
is_kimi_k3_multimodal_config, but they do not validate the newload_pretrained_configbranch. Add tests that mockPretrainedConfig.get_config_dictand verify that a composite checkpoint createsKimiK3Config, rebuilds both typed sub-configurations, and setsarchitecturestoKimiK3ForConditionalGeneration. Add alanguage_model_onlycase that verifies flattening toKimiLinearConfig.Without these tests, branch ordering or
from_dict()regressions fail only when loading a checkpoint. As per path instructions, “Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM.”🤖 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_k3_config_routing.py` around lines 29 - 100, Extend the loader tests around load_pretrained_config, mocking PretrainedConfig.get_config_dict to cover composite checkpoints and language_model_only configurations. Verify composite loading returns KimiK3Config with typed vision_config and text_config sub-configurations and architectures set to KimiK3ForConditionalGeneration; verify the language_model_only path flattens to KimiLinearConfig.Source: Path instructions
tests/unittest/others/test_lm_eval.py (2)
1654-1774: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd return annotations to the new test functions.
Add
-> Noneto each addedtest_k3_*function. The coding guidelines require annotations on every function.As per coding guidelines: “Annotate every function.”
🤖 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/others/test_lm_eval.py` around lines 1654 - 1774, Add a -> None return annotation to every newly added test_k3_* function in this diff, including test_k3_channel_bare_letter, test_k3_channel_real_samples_recovered, test_k3_channel_parenthesized_answer, test_k3_channel_answer_is_phrase, test_k3_truncated_after_channel_open, test_k3_last_channel_wins, test_k3_no_channel_bare_letter_falls_back, test_k3_no_channel_truncated_thinking_does_not_crash, test_k3_empty_input, test_k3_scrubs_residual_special_tokens, test_k2_5_strip_thinking_path_unchanged, and test_k3_channel_extracting_to_nothing_falls_back_to_cascade.Source: Coding guidelines
1654-1774: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the 12 added test functions with
-> None.Coverage is sufficient for channel extraction, truncation, repeated channels, token cleanup, fallback behavior, and K2.5 compatibility.
tests/integration/test_lists/test-db/l0_cpu.ymlincludesunittest/others, which coverstests/unittest/others/test_lm_eval.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/others/test_lm_eval.py` around lines 1654 - 1774, Annotate all 12 newly added test functions in test_lm_eval.py with an explicit return type of None, including test_k3_channel_bare_letter and the other test_k3_* and K2.5 compatibility tests in this diff.Sources: Coding guidelines, Path instructions
🤖 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/kimi_k3/run_dspark_acceptance.sbatch`:
- Around line 122-152: Update run_leg to preserve argument boundaries when
forwarding its shifted arguments into the nested bash -c command. Encode each
argument shell-safely, such as with Bash printf -v %q, before constructing the
command, then use the encoded argument list instead of unquoted $* so
whitespace-containing values like --drafter paths remain single arguments.
In `@examples/kimi_k3/run_eval_kimi_k3.sbatch`:
- Around line 231-234: Update the mmmu TASK_CMD output_path value to be a
directory name rather than a .json file path, removing the .json suffix while
preserving the existing job ID-based naming and LmEvalEvaluator.save_results
behavior.
- Around line 317-333: Make the virtual-environment path configurable across the
Kimi K3 scripts: in examples/kimi_k3/run_eval_kimi_k3.sbatch lines 317-333,
define the shared environment-path setting and use it for PATH and trtllm-eval;
in examples/kimi_k3/run_dspark_acceptance.sbatch lines 143-148, use the same
setting for PATH; and in examples/kimi_k3/README.md lines 24-28, document the
default .venv-3.12 commands and instruct users to set the configuration to the
matching .venv-<major>.<minor> path for other container Python versions.
In `@tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py`:
- Around line 1-104: Extend TestIsKimiK3MultimodalConfig with loader-level tests
that invoke load_pretrained_config on representative Kimi K3 configs and
validate KimiK3Config.from_dict behavior. Assert multimodal configs retain the
composite KimiK3 configuration and assign KimiK3ForConditionalGeneration, while
text-only or invalid subconfigs follow the intended text-only route. Keep the
existing predicate and replication tests, adding coverage for the actual loading
path rather than testing only helper functions.
---
Nitpick comments:
In `@tensorrt_llm/_torch/configs/kimi_k3.py`:
- Around line 34-68: Complete the annotations across all affected sites: in
tensorrt_llm/_torch/configs/kimi_k3.py:34-68, add -> None to the configuration
constructor and change merge_kernel_size from bare tuple to a precise built-in
generic; in tensorrt_llm/_torch/configs/kimi_k3.py:121-129, add -> None to the
other constructor and parameterize its nested dict annotations; in
tensorrt_llm/_torch/pyexecutor/config_utils.py:516, replace the bare dict
parameter with a parameterized configuration mapping type; and in
tensorrt_llm/_torch/models/modeling_kimi_k3_vl.py:265-272, annotate
PatchMergerMLPV2.forward so x and its tensor-or-list return value use precise
types, preferring Python 3.10 built-in generics and | unions.
In `@tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py`:
- Around line 29-100: Extend the loader tests around load_pretrained_config,
mocking PretrainedConfig.get_config_dict to cover composite checkpoints and
language_model_only configurations. Verify composite loading returns
KimiK3Config with typed vision_config and text_config sub-configurations and
architectures set to KimiK3ForConditionalGeneration; verify the
language_model_only path flattens to KimiLinearConfig.
In `@tests/unittest/others/test_lm_eval.py`:
- Around line 1654-1774: Add a -> None return annotation to every newly added
test_k3_* function in this diff, including test_k3_channel_bare_letter,
test_k3_channel_real_samples_recovered, test_k3_channel_parenthesized_answer,
test_k3_channel_answer_is_phrase, test_k3_truncated_after_channel_open,
test_k3_last_channel_wins, test_k3_no_channel_bare_letter_falls_back,
test_k3_no_channel_truncated_thinking_does_not_crash, test_k3_empty_input,
test_k3_scrubs_residual_special_tokens, test_k2_5_strip_thinking_path_unchanged,
and test_k3_channel_extracting_to_nothing_falls_back_to_cascade.
- Around line 1654-1774: Annotate all 12 newly added test functions in
test_lm_eval.py with an explicit return type of None, including
test_k3_channel_bare_letter and the other test_k3_* and K2.5 compatibility tests
in this diff.
🪄 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: cca152b6-b4af-48a4-bbc8-468c82166e6d
📒 Files selected for processing (19)
.gitignoredocs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.mdexamples/kimi_k3/README.mdexamples/kimi_k3/perf_sweep/acc_sweep.sbatchexamples/kimi_k3/run_dspark_acceptance.sbatchexamples/kimi_k3/run_eval_kimi_k3.sbatchexamples/kimi_k3/run_gsm8k_kimi_k3.sbatchtensorrt_llm/_torch/configs/__init__.pytensorrt_llm/_torch/configs/kimi_k3.pytensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/_arch_index.pytensorrt_llm/_torch/models/modeling_kimi_k25.pytensorrt_llm/_torch/models/modeling_kimi_k3_vl.pytensorrt_llm/_torch/models/modeling_kimi_linear.pytensorrt_llm/_torch/pyexecutor/config_utils.pytensorrt_llm/evaluate/lm_eval.pytensorrt_llm/evaluate/post_processing.pytests/unittest/_torch/modeling/test_kimi_k3_config_routing.pytests/unittest/others/test_lm_eval.py
💤 Files with no reviewable changes (1)
- examples/kimi_k3/run_gsm8k_kimi_k3.sbatch
- run_dspark_acceptance.sbatch: %q-quote the arguments run_leg forwards into the nested bash -c string, so values containing whitespace (e.g. a drafter path) keep their word boundaries. - run_eval_kimi_k3.sbatch: pass --output_path without the .json suffix — the evaluator treats it as a directory, so the old value created a directory literally named *.json. - Make the job virtual-environment path configurable via TRTLLM_VENV in both Slurm scripts (defaulting to the repository-root .venv-3.12) and document the export in the README, so the documented .venv-<major>.<minor> substitution actually works for the jobs. - Add loader-level Kimi K3 routing tests: config.json -> load_pretrained_config for the composite VLM (KimiK3Config + KimiK3ForConditionalGeneration architectures), the language_model_only opt-out, and the text-only kimi_linear flatten. Signed-off-by: Michal Guzek <mguzek@nvidia.com>
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py (1)
109-114: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd annotations to the new functions.
Annotate
_load_config_from_dictwith a precise input type andPretrainedConfigreturn type. Add-> Noneto each test method. Run the configured type checker andpytest tests/unittest/after the change.As per coding guidelines: “Annotate every function, use
Nonefor procedures, ... use precise ... types.”Also applies to: 125-144
🤖 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_k3_config_routing.py` around lines 109 - 114, Annotate `_load_config_from_dict` with a precise configuration-dictionary input type and `PretrainedConfig` return type, and add `-> None` to every affected test method in the surrounding range. Run the configured type checker and `pytest tests/unittest/` to verify the annotations.Source: Coding guidelines
🤖 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.
Nitpick comments:
In `@tests/unittest/_torch/modeling/test_kimi_k3_config_routing.py`:
- Around line 109-114: Annotate `_load_config_from_dict` with a precise
configuration-dictionary input type and `PretrainedConfig` return type, and add
`-> None` to every affected test method in the surrounding range. Run the
configured type checker and `pytest tests/unittest/` to verify the annotations.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 7402af67-bf3e-4fe0-a275-3c114a13ca66
📒 Files selected for processing (4)
examples/kimi_k3/README.mdexamples/kimi_k3/run_dspark_acceptance.sbatchexamples/kimi_k3/run_eval_kimi_k3.sbatchtests/unittest/_torch/modeling/test_kimi_k3_config_routing.py
🚧 Files skipped from review as they are similar to previous changes (3)
- examples/kimi_k3/run_dspark_acceptance.sbatch
- examples/kimi_k3/run_eval_kimi_k3.sbatch
- examples/kimi_k3/README.md
Add the precise input/return annotations CodeRabbit's trivial nitpick requested for _load_config_from_dict and -> None on the new loader test methods, per the coding guidelines. Signed-off-by: Michal Guzek <mguzek@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #65420 [ run ] triggered by Bot. Commit: |
|
PR_Github #65420 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #65454 [ run ] triggered by Bot. Commit: |
|
PR_Github #65454 [ run ] completed with state
|
|
/bot run |
|
PR_Github #65644 [ run ] triggered by Bot. Commit: |
|
PR_Github #65644 [ run ] completed with state |
|
/bot run |
|
PR_Github #65715 [ run ] triggered by Bot. Commit: |
|
PR_Github #65715 [ run ] completed with state |
asfiyab-nvidia
left a comment
asfiyab-nvidia
left a comment
There was a problem hiding this comment.
LGTM. Approval based on runtime module review
9 participants
Dev Engineer Review
QA Engineer Review
Added test functions in:
tests/unittest/_torch/modeling/test_kimi_k3_config_routing.pytests/unittest/others/test_lm_eval.pyNo corresponding entries in
tests/integration/test_lists/test-db/ortests/integration/test_lists/are identified. Test-list coverage requires follow-up.Verdict: needs follow-up
Description
Adds the multi-modal (vision-language) portion of Kimi-K3 to the PyTorch backend.
tensorrt_llm/_torch/configs/kimi_k3.py— K3 multimodal configtensorrt_llm/_torch/models/modeling_kimi_k3_vl.py— K3 VL modelconfig_utils.py, model/config__init__.py,modeling_kimi_k25.py,modeling_kimi_linear.pyevaluate/lm_eval.py+evaluate/post_processing.pytests/unittest/others/test_lm_eval.pyTest Coverage
tests/unittest/others/test_lm_eval.pyMMMU-val accuracy (GB300, DEP16, full 900 samples)
trtllm-eval ... mmmu --post_process_fn kimi_k3_mmmuon 4 nodes x 4 GPUs (TP16 + attention-DP + EP16), bf16 K3 VL checkpoint,--max_output_length 16384 --preserve_caller_max_tokens,--max_seq_len 24576. Kimi K2.5 NVFP4 reference inreferences/mmmu.yamlis 81.56.feat/kimi_k3basefeat/kimi_k3tip (+ review fixes)mainbase, head03e58c9f14All three agree within their error bars. TEP16 (
--parallel tep) was additionally validated at 85.44 ± 1.15, and text-only GSM8K at 96.66 ± 0.49 (vs the 96.47 ± 0.5 deployment-guide expectation).Note: the
main-base run pinskv_cache_config.use_kv_cache_manager_v2: false— with main's KVCM-V2 default for Kimi, requests beyond ~8k tokens currently fail with"User-provided base page indices is too short"(kv_cache_manager_v2/kvCache.cpp:2274); reported separately.