kvcache-ai/sglang

Fork 0

mirror of https://github.com/kvcache-ai/sglang.git synced 2026-11-05 01:07:39 +00:00

[PR #97] fix(kt): apply the KT GPU-expert remap when there is no expert-location metadata #97

New issue

Open

opened 2026-10-29 21:07:10 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-29 21:07:10 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/sglang/pull/97
Author: @05yuki
Created: 9/25/2026
Status: 🔄 Open

Base: main ← Head: fix/kt-gpu-expert-remap-no-metadata

📝 Commits (1)

f3080b4 fix(kt): apply the KT GPU-expert remap when there is no expert-location metadata

📊 Changes

1 file changed (+28 additions, -19 deletions)

View changed files

📝 python/sglang/srt/layers/moe/fused_moe_triton/layer.py (+28 -19)

📄 Description

Problem

With --kt-num-gpu-experts N and any placement other than logical experts 0..N-1 (--kt-expert-placement-strategy frequency, random, or a placement file), the resident GPU experts hold the wrong weights.

FusedMoE.weight_loader only maps a logical expert to its KT GPU slot inside _weight_loader_physical, and that function is only reached when get_global_expert_location_metadata() is set. When it is None, which is the case for a plain KT launch, the loader does the global-to-local mapping and calls _weight_loader_impl directly. Logical expert e is then written into GPU slot e, and experts past the slot count are dropped by the bounds check. The wrapper still routes its selected experts (logical_to_gpu_index) to those slots, so every token routed to a GPU expert is computed with another expert's weights.

The uniform placement is unaffected, because its map is the identity. That is probably why this has not shown up.

How it was found

GLM-5.3 (FP8 checkpoint, NVFP4 CPU experts, TP2, 2x RTX 5060 Ti on a dual EPYC 7452), --kt-num-gpu-experts 4 --kt-expert-placement-strategy frequency. Perplexity rose with the share of routing that went to the GPU (hot-4 covers 12.3% of picks). I dumped the resident expert parameters in process_weights_after_loading and compared them byte for byte with the checkpoint:

layer 10, both TP ranks, before and after post-processing
gpu_index_to_logical [76, 145, 262, 268]
slot 0: claims 76   -> holds expert 0
slot 1: claims 145  -> holds expert 1
slot 2: claims 262  -> holds expert 2
slot 3: claims 268  -> holds expert 3

After the fix, all four slots match the claimed experts' FP8 bytes (w13 and w2, both ranks).

Change

The KT remap moves into FusedMoE._map_kt_gpu_expert_id() and is called from both branches: the existing _weight_loader_physical path, which behaves as before, and the no-metadata path, which now remaps too. Parameters marked _sglang_require_global_experts are left alone on both paths, as before.

The behaviour for fused shared experts does not change. Their local id is outside gpu_experts_mask, so they are skipped under KT on both paths, as _weight_loader_physical already did.

Results

Paired NLL over 7,664 tokens (fixed held-out text, teacher-forced, per-token comparison against the same server with 0 GPU experts). GPU prefill was off so that every token takes the hybrid path.

config	mean NLL	delta vs 0 GPU experts	z
0 GPU experts	2.5580
frequency hot-4, before	5.4788	+2.9208	+95.6
frequency hot-4, after	2.5478	-0.0102	-2.0

The small gain after the fix is expected: here the GPU experts run on the original FP8 while the CPU experts run on NVFP4. Decode with hot-4 goes from 17.2 to 18.5 tok/s (3 x 1,500-token completions each).

🤖 Generated with Claude Code

https://claude.ai/code/session_01R57AidCMwH17PTFpX7HfQ5

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/kvcache-ai/sglang/pull/97 **Author:** [@05yuki](https://github.com/05yuki) **Created:** 9/25/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/kt-gpu-expert-remap-no-metadata` --- ### 📝 Commits (1) - [`f3080b4`](https://github.com/kvcache-ai/sglang/commit/f3080b4238e0ddf845c8911306a1f0bca623cef9) fix(kt): apply the KT GPU-expert remap when there is no expert-location metadata ### 📊 Changes **1 file changed** (+28 additions, -19 deletions) <details> <summary>View changed files</summary> 📝 `python/sglang/srt/layers/moe/fused_moe_triton/layer.py` (+28 -19) </details> ### 📄 Description ## Problem With `--kt-num-gpu-experts N` and any placement other than logical experts `0..N-1` (`--kt-expert-placement-strategy frequency`, `random`, or a placement file), the resident GPU experts hold the wrong weights. `FusedMoE.weight_loader` only maps a logical expert to its KT GPU slot inside `_weight_loader_physical`, and that function is only reached when `get_global_expert_location_metadata()` is set. When it is `None`, which is the case for a plain KT launch, the loader does the global-to-local mapping and calls `_weight_loader_impl` directly. Logical expert `e` is then written into GPU slot `e`, and experts past the slot count are dropped by the bounds check. The wrapper still routes its selected experts (`logical_to_gpu_index`) to those slots, so every token routed to a GPU expert is computed with another expert's weights. The `uniform` placement is unaffected, because its map is the identity. That is probably why this has not shown up. ## How it was found GLM-5.3 (FP8 checkpoint, NVFP4 CPU experts, TP2, 2x RTX 5060 Ti on a dual EPYC 7452), `--kt-num-gpu-experts 4 --kt-expert-placement-strategy frequency`. Perplexity rose with the share of routing that went to the GPU (hot-4 covers 12.3% of picks). I dumped the resident expert parameters in `process_weights_after_loading` and compared them byte for byte with the checkpoint: ``` layer 10, both TP ranks, before and after post-processing gpu_index_to_logical [76, 145, 262, 268] slot 0: claims 76 -> holds expert 0 slot 1: claims 145 -> holds expert 1 slot 2: claims 262 -> holds expert 2 slot 3: claims 268 -> holds expert 3 ``` After the fix, all four slots match the claimed experts' FP8 bytes (w13 and w2, both ranks). ## Change The KT remap moves into `FusedMoE._map_kt_gpu_expert_id()` and is called from both branches: the existing `_weight_loader_physical` path, which behaves as before, and the no-metadata path, which now remaps too. Parameters marked `_sglang_require_global_experts` are left alone on both paths, as before. The behaviour for fused shared experts does not change. Their local id is outside `gpu_experts_mask`, so they are skipped under KT on both paths, as `_weight_loader_physical` already did. ## Results Paired NLL over 7,664 tokens (fixed held-out text, teacher-forced, per-token comparison against the same server with 0 GPU experts). GPU prefill was off so that every token takes the hybrid path. | config | mean NLL | delta vs 0 GPU experts | z | |---|---|---|---| | 0 GPU experts | 2.5580 | | | | frequency hot-4, before | 5.4788 | +2.9208 | +95.6 | | frequency hot-4, after | 2.5478 | -0.0102 | -2.0 | The small gain after the fix is expected: here the GPU experts run on the original FP8 while the CPU experts run on NVFP4. Decode with hot-4 goes from 17.2 to 18.5 tok/s (3 x 1,500-token completions each). 🤖 Generated with [Claude Code](https://claude.com/claude-code) https://claude.ai/code/session_01R57AidCMwH17PTFpX7HfQ5 --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-29 21:07:10 +00:00

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

kvcache-ai/sglang#97

No description provided.

Rows
Columns