Repository navigation
[Perf] Gluon attention backend with large head dimensions #1322
Description
Add an opt-in Gluon backend for attention, preserving the current causal attention contract: dense and packed variable-length inputs, MHA/GQA/MQA, unequal query/key and value dimensions, logarithmic gates, sliding windows, sink bias, and gradients. Include the corresponding packed single-query decoding features and cover head dimensions through 512.
The implementation will use explicit shared-memory staging, asynchronous data movement and MMA scheduling. It will preserve current operand/accumulation precision, the public API, existing tolerances, and default backend selection. Unsupported hardware/dtypes retain existing fallback behavior.
Validation will run the unchanged attention and sink tests plus dispatch and expanded boundary tests against fp32 references. Performance acceptance requires repeatable same-device improvements over the original Triton route, with dense/varlen forward, backward, and combined timings; unsupported baseline dimensions will be reported separately without invented speedup ratios. Benchmark results will include shapes, dtypes, software versions, and methodology.
Implementation references: the upstream Triton Gluon TMA/WGMMA/tcgen05 tutorials and attention example, adapted to the installed API and this operator's contract.
Activity
yzhangcs commentedon Oct 3, 2026
The additional bf16 + gate + sink + packed-sequence test exposes a sink-gradient error in both the existing Triton implementation and the initial Gluon implementation: relative RMS error 0.005750 and 0.005751 respectively against the fp32 reference (budget 0.005). Outputs and all other gradients pass. The existing sink backward uses
delta = sum(round_to_input_dtype(O) * dO), so output rounding propagates into the scalar sink-gradient reduction.Proposed numerical detail for the new opt-in backend: accumulate
sum(P * (dO @ V.T))in fp32 while recomputing the existing backward probabilities, then multiply by the negative sink probability and reduce across queries. This computes the same softmax derivative directly from backward operands without the rounded output intermediate. QK, dOV, probability, and reduction arithmetic remain fp32; tensor-core operands and accumulations, output dtypes, and the default Triton implementation remain unchanged. No test tolerance is relaxed. This may modestly change the opt-in sink gradient relative to the original implementation while improving agreement with the reference. Existing tests and explicit route parity remain required.yzhangcs commentedon Oct 5, 2026
The next backward experiment will fuse dQ into the existing dK/dV pass for aligned head dimensions up to 128. Each score and dP tile will then be computed once; dQ partials will accumulate into a zero-initialized fp32 buffer and convert to the input dtype after the kernel. P/dS operand casts, fp32 dot accumulation, reference tests, and tolerances remain unchanged. Gate and sink partials use the same fp32 reductions. Deterministic mode and ineligible shapes retain the existing two-pass implementation.
This is an experiment, not an accepted speedup. The frozen correctness gate and full feature matrix must pass before timing, and the existing benchmark runner will compare the complete forward/backward operation, including zeroing and conversion overhead. PR #1324 remains draft pending reproducible gains.