REVIEW 4 major objections 4 minor 34 references
FLASH-D: FlashAttention with Hidden Softmax Division
T0 review · 4 major / 4 minor · reviewed 2026-08-07 · deepseek-v4-flash
Pith's one-line read FlashAttention can be rewritten exactly so each softmax weight is a sigmoid of a neighboring score difference plus a log-weight term, removing max subtraction and simplifying hardware.
desk verdict Correct core identity and real hardware savings, but a sign typo in Algorithm 3 and an unsafe static skip-rule overclaim make this a revise, not an accept. read the letter →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
The practical gain is hardware simplicity. The kernel no longer needs a separate division unit, a running maximum unit, or a running sum-of-exponentials unit. The paper implements both FlashAttention2 and FLASH-D as unrolled systolic kernels in 28nm ASIC technology and reports average area and power reductions of about 23% and 20%, at the same clock frequency and with the same dataflow. The authors also observe that in real LLM benchmarks many output updates can be skipped if the sigmoid argument is strongly saturated, creating additional potential savings. The paper's hardware uses 8-segment piecewise-linear approximations for sigmoid and logarithm, so the exact mathematical equivalence is only part of the story; numerical equivalence at the hardware level needs separate verification.
Extended reading notes
Core claim
The paper claims that FLASH-D is a mathematically equivalent reformulation of the FlashAttention forward pass, replacing the recursive maximum, sum-of-exponentials, and division with the sigmoid recurrence w_i = sigmoid(s_i - s_{i-1} + ln w_{i-1}), eliminating the need for max subtraction, and that a 28nm ASIC implementation reduces average area by 22.8% and power by 20.3% versus a FlashAttention2 kernel at the same performance. If correct, Alg. 3 produces identical outputs to Alg. 1 in exact arithmetic and enables simpler hardware for exact attention.
Load-bearing premise
The paper assumes that the static range check on the score difference si - si-1 outside [-6, 11] is sufficient to saturate the sigmoid and safely simplify output updates. This ignores the ln wi-1 term in the sigmoid argument: if wi-1 is extremely small, ln wi-1 is very negative, and a large positive score difference can still yield a weight near 0, not 1. Forcing wi=1 in that case discards the accumulated output, so the claimed 'without affecting the outcome' is not guaranteed. This assumption enters in Section III-C, Fig. 2, and the output-update simplification described with Alg. 3, and it is used again in Section V.B.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
Summary. The paper derives a mathematically equivalent reformulation of the FlashAttention forward pass, replacing the online maximum, running sum-of-exponentials, and explicit division with the recurrence w_i = sigma(s_i - s_{i-1} + ln w_{i-1}) and the output update o_i = o_{i-1}(1-w_i) + v_i w_i. It further claims that score differences outside [-6, 11] allow safe skipping of the sigmoid and of the output update, that these skips are always beneficial 'without affecting the outcome,' and that a 28nm ASIC implementation of FLASH-D reduces average area by 22.8% and power by 20.3% relative to a parallel FlashAttention2 kernel. The exact-arithmetic derivation in Section III is correct and self-contained, but the static skip rule is not exact, the pseudo-code has a sign inconsistency with the derived recurrence, and the hardware comparison does not fully account for the approximations used in the FLASH-D nonlinear units.
Significance. If the claims were fully supported, this would be a useful contribution: the algebraic reformulation is elegant, preserves the tiled/IO-friendly structure of FlashAttention, removes the need for maximum subtraction, and hides the softmax division inside sigmoid/log evaluations. The paper gives a step-by-step derivation, provides public code, and reports concrete 28nm synthesis results, which are verifiable artifacts. However, the central 'no tradeoff' claim rests on a static saturation criterion that is not valid in the upper-skip branch, because the sigmoid argument includes the negative quantity ln w_{i-1}. The exact equivalence is also only for ideal real-arithmetic sigmoid/log functions; the implemented hardware uses 8-segment PWL approximations. These issues do not invalidate the core algebraic identity, but they do mean the advertised simplifications and the hardware-savings story need to be reworked before the paper can be accepted.
major comments (4)
- [Section III-B and Algorithm 3] Equation (11) correctly derives w_i = sigma(s_i - s_{i-1} + ln w_{i-1}), but Algorithm 3, line 5, writes w_i <- sigma(s_i - s_{i-1} - ln w_{i-1}). These two recurrences are different, and the printed minus sign does not reproduce the baseline softmax in exact arithmetic. This is not merely cosmetic: the sign of ln w_{i-1} determines whether the sigmoid input is shifted left or right, and it also affects the saturation arguments in Section III-C. Please correct the algorithm and state clearly whether the Python and hardware implementations use Eq. (11) or the printed line 5.
- [Section III-C, Fig. 2, and Section V.B] The static saturation rule based only on s_i - s_{i-1} is not valid for the upper skip. Since the sigmoid argument is s_i - s_{i-1} + ln w_{i-1} and ln w_{i-1} <= 0, the condition s_i - s_{i-1} >= 11 does not imply that the sigmoid is near 1. For example, if w_{i-1} = 1e-20, then ln w_{i-1} is about -46, so a score difference of 20 yields a sigmoid argument of about -26 and w_i is near 0, whereas the rule in Section III-C would set w_i = 1 and replace the accumulated output with v_i, discarding all previous contributions. The lower-skip branch is conservative because ln w_{i-1} makes the argument more negative, but the upper branch is not. Consequently, Table I overcounts valid skips and the claim that these skips are always a 'win scenario' without 'any tradeoff' is not supported. The paper needs a weight-aware criterion (for example, checking s_i - s_{i-1} + ln w_{i-1}) and a forward error bound for the skipped updates.
- [Section IV-B and Section V.A] The hardware comparison is not fully apples-to-apples because the baseline FlashAttention2 nonlinear units are not described. FLASH-D evaluates sigmoid and natural logarithm with 8-segment PWL approximations, while the paper does not state how the FlashAttention2 baseline implements exponential and division. If the baseline uses a different, more expensive or more accurate nonlinear function implementation, the reported 22.8% area and 20.3% power reductions partly reflect the cost of approximation rather than the benefit of the reformulation. Please disclose the baseline's exponential/division hardware, report its approximation error, and compare FLASH-D against an equally accurate baseline, or explicitly characterize the accuracy-cost trade-off.
- [Abstract and Section III-C] The abstract claims 'a reduction in computational cost without introducing numerical approximations to the FlashAttention kernel.' This is accurate only for the exact-arithmetic recurrence with ideal sigmoid and log functions; the implemented hardware uses 8-segment PWL approximations, and the static skip rule is itself an approximation. Please rephrase the claims to distinguish the exact algebraic equivalence from the approximate hardware implementation, and avoid stating that the hardware introduces no numerical approximations.
minor comments (4)
- [Section III-A] The text says 'From line 5 of Alg. 1' when referring to the output recursion, but in Algorithm 1 the output update is line 6; the cross-reference should be corrected.
- [Fig. 2] The caption says the four graphs correspond to four different values of w_{i-1}, but only w_{i-1} = 0.99 is identified in the text; the other three values should be given in the caption or legend.
- [Table I] The table formatting is confusing: entries such as 'Microsoft / 0.8%' combine the model vendor and the measured percentage without clear column separation, making the table hard to read.
- [Section V.A] The repository is mentioned as publicly available, but the manuscript does not provide a version or commit identifier; adding one would improve reproducibility.
Assumptions & free parameters
free parameters (3)
- PWL coefficients for sigmoid approximation (8 segments) =
not reported
- PWL coefficients for natural logarithm approximation (8 segments) =
not reported
- Sigmoid saturation thresholds [-6, 11] =
lower -6, upper 11
assumptions (4)
- standard math Exact real-number algebra is associative and commutative enough that the reordered computation in Alg. 3 equals Alg. 1.
- domain assumption 8-segment PWL approximations of sigmoid and log are accurate enough for LLM inference outputs.
- ad hoc to paper A score difference outside [-6, 11] guarantees sigmoid saturation regardless of the previous weight.
- domain assumption BF16 and FP8-E4M3 reduced-precision formats preserve the claimed equivalence.
Cite this review
Pith. "Pith review of FLASH-D: FlashAttention with Hidden Softmax Division." pith.science (2026). https://pith.science/paper/R3BCBVYC
@misc{pith2026250514201,
author = {Pith},
title = {Pith review of: FLASH-D: FlashAttention with Hidden Softmax Division},
year = {2026},
howpublished = {\url{https://pith.science/paper/R3BCBVYC}},
note = {Machine review of arXiv:2505.14201}
}
read the original abstract
The transformer's attention mechanism has revolutionized AI and machine learning, with its efficient computation being crucial to its performance. However, calculating attention involves matrix operations interspersed with softmax rescaling, which inherently slows down computation and requires processing the entire input sequence. Building on online softmax computation, FlashAttention integrates softmax calculation with matrix arithmetic, enabling tiled computation independent of sequence length. While optimized for GPUs, FlashAttention's simplicity makes it amenable to direct hardware acceleration. This work re-evaluates the core FlashAttention kernel, presenting FLASH-D a mathematically equivalent, yet simplified, formulation that achieves: (a) hiding softmax division within other non-linear function evaluations; (b) inherently numerically stable computation of exponentials, eliminating the need for maximum value subtraction; and (c) a reduction in computational cost without introducing numerical approximations to the FlashAttention kernel. Importantly, the essential FlashAttention properties that facilitate efficient tiled implementation are fully preserved. Hardware implementation results at 28nm demonstrate that this proposed formulation achieves a 22.8% reduction in area and a 20.3% reduction in power, on average, compared to state-of-the-art parallel hardware architectures without any performance penalty.
Figures
Reference graph
Works this paper leans on
-
[1]
An image is worth 16x16 words: Transformers for image recognition at scale,
A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly et al., “An image is worth 16x16 words: Transformers for image recognition at scale,” arXiv preprint arXiv:2010.11929 , 2020
arXiv 2010
-
[2]
Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning,
D. Guo, D. Yang, H. Zhang, J. Song, R. Zhang, R. Xu, Q. Zhu, S. Ma, P. Wang, X. Bi et al., “Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning,” arXiv preprint arXiv:2501.12948 , 2025
arXiv 2025
-
[3]
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin, “Attention is all you need,” in Intern. Conf. on Neural Information Processing Systems (NIPS) , 2017, p. 6000–6010
work page 2017
-
[4]
Longformer: The long- document transformer,
I. Beltagy, M. E. Peters, and A. Cohan, “Longformer: The long- document transformer,” arXiv preprint arXiv:2004.05150 , 2020
arXiv 2004
-
[5]
Generating long sequences with sparse transformers,
R. Child, S. Gray, A. Radford, and I. Sutskever, “Generating long sequences with sparse transformers,” arXiv preprint arXiv:1904.10509 , 2019
arXiv 1904
-
[6]
Transformers are rnns: Fast autoregressive transformers with linear attention,
A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret, “Transformers are rnns: Fast autoregressive transformers with linear attention,” in Intern. conference on machine learning , 2020, pp. 5156–5165
work page 2020
-
[7]
Low-rank Attention Side-Tuning for Parameter-Efficient Fine-Tuning
N. Tang, M. Fu, K. Zhu, and J. Wu, “Low-rank attention side-tuning for parameter-efficient fine-tuning,” arXiv preprint arXiv:2402.04009, 2024
work page Pith review arXiv 2024
-
[8]
A3: Accelerating attention mechanisms in neural networks with approximation,
T. J. Ham, S. J. Jung, S. Kim, Y . H. Oh, Y . Park, Y . Song, J.-H. Park, S. Lee, K. Park, J. W. Lee, and D.-K. Jeong, “A3: Accelerating attention mechanisms in neural networks with approximation,” in IEEE Intern. Symp. on High-Performance Computer Architecture (HPCA) , 2020, p. 328–341
work page 2020
Show all 34 references
-
[9]
A 95.6-TOPS/W deep learning inference accelerator with per-vector scaled 4-bit quantization in 5 nm,
B. Keller, R. Venkatesan, S. Dai, S. G. Tell, B. Zimmer, C. Sakr, W. J. Dally, C. T. Gray, and B. Khailany, “A 95.6-TOPS/W deep learning inference accelerator with per-vector scaled 4-bit quantization in 5 nm,” IEEE Journal of Solid-State Circuits, vol. 58, no. 4, p. 1129–1141, 2023
2023
-
[10]
Hardware accelerator for multi-head attention and position-wise feed-forward in the trans- former,
S. Lu, M. Wang, S. Liang, J. Lin, and Z. Wang, “Hardware accelerator for multi-head attention and position-wise feed-forward in the trans- former,” in IEEE Intern. System-on-Chip Conference (SOCC) , 2020, pp. 84–89
2020
-
[11]
Mnnfast: a fast and scalable system architecture for memory-augmented neural networks,
H. Jang, J. Kim, J.-E. Jo, J. Lee, and J. Kim, “Mnnfast: a fast and scalable system architecture for memory-augmented neural networks,” in Intern. Symp. on Computer Architecture (ISCA) , 2019, p. 250–263
2019
-
[12]
COSA plus: Enhanced co-operative systolic arrays for attention mechanism in transformers,
Z. Wang, G. Wang, and G. He, “COSA plus: Enhanced co-operative systolic arrays for attention mechanism in transformers,” IEEE Trans. on Computer-Aided Design of Integrated Circuits and Systems (TCAD) , vol. 44, no. 2, p. 723–736, 2025
2025
-
[13]
ELSA: Hardware-software co-design for efficient, lightweight self- attention mechanism in neural networks,
T. J. Ham, Y . Lee, S. H. Seo, S. Kim, H. Choi, S. J. Jung, and J. W. Lee, “ELSA: Hardware-software co-design for efficient, lightweight self- attention mechanism in neural networks,” in Intern. Symp. on Computer Architecture (ISCA), 2021, p. 692–705
2021
-
[14]
TSAcc: An efficient tempo-spatial similarity aware accelerator for attention acceleration,
Z. Song, C. Qi, Y . Yao, P. Zhou, Y . Zi, N. Wang, and X. Liang, “TSAcc: An efficient tempo-spatial similarity aware accelerator for attention acceleration,” in ACM/IEEE Design Automation Conference , 2024
2024
-
[15]
X-former: In-memory acceleration of transformers,
S. Sridharan, J. R. Stevens, K. Roy, and A. Raghunathan, “X-former: In-memory acceleration of transformers,” IEEE Transactions on Very Large Scale Integration (VLSI) Systems , vol. 31, no. 8, pp. 1223–1233, 2023
2023
-
[16]
Flashattention: Fast and memory-efficient exact attention with IO-awareness,
T. Dao, D. Fu, S. Ermon, A. Rudra, and C. R ´e, “Flashattention: Fast and memory-efficient exact attention with IO-awareness,” Advances in neural information processing systems, vol. 35, pp. 16 344–16 359, 2022
2022
-
[17]
Flashattention-2: Faster attention with better parallelism and work partitioning,
T. Dao, “Flashattention-2: Faster attention with better parallelism and work partitioning,” arXiv preprint arXiv:2307.08691 , 2023
2023 arXiv
-
[18]
Self-attention does not need O(n2) memory,
M. N. Rabe and C. Staats, “Self-attention does not need O(n2) memory,” arXiv preprint arXiv:2112.05682 , 2021
2021 arXiv
-
[19]
Hardware-efficient softmax approximation for self-attention networks,
N. A. Koca, A. T. Do, and C.-H. Chang, “Hardware-efficient softmax approximation for self-attention networks,” in Intern. Symp. on Circuits and Systems (ISCAS) , 2023, p. 1–5
2023
-
[20]
Language models are unsupervised multitask learners,
A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever et al., “Language models are unsupervised multitask learners,” OpenAI blog , p. 9, 2019
2019
-
[21]
Fastervit: Fast vision transformers with hierarchical attention,
A. Hatamizadeh, G. Heinrich, H. Yin, A. Tao, J. M. Alvarez, J. Kautz, and P. Molchanov, “Fastervit: Fast vision transformers with hierarchical attention,” in Intern. Conf. on Learning Representations (ICLR) , 2024
2024
-
[22]
Consmax: Hardware-friendly alternative softmax with learnable parameters,
S. Liu, G. Tao, Y . Zou, D. Chow, Z. Fan, K. Lei, B. Pan, D. Sylvester, G. Kielian, and M. Saligane, “Consmax: Hardware-friendly alternative softmax with learnable parameters,” arXiv preprint arXiv:2402.10930 , 2024
2024 arXiv
-
[23]
Online normalizer calculation for softmax,
M. Milakov and N. Gimelshein, “Online normalizer calculation for softmax,” arXiv preprint arXiv:1805.02867 , 2018
2018 arXiv
-
[24]
Edge- BERT: sentence-level energy optimizations for latency-aware multi-task NLP inference,
T. Tambe, C. Hooper, L. Pentecost, T. Jia, E.-Y . Yang, M. Donato, V . Sanh, P. Whatmough, A. M. Rush, D. Brooks, and G.-Y . Wei, “Edge- BERT: sentence-level energy optimizations for latency-aware multi-task NLP inference,” in Intern. Symposium on Microarchitecture (MICRO) , 2...
2021
-
[25]
Online alignment and ad- dition in multiterm floating-point adders,
K. Alexandridis and G. Dimitrakopoulos, “Online alignment and ad- dition in multiterm floating-point adders,” IEEE Transactions on Very Large Scale Integration (VLSI) Systems , vol. 33, no. 4, pp. 1182–1186, 2025
2025
-
[26]
Templatized fused vector floating-point dot product for high-level synthesis,
D. Filippas, C. Nicopoulos, and G. Dimitrakopoulos, “Templatized fused vector floating-point dot product for high-level synthesis,” Journal of Low Power Electronics and Applications , vol. 12, no. 4, p. 56, 2022
2022
-
[27]
SOLE: hardware- software co-design of softmax and layernorm for efficient transformer inference,
W. Wang, S. Zhou, W. Sun, P. Sun, and Y . Liu, “SOLE: hardware- software co-design of softmax and layernorm for efficient transformer inference,” in IEEE/ACM Intern. Conference on Computer Aided Design (ICCAD), 2023, p. 1–9
2023
-
[28]
A pseudo-softmax function for hardware-based high speed image classification,
G. C. Cardarilli, L. Di Nunzio, R. Fazzolari, D. Giardino, A. Nannarelli, M. Re, and S. Span `o, “A pseudo-softmax function for hardware-based high speed image classification,” Scientific Reports, vol. 11, 2021
2021
-
[29]
C. F. Jekel and G. Venter, pwlf: A Python Library for Fitting 1D Continuous Piecewise Linear Functions , 2019
2019
-
[30]
A study of bfloat16 for deep learning training,
D. Kalamkar, D. Mudigere, N. Mellempudi, D. Das, K. Banerjee, S. Avancha, D. T. V ooturi, N. Jammalamadaka, J. Huang, H. Yuen et al. , “A study of bfloat16 for deep learning training,” arXiv preprint arXiv:1905.12322, 2019
1905 arXiv
-
[31]
Fp8 formats for deep learning,
P. Micikevicius, D. Stosic, N. Burgess, M. Cornea, P. Dubey, R. Grisen- thwaite, S. Ha, A. Heinecke, P. Judd, J. Kamalu et al., “Fp8 formats for deep learning,” arXiv preprint arXiv:2209.05433 , 2022
2022 arXiv
-
[32]
llama2.c,
A. Karpathy, “llama2.c,” https://github.com/karpathy/llama2.c.git, 2023
2023
-
[33]
Transformers: State-of-the-art natural language processing,
T. Wolf, L. Debut, V . Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen, C. Ma, Y . Jernite, J. Plu, C. Xu, T. L. Scao, S. Gugger, M. Drame, Q. Lhoest, and A. M. Rush, “Transformers: State-of-the-art nat...
2020
-
[34]
Promptbench: A unified library for evaluation of large language models,
K. Zhu, Q. Zhao, H. Chen, J. Wang, and X. Xie, “Promptbench: A unified library for evaluation of large language models,” arXiv preprint arXiv:2312.07910, 2023
2023 arXiv
Reviewed August 7, 2026 · model on record in the stance chip above.
Discussion (0). Continue with ORCID to comment.