[ormandj]

Motivation

#103 removed the SM120 dense and paged MQA logits kernels, which served the old attention API. Since then, fp8_fp4_mqa_logits and fp8_fp4_paged_mqa_logits have no SM120 path on dev.

Modifications

• Kernels and host launcher: restored from deepseek-ai/DeepGEMM nv_dev at b64107f2, which has fixes made after Merge upstream main (26/09/30 release + #471) into dev #103's copy:
  • the next Q tile follows the scheduler's actual next task;
  • a shared-memory fence orders reads before a stage is reused;
  • paged KV accepts 32-, 128- and 256-token pages as well as 64.
• SM120 on the new API:
  • FP32 logits;
  • compressed, uncleaned dense logits;
  • FP8 with one float scale per KV token, FP4 with MX scales;
  • paged requests varlen or fixed next_n.
• Limits:
  • 16, 32 or 64 heads (8 was accepted before Merge upstream main (26/09/30 release + #471) into dev #103);
  • no dense schedule_meta;
  • varlen metadata fits about 8K query tokens.
• Other architectures: SM90 and SM100 are unchanged.
• Tests: test_attention.py covers SM120:
  • FP8 and MXFP4 dense logits, plus short-query cases with nonzero starts;
  • paged logits, varlen and fixed next_n 1-5, all four page sizes, with empty contexts and a batch whose contexts are all empty.

Validation

On four RTX PRO 6000 Blackwell (SM120, CUDA 13.0):

• tests/test_attention.py passes: 96 dense, 36 short-query and 1,680 paged cases (480 varlen), plus the all-empty paged case.
• Performance: dense kernel time against the pre-Merge upstream main (26/09/30 release + #471) into dev #103 kernels (cbe1727e) at 2K/8K queries x 8K/64K keys and 32/64 heads, timed with CUDA events (median of 15 per shape, two runs):
  • FP8 is within -0.1% to +2.6%;
  • MXFP4 is within -0.7% to +0.9%.

Developed with AI assistance.

[zcnrex]

Tested on B200 (SM100, CUDA 13.0) to check that SM90/SM100 are unaffected. SM120 itself is not tested here.

Setup: this PR's head 51ceada merged cleanly onto current dev (d0a3af7), built with build_sgl_deep_gemm.sh.

• tests/test_attention.py: exits 0, with 1789 timed cases and no errors or assertion failures. This covers dense, short-query, all-empty paged, paged and sparse MQA logits on SM100.
• sglang DSA/DSV4 indexer suite (13 test files, run on the [Refactor] [NVIDIA] attention: Migrate DSA/DSV4 indexers to DeepGEMM 0.2.5.post2 sglang#44350 branch, which moves the indexers onto the fp8_fp4_* API): 524 passed, 29 skipped. That matches plain dev exactly. The skips are AMD/Hopper-only cases and the SM100 metadata-format comparison.
• The shared host checks in csrc/apis/attention.hpp keep the same conditions for arch_major 9 and 10. Only == 12 branches are added.

🤖 Generated with Claude Code