[PR #452] perf: Use full-rate FADD relu in SM90 FP8 MQA logits epilogue #457

New issue

Open

opened 2026-10-26 21:18:03 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:18:03 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/452
Author: @Chtholly-Boss
Created: 9/18/2026
Status: 🔄 Open

Base: main ← Head: sm90-mqa-logits-relu-fadd

📝 Commits (1)

ee53a16 Use full-rate FADD relu in SM90 FP8 MQA logits epilogue

📊 Changes

1 file changed (+3 additions, -3 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_mqa_logits.cuh (+3 -3)

📄 Description

Summary

Replace fmaxf(x, 0) with x + fabsf(x) in the SM90 FP8 MQA logits epilogue, recovering the 2× throughput gap between FMNMX (half-rate) and FADD (full-rate) on Hopper. The factor-of-2 introduced by the identity is resolved by folding 0.5f into the existing scale_kv multiply — zero extra instructions, bit-exact numerics.

Background

On SM90, fmaxf(x, 0) compiles to a single FMNMX, which issues at 64 ops/clk/SM — half the rate of FADD/FFMA (128 ops/clk/SM). Meanwhile x + |x| compiles to a single FADD with the absolute-value operand modifier, which is a free encoding bit in the instruction word:

FADD R3, R3, |R3|    ; one instruction, full rate (cuobjdump, sm_90a)

The relu executes once per (token, head, kv) score element in the MQA logits epilogue, so at 64 vs 128 ops/clk/SM the FMNMX form doubles the FP32 pipe time spent on relu alone.

Since x + |x| = 2 · relu(x) holds exactly for all finite IEEE-754 floats (+0 for x ≤ 0), the doubled sum is simply halved by folding * 0.5f into the already-existing scale_kv multiplication (an exponent decrement, also exact):

// Before:
return fmaxf(shifted_accum[j], 0) * weights[i][(j / 4) * 2 + (j & 1)];
...
float v_0 = (sum[0] + sum[1]) * scale_kv_0;

// After:
return (shifted_accum[j] + fabsf(shifted_accum[j])) * weights[i][(j / 4) * 2 + (j & 1)];
...
float v_0 = (sum[0] + sum[1]) * (scale_kv_0 * 0.5f);

The same technique is already used by the SM100 kernel (sm100_mqa_logits.cuh, packed FADD2 variant); this PR brings it to SM90. Verified via cuobjdump -sass: zero FMNMX across all JIT-instantiated sm90_fp8_mqa_logits cubins.

Performance (H200, 132 SMs, SM90A)

Production indexer shapes, best of 3 reps, before → after:

Shape	Before	After	Gain
H=32, D=128, SQ=512, SKV=140k, FP32 logits	1018 TFLOPS	1093 TFLOPS	+7.4%
SQ=512, SKV=130560, H=64, BF16, CP=1	1042 TFLOPS	1126 TFLOPS	+8.1%
SQ=512, SKV=130560, H=32, BF16, CP=0	941 TFLOPS	989 TFLOPS	+5.1%
SQ=2048, SKV=8192, H=64, BF16, CP=1	980 TFLOPS	1027 TFLOPS	+4.8%

Full tests/test_attention.py sweep (192 contiguous MQA cases on SM90): avg +3.7% (SK=65536 subset +4.4%, H=64 subset +5.5%, best case +11.3%). Gains scale with the relu's share of kernel time: larger SKV, more heads, smaller head dim → bigger win. Paged MQA logits (separate kernel, untouched) is unaffected.

Correctness

x + |x| = 2 · relu(x) exactly for all finite IEEE-754 floats (x > 0: x + x = 2x; x ≤ 0: x + (−x) = +0); the ÷2 via * 0.5f is an exact exponent decrement.
All 192 SM90 cases in tests/test_attention.py pass (numeric diff vs simulated reference, 20× bitwise self-consistency, mask equality) — verified on H200.
NaN semantics differ (x + |x| propagates NaN; fmaxf(NaN, 0) returns 0), but WGMMA accumulator outputs are finite in practice; this matches the SM100 kernel's existing behavior.

Changes

deep_gemm/include/deep_gemm/impls/sm90_fp8_mqa_logits.cuh — 3 lines modified (transform lambda + two * 0.5f folds).

_{🔄 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/deepseek-ai/DeepGEMM/pull/452 **Author:** [@Chtholly-Boss](https://github.com/Chtholly-Boss) **Created:** 9/18/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `sm90-mqa-logits-relu-fadd` --- ### 📝 Commits (1) - [`ee53a16`](https://github.com/deepseek-ai/DeepGEMM/commit/ee53a16b089d441cbfc4020cd558c1f05f623d97) Use full-rate FADD relu in SM90 FP8 MQA logits epilogue ### 📊 Changes **1 file changed** (+3 additions, -3 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_mqa_logits.cuh` (+3 -3) </details> ### 📄 Description ## Summary Replace `fmaxf(x, 0)` with `x + fabsf(x)` in the SM90 FP8 MQA logits epilogue, recovering the 2× throughput gap between FMNMX (half-rate) and FADD (full-rate) on Hopper. The factor-of-2 introduced by the identity is resolved by folding `0.5f` into the existing `scale_kv` multiply — zero extra instructions, bit-exact numerics. ## Background On SM90, `fmaxf(x, 0)` compiles to a single FMNMX, which issues at 64 ops/clk/SM — half the rate of FADD/FFMA (128 ops/clk/SM). Meanwhile `x + |x|` compiles to a single FADD with the absolute-value operand modifier, which is a free encoding bit in the instruction word: ``` FADD R3, R3, |R3| ; one instruction, full rate (cuobjdump, sm_90a) ``` The relu executes once per `(token, head, kv)` score element in the MQA logits epilogue, so at 64 vs 128 ops/clk/SM the FMNMX form doubles the FP32 pipe time spent on relu alone. Since `x + |x| = 2 · relu(x)` holds exactly for all finite IEEE-754 floats (`+0` for x ≤ 0), the doubled sum is simply halved by folding `* 0.5f` into the already-existing `scale_kv` multiplication (an exponent decrement, also exact): ```cpp // Before: return fmaxf(shifted_accum[j], 0) * weights[i][(j / 4) * 2 + (j & 1)]; ... float v_0 = (sum[0] + sum[1]) * scale_kv_0; // After: return (shifted_accum[j] + fabsf(shifted_accum[j])) * weights[i][(j / 4) * 2 + (j & 1)]; ... float v_0 = (sum[0] + sum[1]) * (scale_kv_0 * 0.5f); ``` The same technique is already used by the SM100 kernel (`sm100_mqa_logits.cuh`, packed `FADD2` variant); this PR brings it to SM90. Verified via `cuobjdump -sass`: zero FMNMX across all JIT-instantiated `sm90_fp8_mqa_logits` cubins. ## Performance (H200, 132 SMs, SM90A) Production indexer shapes, best of 3 reps, before → after: | Shape | Before | After | Gain | |---|---|---|---| | H=32, D=128, SQ=512, SKV=140k, FP32 logits | 1018 TFLOPS | 1093 TFLOPS | +7.4% | | SQ=512, SKV=130560, H=64, BF16, CP=1 | 1042 TFLOPS | 1126 TFLOPS | +8.1% | | SQ=512, SKV=130560, H=32, BF16, CP=0 | 941 TFLOPS | 989 TFLOPS | +5.1% | | SQ=2048, SKV=8192, H=64, BF16, CP=1 | 980 TFLOPS | 1027 TFLOPS | +4.8% | Full `tests/test_attention.py` sweep (192 contiguous MQA cases on SM90): **avg +3.7%** (SK=65536 subset +4.4%, H=64 subset +5.5%, best case +11.3%). Gains scale with the relu's share of kernel time: larger SKV, more heads, smaller head dim → bigger win. Paged MQA logits (separate kernel, untouched) is unaffected. ## Correctness - `x + |x| = 2 · relu(x)` exactly for all finite IEEE-754 floats (x > 0: `x + x = 2x`; x ≤ 0: `x + (−x) = +0`); the ÷2 via `* 0.5f` is an exact exponent decrement. - All 192 SM90 cases in `tests/test_attention.py` pass (numeric diff vs simulated reference, 20× bitwise self-consistency, mask equality) — verified on H200. - NaN semantics differ (`x + |x|` propagates NaN; `fmaxf(NaN, 0)` returns 0), but WGMMA accumulator outputs are finite in practice; this matches the SM100 kernel's existing behavior. ## Changes `deep_gemm/include/deep_gemm/impls/sm90_fp8_mqa_logits.cuh` — 3 lines modified (transform lambda + two `* 0.5f` folds). --- <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-26 21:18:03 +00:00

bsakaguchi referenced this issue

2026-10-26 21:18:06 +00:00

[PR #457] [Fix] Honor deterministic algorithms in SM90 bmk,bnk->mn einsum #462