[PR #70] fix(nsa): tilelang sparse kernel hardcoded head offset + sm_120 smem-fit geometry #69

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/70
Author: @architehc
Created: 7/23/2026
Status: 🔄 Open

Base: main ← Head: sm120-tilelang-head-offset

📝 Commits (1)

469966f fix(nsa): tilelang sparse kernel head offset + sm_120 smem-fit head split

📊 Changes

1 file changed (+32 additions, -14 deletions)

View changed files

📝 python/sglang/srt/layers/attention/nsa/tilelang_kernel.py (+32 -14)

📄 Description

`sparse_attention_fwd_kernel_v1` hardcodes `(bx % REPLICATE_H) * 64` for the head offset, which breaks any head split where `H_per_block != 64`. Fixed to `* H_per_block`.
Adds an sm_120 geometry: `H_per_block=16`, `threads=128`, v1 with `num_stages=1` — fits the 99 KB shared-memory budget of consumer/workstation Blackwell (the default geometries need 154–226 KB and only fit Hopper/SM100).
`kv_dtype` parametrization so the kernel can load fp8 KV pools directly.
Drops the `wg_wait` kwarg removed in tilelang ≥ 0.1.9.

With this, tilelang NSA prefill and decode run on sm_120 (validated with tilelang 0.1.12 + apache-tvm-ffi 0.1.11 + nvidia-cutlass-dsl 4.5.0, serving GLM-5.2 at 1M context).

🤖 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/70 **Author:** [@architehc](https://github.com/architehc) **Created:** 7/23/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `sm120-tilelang-head-offset` --- ### 📝 Commits (1) - [`469966f`](https://github.com/kvcache-ai/sglang/commit/469966f1f91e2efc6ac7db03794e438177583479) fix(nsa): tilelang sparse kernel head offset + sm_120 smem-fit head split ### 📊 Changes **1 file changed** (+32 additions, -14 deletions) <details> <summary>View changed files</summary> 📝 `python/sglang/srt/layers/attention/nsa/tilelang_kernel.py` (+32 -14) </details> ### 📄 Description - \`sparse_attention_fwd_kernel_v1\` hardcodes \`(bx % REPLICATE_H) * 64\` for the head offset, which breaks any head split where \`H_per_block != 64\`. Fixed to \`* H_per_block\`. - Adds an sm_120 geometry: \`H_per_block=16\`, \`threads=128\`, v1 with \`num_stages=1\` — fits the 99 KB shared-memory budget of consumer/workstation Blackwell (the default geometries need 154–226 KB and only fit Hopper/SM100). - \`kv_dtype\` parametrization so the kernel can load fp8 KV pools directly. - Drops the \`wg_wait\` kwarg removed in tilelang ≥ 0.1.9. With this, tilelang NSA prefill and decode run on sm_120 (validated with tilelang 0.1.12 + apache-tvm-ffi 0.1.11 + nvidia-cutlass-dsl 4.5.0, serving GLM-5.2 at 1M context). 🤖 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#69

No description provided.

Rows
Columns