[TRTLLM-15215][fix] Kimi K3: make the FP8 weight-read master switch opt-in - #17446
Merged
Merged
Conversation
…pt-in KIMI_K3_FP8_WEIGHT_READ defaulted to "1", so any Kimi K3 run on sm_100-family hardware served the replicated MoE-MLP, KDA q/k/v/g/o and MLA projections from an FP8 copy of their weights. The FP8 read is lossy relative to BF16 and there was no opt-in: a stock run silently traded accuracy for decode bandwidth, and the published accuracy numbers are measured with it off. Default the master switch to "0". The KDA and KDA-glue sub-switches only ever narrow an already-enabled master, so they stay default-on and are inert while the master is off; the SM100 gate is unchanged. Resolution of the three switches moves into _resolve_fp8_weight_read_gates() so the defaults are assertable without loading a model. The new test pins them, including that a sub-switch cannot enable FP8 reads on its own and that the SM gate still wins. Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
WalkthroughThe change centralizes Kimi FP8 weight-read gate resolution. FP8 reads now require explicit opt-in on SM100 hardware. Tests cover default, sub-gate, and non-Blackwell behavior. ChangesFP8 weight-read gates
Estimated code review effort: 2 (Simple) | ~10 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/models/modeling_kimi_linear.py (1)
2587-2587: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winUse
kda_fp8for the KDA conversion condition.Line 2587 resolves
kda_fp8, but line 2624 re-reads_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV. Useif kda_fp8:there. This keepsload_weights()aligned with the centralized resolver and makes resolver tests represent loader behavior.Proposed fix
- if os.environ.get(_KIMI_K3_FP8_WEIGHT_READ_KDA_ENV, "1") != "0": + if kda_fp8:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_kimi_linear.py` at line 2587, Update load_weights() to use the resolved kda_fp8 value from _resolve_fp8_weight_read_gates() for the KDA conversion condition instead of re-reading _KIMI_K3_FP8_WEIGHT_READ_KDA_ENV. Preserve the existing conversion behavior while ensuring it follows the centralized resolver.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/models/modeling_kimi_linear.py`:
- Around line 191-193: Update the fp8_weight_read master-switch evaluation in
the surrounding model configuration logic so KIMI_K3_FP8_WEIGHT_READ enables FP8
only when its value is nonempty and not "0"; preserve the existing SM capability
check and dependent kda_fp8/kda_glue_fp8 behavior. Add a regression test
covering an empty master environment value and asserting FP8 weight reads remain
disabled.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/modeling_kimi_linear.py`:
- Line 2587: Update load_weights() to use the resolved kda_fp8 value from
_resolve_fp8_weight_read_gates() for the KDA conversion condition instead of
re-reading _KIMI_K3_FP8_WEIGHT_READ_KDA_ENV. Preserve the existing conversion
behavior while ensuring it follows the centralized resolver.
🪄 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: 2c109ae9-54cb-4479-befb-92f9da38a911
📒 Files selected for processing (2)
tensorrt_llm/_torch/models/modeling_kimi_linear.pytests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py
Signed-off-by: Brian Nguyen <brnguyen@nvidia.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64865 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64865 [ run ] completed with state |
pranav-nvidia
approved these changes
Aug 10, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
KIMI_K3_FP8_WEIGHT_READdefaulted to"1", so any Kimi K3 run on sm_100-familyhardware served the replicated MoE-MLP, KDA q/k/v/g/o and MLA projections from an FP8
copy of their weights. The FP8 read is lossy relative to BF16 and there was no opt-in:
a stock run silently traded accuracy for decode bandwidth, and the published accuracy
numbers are measured with it off.
Default the master switch to
"0". The KDA and KDA-glue sub-switches only ever narrowan already-enabled master, so they stay default-on and are inert while the master is
off; the SM100 gate is unchanged.
Resolution of the three switches moves into
_resolve_fp8_weight_read_gates()so thedefaults are assertable without loading a model. Previously the expression was inline in
load_weights, which meant the only way to test it was to load a checkpoint — so inpractice it was untested, and it drifted.
Test Coverage
New
tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py(7 tests,CPU-only, patches the SM check): the master switch defaults off; setting it enables it
and the default-on sub-gates; each sub-gate narrows an enabled master; a sub-gate cannot
enable FP8 reads on its own; and the SM gate still wins over an explicit opt-in.
PR Checklist
[JIRA/NVBUG/None][type] SummaryDev Engineer Review
KIMI_K3_FP8_WEIGHT_READfrom default-on to default-off._resolve_fp8_weight_read_gates()to centralize SM100 and environment-variable gate resolution.QA Engineer Review
tests/unittest/_torch/modeling/test_kimi_k3_fp8_weight_read_gates.py.tests/integration/test_lists/coverage entry was identified.