[PR #68] fix(moe): sm_120 topk_sigmoid writes only 1408 rows of topk_weights — pure-torch routing fallback #67

New issue

Open

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

bsakaguchi commented

2026-10-29 21:06:54 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/kvcache-ai/sglang/pull/68
Author: @architehc
Created: 7/23/2026
Status: 🔄 Open

Base: main ← Head: sm120-topk-sigmoid-row-truncation

📝 Commits (1)

994dab4 fix(moe): pure-torch sigmoid routing fallback for sm_120 topk_sigmoid row truncation

📊 Changes

1 file changed (+20 additions, -10 deletions)

View changed files

📝 python/sglang/srt/layers/moe/topk.py (+20 -10)

📄 Description

Bug

`sgl_kernel.topk_sigmoid` (0.3.21) on sm_120 (RTX PRO 6000 Blackwell Workstation) writes `topk_weights` for only the first 1408 rows — `topk_ids` are fully written, so tail rows silently consume stale allocator memory as routing weights. Because the same-size buffer is recycled every layer, tail rows first reuse the previous layer's weights (finite → subtle quality degradation), then produce NaNs once poisoned memory rotates in (presents as `!!!!` output in long prefills).

Standalone kernel tests pass because fresh zero pages mask the unwritten rows; we reproduced the corruption byte-exact by pre-poisoning the output buffers before the call.

Fix

`fused_topk`: fall back to eager sigmoid routing (sigmoid → +bias topk → gather → eps renorm) when `M` exceeds a capability-gated row bound on sm_120 (`SGLANG_SIGMOID_FUSED_MAXROWS`, default 1408; `0` disables the fused kernel entirely). Routing cost is negligible relative to the corruption.
`1e-20` eps on the three renormalize sites to kill 0/0 NaNs independently.

The real fix belongs in the sgl_kernel CUDA kernel's row coverage; this makes sglang correct on sm_120 meanwhile.

Validated serving GLM-5.2 at up to 1M context on the RTX PRO 6000 (full write-up: https://github.com/architehc/blackwell_ktransformers).

🤖 Generated with Claude Code

_{🔄 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/68 **Author:** [@architehc](https://github.com/architehc) **Created:** 7/23/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `sm120-topk-sigmoid-row-truncation` --- ### 📝 Commits (1) - [`994dab4`](https://github.com/kvcache-ai/sglang/commit/994dab4c3432640bdae543439cc27d8a9be468a8) fix(moe): pure-torch sigmoid routing fallback for sm_120 topk_sigmoid row truncation ### 📊 Changes **1 file changed** (+20 additions, -10 deletions) <details> <summary>View changed files</summary> 📝 `python/sglang/srt/layers/moe/topk.py` (+20 -10) </details> ### 📄 Description ## Bug \`sgl_kernel.topk_sigmoid\` (0.3.21) on **sm_120** (RTX PRO 6000 Blackwell Workstation) writes \`topk_weights\` for only the **first 1408 rows** — \`topk_ids\` are fully written, so tail rows silently consume stale allocator memory as routing weights. Because the same-size buffer is recycled every layer, tail rows first reuse the *previous layer's* weights (finite → subtle quality degradation), then produce NaNs once poisoned memory rotates in (presents as \`!!!!\` output in long prefills). Standalone kernel tests pass because fresh zero pages mask the unwritten rows; we reproduced the corruption byte-exact by pre-poisoning the output buffers before the call. ## Fix - \`fused_topk\`: fall back to eager sigmoid routing (sigmoid → +bias topk → gather → eps renorm) when \`M\` exceeds a capability-gated row bound on sm_120 (\`SGLANG_SIGMOID_FUSED_MAXROWS\`, default 1408; \`0\` disables the fused kernel entirely). Routing cost is negligible relative to the corruption. - \`1e-20\` eps on the three renormalize sites to kill 0/0 NaNs independently. The real fix belongs in the sgl_kernel CUDA kernel's row coverage; this makes sglang correct on sm_120 meanwhile. Validated serving GLM-5.2 at up to 1M context on the RTX PRO 6000 (full write-up: https://github.com/architehc/blackwell_ktransformers). 🤖 Generated with [Claude Code](https://claude.com/claude-code) --- <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:06:54 +00:00

bsakaguchi referenced this issue

2026-10-29 21:06:55 +00:00

[PR #67] [CLOSED] [Elastic EP] Fix recovery lifecycle and add manual coverage #70