Repository navigation
[SM100] mqa_logits: faster bf16-weight paths, including the dev contiguous kernel for MX scales (release/v0.2.3) - #117
Merged
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Brings the faster SM100
fp8_fp4_mqa_logitspaths torelease/v0.2.3, for a post release before the next minor.devSM100 contiguous kernel assm100_bf16_mqa_logits(newimpls/,layout/andscheduler/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 thedevkernel are left out;ptx::st_global_low_bf16is added for its epilogue.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.