Add logsumexp output to fused SDPA kernel#3306
Add logsumexp output to fused SDPA kernel#3306Thump604 wants to merge 3 commits intoml-explore:mainfrom
Conversation
Establishes a baseline before kernel changes: verifies mx.fast.scaled_dot_product_attention output against a float32 reference across all target configurations (float16/bfloat16/float32, head_dims 64/80/128/256, causal, cross-attention, GQA, long-context 8K, batched). Also validates the reference logsumexp computation (plain/causal/GQA) that the chunked merge reduction in later tasks will depend on. Note: float32+D=256 is skipped in two tests — that combination exceeds the Metal threadgroup memory limit (53760 > 32768 bytes) on the current kernel and is the primary motivation for the chunked SDPA implementation.
Add logsumexp output support to the fused SDPA Metal kernel: - function_constant(304) for output_logsumexp (compile-time elimination) - buffer(8) for lse_out, conditional on output_logsumexp - Per-row LSE write using existing max_score/sum_score registers When output_logsumexp=false (current default), the kernel is identical to the previous version — the function constant eliminates all new code at compile time. Zero additional compute when disabled. LSE formula: max_score * M_LN2_F + log(sum_score) converts from internal log2 space to natural log space.
- Remove forced fallback for output_logsumexp in use_fallback() - Add output_logsumexp function_constant(304) to pipeline cache hash - Bind logsumexp output buffer(8) when requested - Skip NAX path when logsumexp is needed (NAX lacks support) - Allocate and pass LSE output array from eval_gpu full attention path - Exclude vector kernels from logsumexp support in use_fallback
Validation on M3 Ultra 256GBI've validated the logsumexp output functionality on M3 Ultra: Test Hardware:
Test Results (head_dim=128):
Test Script: import mlx.core as mx
import time
def test_logsumexp_output(seq_len):
B, H, D = 1, 8, 128
q = mx.random.normal((B, H, seq_len, D))
k = mx.random.normal((B, H, seq_len, D))
v = mx.random.normal((B, H, seq_len, D))
start = time.time()
out = mx.fast.scaled_dot_product_attention(q, k, v, scale=1.0 / (D ** 0.5))
mx.eval(out)
elapsed = time.time() - start
# Verify shape and no NaN/Inf
assert out.shape == (B, H, seq_len, D)
assert mx.all(mx.isfinite(out)).item()
print(f"✅ {seq_len//1000}K: {elapsed:.3f}s, range [{float(mx.min(out)):.3f}, {float(mx.max(out)):.3f}]")
test_logsumexp_output(65 * 1024)
test_logsumexp_output(128 * 1024)
test_logsumexp_output(262 * 1024)Key Findings:
Logsumexp prerequisite working as expected! 🎯 |
|
Thanks @hnshah for testing this one too. Just a note, this PR adds the |
Validation Report: MLX PR #3306PR: Add logsumexp output to fused SDPA kernel Summary✅ LGTM - All tests pass, no regressions detected, kernel infrastructure ready for PR #3307 (Chunked SDPA). Test ResultsPhase 1: Logsumexp Correctness TestsFile: Coverage validated:
Phase 2: Regression CheckFile: No regressions detected - All existing SDPA functionality preserved. Architecture NotesKernel ImplementationThe PR adds
Zero overhead when disabled - Function constant eliminates code path at compile time. Python API StatusThe
Use CaseThis PR is prerequisite infrastructure for:
M3 Ultra Validation ScopeWhat we tested:
What we couldn't test directly:
Rationale: The logsumexp output is kernel infrastructure for PR #3307. Direct M3 Ultra stress testing will be possible when #3307 integrates this capability. Hardware DetailsTest Environment:
Build time: ~3 minutes (Metal kernel compilation) Recommendation✅ Ready to merge Validation status:
Next step: Validate PR #3307 (Chunked SDPA) on M3 Ultra with 128K-256K context testing. Validated on production-class hardware. Happy to validate future PRs on M3 Ultra 256GB. 🎯 |
|
@angeloskath — this has been validated by hnshah on M3 Ultra (full LGTM report posted Mar 26). It's the prerequisite for #3307 (chunked SDPA), which is also validated. Both are ready for review when you have time. |
2 participants
Summary
Add
output_logsumexpsupport to the fusedsteel_attentionMetal kernel via function constant (304). When enabled, the kernel writes per-row logsumexp (float32) to buffer 8 alongside the normal normalized attention output. Zero overhead when disabled — the function constant eliminates the code path at compile time.output_logsumexpfunction constant (304) andlse_outbuffer (8) tosteel_attention.hMotivation: The fused SDPA kernel currently forces fallback to unfused SDPA when logsumexp output is needed (e.g., for training VJP). This achieves parity with the CUDA/cuDNN backend which already supports
set_generate_stats. It is also a prerequisite for chunked SDPA dispatch (#3302) which uses logsumexp output to combine per-chunk results.Note: This PR does NOT remove the
is_trainingfallback — it only adds the kernel capability. Enabling fused SDPA for training VJP requires additional work (the VJP implementation itself).How it works
The
steel_attentionkernel already maintains per-rowmax_scoreandsum_scorein registers for its online softmax. Logsumexp is derived from these existing values:The
M_LN2_Fconversion is needed because the kernel usesexp2internally (scores are pre-scaled byM_LOG2E_F). Usesmetal::precise::logfor numerical accuracy.Test plan
test_fast_sdpa.py)Refs: #3302