[PR #69] feat(nsa): sm_120 support — MQA logits torch/Triton replacements + NSA backend fallbacks #68

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

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

Base: main ← Head: sm120-mqa-logits-nsa-fallbacks

📝 Commits (1)

1ef82ef feat(nsa): sm_120 support - MQA logits replacements + capability-gated backend fallbacks

📊 Changes

4 files changed (+372 additions, -15 deletions)

View changed files

📝 python/sglang/srt/layers/attention/nsa/nsa_indexer.py (+61 -7)
➕ python/sglang/srt/layers/attention/nsa/sm120_mqa_logits.py (+90 -0)
➕ python/sglang/srt/layers/attention/nsa/sm120_mqa_logits_triton.py (+85 -0)
📝 python/sglang/srt/layers/attention/nsa_backend.py (+136 -8)

📄 Description

Enables the NSA attention path on sm_120 (RTX PRO 6000 Blackwell Workstation), where deep_gemm's `fp8_mqa_logits` / `fp8_paged_mqa_logits` raise "Unsupported architecture".

`sm120_mqa_logits.py` + `sm120_mqa_logits_triton.py`: torch and Triton replacements for the ragged and paged MQA logits. The Triton ragged kernel uses an fp8 tensor-core dot with a K-stripe inner loop and reaches ~306 effective TFLOPS at rows=2048 / L=900k on sm_120; int64 output offsets fix real overflow at rows×L > 2^31 elements.
Probe-once dispatchers in `nsa_indexer.py` (first-call probe, then genuine runtime failures propagate instead of silently degrading to torch).
NSA backend fallbacks: capability-gated torch `fast_topk_v2`; CUDA-graph-safe per-q-token page-table expansion (searchsorted + index_select — `repeat_interleave`'s data-dependent shapes corrupt outputs under graph capture); frozen-dataclass-safe schedule-metadata assignment; fp8-KV gather-dequant hop in `_forward_tilelang` (dequantizes only the gathered top-k rows).

Validated end-to-end serving GLM-5.2 at up to 1M context on a single RTX PRO 6000 (details: 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/69 **Author:** [@architehc](https://github.com/architehc) **Created:** 7/23/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `sm120-mqa-logits-nsa-fallbacks` --- ### 📝 Commits (1) - [`1ef82ef`](https://github.com/kvcache-ai/sglang/commit/1ef82eff1c902bfe5959dc5ddc134544e40829fc) feat(nsa): sm_120 support - MQA logits replacements + capability-gated backend fallbacks ### 📊 Changes **4 files changed** (+372 additions, -15 deletions) <details> <summary>View changed files</summary> 📝 `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` (+61 -7) ➕ `python/sglang/srt/layers/attention/nsa/sm120_mqa_logits.py` (+90 -0) ➕ `python/sglang/srt/layers/attention/nsa/sm120_mqa_logits_triton.py` (+85 -0) 📝 `python/sglang/srt/layers/attention/nsa_backend.py` (+136 -8) </details> ### 📄 Description Enables the NSA attention path on **sm_120** (RTX PRO 6000 Blackwell Workstation), where deep_gemm's \`fp8_mqa_logits\` / \`fp8_paged_mqa_logits\` raise \"Unsupported architecture\". - \`sm120_mqa_logits.py\` + \`sm120_mqa_logits_triton.py\`: torch and Triton replacements for the ragged and paged MQA logits. The Triton ragged kernel uses an fp8 tensor-core dot with a K-stripe inner loop and reaches ~306 effective TFLOPS at rows=2048 / L=900k on sm_120; int64 output offsets fix real overflow at rows×L > 2^31 elements. - Probe-once dispatchers in \`nsa_indexer.py\` (first-call probe, then genuine runtime failures propagate instead of silently degrading to torch). - NSA backend fallbacks: capability-gated torch \`fast_topk_v2\`; CUDA-graph-safe per-q-token page-table expansion (searchsorted + index_select — \`repeat_interleave\`'s data-dependent shapes corrupt outputs under graph capture); frozen-dataclass-safe schedule-metadata assignment; fp8-KV gather-dequant hop in \`_forward_tilelang\` (dequantizes only the gathered top-k rows). Validated end-to-end serving GLM-5.2 at up to 1M context on a single RTX PRO 6000 (details: 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:55 +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#68

No description provided.

Rows
Columns