Skip to content

Fix Kimi K2.6 compressed-tensors MoE w2 scale loading - #28905

Closed
bingxche wants to merge 1 commit into
mainfrom
bingxche/fix-kimi-k26-w2-scale
Closed

bingxche wants to merge 1 commit into
mainfrom
bingxche/fix-kimi-k26-w2-scale

Conversation

@bingxche

@bingxche bingxche commented Jun 22, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Treat missing/false/static actorder as non-actorder in compressed-tensors WNA16 MoE.
  • Avoid allocating/loading full TP-width w2_weight_scale for checkpoints whose w2 scales are already per-rank, such as Kimi-K2.6.
  • Keep full w2 scale loading for real actorder/g_idx cases.

Validation

  • PYTHONPYCACHEPREFIX=/tmp/sglang-pycache python3 -m py_compile python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py
  • pre-commit hooks from git commit

CI


CI States

Latest PR Test (Base): ❌ Run #28089355031
Latest PR Test (Extra): ❌ Run #28089354818

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request updates compressed_tensors_wNa16_moe.py to properly handle self.actorder when it is an Enum or a mixed-case string. The reviewer suggests normalizing self.actorder in-place on the instance rather than using a local variable, ensuring that other references to self.actorder within the class also benefit from this normalization and evaluate correctly.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment on lines +133 to +138
actorder = getattr(self.actorder, "value", self.actorder)
if isinstance(actorder, str):
actorder = actorder.lower()
load_full_w2 = (
actorder not in (None, False, "", "static") and self.group_size != -1
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

By normalizing self.actorder in-place rather than using a local variable, we ensure that other references to self.actorder in this class also benefit from the normalization.

Specifically:

  1. On line 145, self.is_k_full = (not self.actorder) ... will evaluate correctly. If self.actorder is an Enum member (e.g., ActOrder.NONE), not self.actorder would incorrectly evaluate to False because Python Enum members are always truthy.
  2. On line 303, if self.actorder == "group": will correctly match the string value instead of failing when compared against an Enum member or a mixed-case string.
Suggested change
actorder = getattr(self.actorder, "value", self.actorder)
if isinstance(actorder, str):
actorder = actorder.lower()
load_full_w2 = (
actorder not in (None, False, "", "static") and self.group_size != -1
)
self.actorder = getattr(self.actorder, "value", self.actorder)
if isinstance(self.actorder, str):
self.actorder = self.actorder.lower()
load_full_w2 = (
self.actorder not in (None, False, "", "static") and self.group_size != -1
)

Experiment: the K2.6 w2 _load_w2 narrow IndexError on AMD is non-deterministic
(different TP rank reported first across days: got 192/320/448, plus a segfault
day), with the load/shard code and the HF checkpoint unchanged. This points to a
race in multithreaded weight loading rather than a deterministic sharding bug.
Disable enable_multithread_load for both kimi-k26 accuracy jobs to test whether
the IndexError disappears.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@bingxche
bingxche force-pushed the bingxche/fix-kimi-k26-w2-scale branch from 463c6d6 to ac6ae22 Compare June 24, 2026 09:36
@github-actions github-actions Bot added the amd label Jun 24, 2026
@bingxche

Copy link
Copy Markdown
Collaborator Author

Experiment: disable enable_multithread_load for Kimi-K2.6 (mi325 + mi355x)

Force-pushed this branch to a single recipe-only commit (ac6ae22) that flips enable_multithread_load to false in both K2.6 accuracy recipes — nothing else changed.

Why: the K2.6 _load_w2 narrow IndexError on AMD looks non-deterministic, not a deterministic sharding bug:

  • Same load/shard code (_load_w2 narrow unchanged since May 20) and same HF checkpoint (last modified 2026-05-19), yet mi325 failed 6 days straight then passed clean on 6/23 (acc 0.942).
  • The failing narrow reports a different TP rank first each daygot 192 / 448 / 192 / 320 / 320 (= 64 × rank for rank 3/7/3/5/5), plus one day it was a segfault in fill_kernel instead. Whoever loses the race raises first → smells like a multithreaded-load race.

This run disables multithreaded loading to see if the IndexError disappears.

Triggered run (mi325 + mi355x kimi-k2.6 only): https://github.com/sgl-project/sglang/actions/runs/28089396987

Note: mi355x has separately been blocked at the hf_xet download stall (7200s timeout) before load even starts, so its result here may still be masked by that.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant