Pith. sign in

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 →

arxiv 2505.14201 v1 pith:R3BCBVYC submitted 2025-05-20 cs.LG cs.AIcs.AR

classification cs.LGcs.AIcs.AR
keywords flashattentioncomputationsoftmaxhardwarereductionachievesattentiondivision
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

Attention in transformers computes, for each query, a weighted average of value vectors, where the weights are softmax scores over all keys. FlashAttention computes these scores tile by tile using an online softmax, which keeps a running maximum and a running sum of exponentials. FLASH-D replaces that online-softmax bookkeeping with a direct recurrence for the weight of each new value vector. The weight at step i is a sigmoid of the difference between the current and previous attention scores, plus the logarithm of the previous weight. This is an exact algebraic identity, not an approximation: the running max and the running denominator are both absorbed into the sigmoid argument.

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.

Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

4 major / 4 minor

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)
  1. [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.
  2. [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.
  3. [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.
  4. [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)
  1. [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.
  2. [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.
  3. [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.
  4. [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 3 free parameters · 4 assumptions · 0 invented entities

The core mathematical derivation adds no physical entities and no fitted constants; it is a pure algebraic reformulation. The hardware implementation introduces PWL fit coefficients and hand-chosen saturation thresholds, which affect the secondary claims about accuracy and energy savings. The ad hoc assumption about the static range rule is the most fragile part of the presentation.

free parameters (3)
  • PWL coefficients for sigmoid approximation (8 segments) = not reported
    Fitted offline with the pwlf library; these coefficients determine the hardware area, power, and approximation error of the sigmoid in FLASH-D.
  • PWL coefficients for natural logarithm approximation (8 segments) = not reported
    Fitted offline with pwlf to compute ln w_{i-1}; same impact on hardware cost and numerical accuracy.
  • Sigmoid saturation thresholds [-6, 11] = lower -6, upper 11
    Hand-chosen range outside which weights are clamped to near 0 or 1; this is used for the output-update simplification and is not safe for all values of the previous weight.
assumptions (4)
  • standard math Exact real-number algebra is associative and commutative enough that the reordered computation in Alg. 3 equals Alg. 1.
    The derivation relies on standard algebraic manipulation of exponentials and fractions, which is valid in exact arithmetic.
  • domain assumption 8-segment PWL approximations of sigmoid and log are accurate enough for LLM inference outputs.
    The hardware implementation uses PWL approximations, but the paper provides no quantitative error analysis of the resulting attention outputs.
  • ad hoc to paper A score difference outside [-6, 11] guarantees sigmoid saturation regardless of the previous weight.
    Section III-C and Fig. 2 support this only for weights near 1; for tiny w_{i-1}, ln w_{i-1} shifts the sigmoid argument and the conclusion fails.
  • domain assumption BF16 and FP8-E4M3 reduced-precision formats preserve the claimed equivalence.
    The hardware measurements use reduced-precision floating point, but no error comparison between FLASH-D and FlashAttention2 in these formats is reported.

how reviews work

0 comments
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

Figures reproduced from arXiv: 2505.14201 by the authors.

Figure 1
Figure 1. A parallel hardware architecture for FlashAttention2 kernel for [PITH_FULL_IMAGE:figures/full_fig_p003_1.png] view at source ↗
Figure 2
Figure 2. Weight wi function for various values of consecutive attention score differences si − si−1. The four weight graphs correspond to four different values of the previous weight wi−1. To clarify this argument we plot in [PITH_FULL_IMAGE:figures/full_fig_p004_2.png] view at source ↗
Figure 3
Figure 3. A parallel hardware architecture for FLASH-D kernel for multiple [PITH_FULL_IMAGE:figures/full_fig_p005_3.png] view at source ↗
Figures from the paper (2 more)
Figure 4
Figure 4. Figure 4: The hardware area at 28 nm for FLASH-D and FlashAttention2 [PITH_FULL_IMAGE:figures/full_fig_p006_4.png]
Figure 5
Figure 5. Figure 5: The average power for FLASH-D and FlashAttention2 kernel for [PITH_FULL_IMAGE:figures/full_fig_p006_5.png]

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

34 extracted references · 22 canonical work pages

  1. [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

  2. [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

  3. [3]

    Attention is all you need,

    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

  4. [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

  5. [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

  6. [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

  7. [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

  8. [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

Show all 34 references
  1. [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

  2. [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

  3. [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

  4. [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

  5. [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

  6. [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

  7. [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

  8. [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

  9. [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

  10. [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

  11. [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

  12. [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

  13. [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

  14. [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

  15. [23]

    Online normalizer calculation for softmax,

    M. Milakov and N. Gimelshein, “Online normalizer calculation for softmax,” arXiv preprint arXiv:1805.02867 , 2018

  16. [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...

  17. [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

  18. [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

  19. [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

  20. [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

  21. [29]

    C. F. Jekel and G. Venter, pwlf: A Python Library for Fitting 1D Continuous Piecewise Linear Functions , 2019

  22. [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

  23. [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

  24. [32]

    llama2.c,

    A. Karpathy, “llama2.c,” https://github.com/karpathy/llama2.c.git, 2023

  25. [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...

  26. [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

Pith tools

Reviewed August 7, 2026 · model on record in the stance chip above.