Skip to content

[SM100] mqa_logits: faster bf16-weight paths, including the dev contiguous kernel for MX scales (release/v0.2.3) - #117

Merged
hnyls2002 merged 2 commits into
release/v0.2.3from
lsyin/mqa-logits-3wg-v0.2.3
Oct 23, 2026
Merged

hnyls2002 merged 2 commits into
release/v0.2.3from
lsyin/mqa-logits-3wg-v0.2.3

Conversation

@hnyls2002

@hnyls2002 hnyls2002 commented Oct 23, 2026 •

Copy link
Copy Markdown

Brings the faster SM100 fp8_fp4_mqa_logits paths to release/v0.2.3, for a post release before the next minor.

  • MX scales (MXFP4 / MXFP8), BF16 weights, compressed BF16 logits, no cleaning, no schedule metadata: these calls now run a backport of the dev SM100 contiguous kernel as sm100_bf16_mqa_logits (new impls/, layout/ and scheduler/ headers under that name, plus a host launcher). It uses a 384-token KV split with three math warpgroups, separate KV and scale-factor producers, a deeper FP4 pipeline (2 Q / 8 KV stages), and reuses resident KV stages when consecutive Q blocks cover the same splits. The paged entry and the schedule-metadata builder of the dev kernel are left out; ptx::st_global_low_bf16 is added for its epilogue.
  • Other BF16-weight calls without cleaning or a schedule (e.g. per-KV-scale FP8): the existing kernel takes a 384-token split with three math warpgroups instead of 256 with two. The math warpgroups split the CTA register pool four ways (statically checked), and the 384-token split loads each KV stage in 128-token TMA pieces.
  • Every other variant (FP32 weights, cleaned logits, schedule metadata, paged) is unchanged.

The logits are unchanged: both kernels compute the same products and reductions per row; only the KV tiling, warp assignment and pipelining differ. Logits from the new paths match the previous kernel bit for bit on SM100 for MXFP4, MXFP8 and FP8 inputs across several row counts and KV widths, and the repository's MQA logits tests pass, covering every case of the new path plus a sample of the others, the sparse cases and the paged cases.

@hnyls2002 hnyls2002 changed the title [SM100] mqa_logits: three math warpgroups with a 384-token split for bf16 head weights (release/v0.2.3) [SM100] mqa_logits: faster bf16-weight paths, including the dev contiguous kernel for MX scales (release/v0.2.3) Oct 23, 2026
@hnyls2002
hnyls2002 merged commit 5b31b3f into release/v0.2.3 Oct 23, 2026
@Fridge003
Fridge003 deleted the lsyin/mqa-logits-3wg-v0.2.3 branch October 23, 2026 04:19
zcnrex added a commit to zcnrex/DeepGEMM that referenced this pull request Nov 2, 2026
…3 calls to it

The `_legacy` APIs now carry 0.2.3.post2's MQA logits for every architecture: the host entry points and the
SM90, SM100 (including the BF16 contiguous kernel from release sgl-project#117) and SM120 kernels, copied under
`deep_gemm::legacy` with `legacy_`-prefixed JIT names.

Callers written against 0.2.3 run unchanged:
- `fp8_mqa_logits` / `fp8_paged_mqa_logits` are restored on top of the `_legacy` functions.
- `fp8_fp4_mqa_logits` / `fp8_fp4_paged_mqa_logits` take the `_legacy` path when called with a bool
  `clean_logits` after the fixed arguments, or with `clean_logits` / `logits_dtype` keywords (and the dense
  call without `max_seqlen_k`).
- `get_paged_mqa_logits_metadata` without `indices` returns the 0.2.3 schedule on SM100.

`test_attention_legacy.py` runs on all architectures and checks the 0.2.3 call shapes bitwise against the
`_legacy` calls.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
zcnrex added a commit that referenced this pull request Nov 2, 2026
…3 calls to it

The `_legacy` APIs now carry 0.2.3.post2's MQA logits for every architecture: the host entry points and the
SM90, SM100 (including the BF16 contiguous kernel from release #117) and SM120 kernels, copied under
`deep_gemm::legacy` with `legacy_`-prefixed JIT names.

Callers written against 0.2.3 run unchanged:
- `fp8_mqa_logits` / `fp8_paged_mqa_logits` are restored on top of the `_legacy` functions.
- `fp8_fp4_mqa_logits` / `fp8_fp4_paged_mqa_logits` take the `_legacy` path when called with a bool
  `clean_logits` after the fixed arguments, or with `clean_logits` / `logits_dtype` keywords (and the dense
  call without `max_seqlen_k`).
- `get_paged_mqa_logits_metadata` without `indices` returns the 0.2.3 schedule on SM100.

`test_attention_legacy.py` runs on all architectures and checks the 0.2.3 call shapes bitwise against the
`_legacy` calls.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
zcnrex added a commit that referenced this pull request Nov 2, 2026
…llers working (#128)

* Restore the 0.2.3 SM100 MQA logits kernel as `_legacy` APIs

0.2.5 made the SM100 MQA logits kernel MX-only. sglang keeps FP32 index-K
scales as its default, so it still needs the 0.2.3 SM100 kernel, which
accepts FP32 or MX scales, FP32 or BF16 weights/logits, and cleaning.

Add SM100-only `fp8_fp4_mqa_logits_legacy`, `fp8_fp4_paged_mqa_logits_legacy`,
`get_mqa_logits_metadata_legacy` and `get_paged_mqa_logits_metadata_legacy`
with the 0.2.3 signatures. The kernel, schedulers, layout and clean-logits
code are copied unchanged from 3cca390^1 into `legacy/` directories under
the `deep_gemm::legacy` namespace. The current APIs are untouched.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* Freeze the 0.2.3.post2 MQA logits for SM90/SM100/SM120 and route 0.2.3 calls to it

The `_legacy` APIs now carry 0.2.3.post2's MQA logits for every architecture: the host entry points and the
SM90, SM100 (including the BF16 contiguous kernel from release #117) and SM120 kernels, copied under
`deep_gemm::legacy` with `legacy_`-prefixed JIT names.

Callers written against 0.2.3 run unchanged:
- `fp8_mqa_logits` / `fp8_paged_mqa_logits` are restored on top of the `_legacy` functions.
- `fp8_fp4_mqa_logits` / `fp8_fp4_paged_mqa_logits` take the `_legacy` path when called with a bool
  `clean_logits` after the fixed arguments, or with `clean_logits` / `logits_dtype` keywords (and the dense
  call without `max_seqlen_k`).
- `get_paged_mqa_logits_metadata` without `indices` returns the 0.2.3 schedule on SM100.

`test_attention_legacy.py` runs on all architectures and checks the 0.2.3 call shapes bitwise against the
`_legacy` calls.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant