Publications

1. 2026

   FoldAttention: Declared-Reference Softmax for Fast Decode and Deterministic Backward

   Sriman Achanta

   arXiv preprint

   • arXiv
   • PDF
   • Code

   Abstract

   Autoregressive decode repeatedly streams a growing KV cache, making attention a major cost at long context. Existing high-performance kernels use online softmax, which discovers a row's normalization reference as it scans keys. Earlier contributions therefore remain provisional and may require rescaling. We argue that the reference need not be discovered: softmax is invariant to a common shift, so the reference only has to keep the weights in range. We present FoldAttention, an additive formulation of softmax attention that fixes a finite reference $Z_i$ before scanning the KV cache. Each weight $2^{s_{ij}-Z_i}$ is then final when computed, so contributions add across disjoint key ranges and their quotient equals softmax attention in real arithmetic. We use this property to develop two techniques for Hopper decode: (1) final weights gate key and value reads before the bytes are fetched, and a per-call depth $T$ cuts keys below $2^{-T}$ while keeping their mass, and (2) additive partials compose split KV and shared-prefix cascades without rescaling. On H100 at $T=16$, FoldAttention decodes seven real-model generations 1.36-2.30$\times$ faster than the fastest BF16 baseline, and up to 3.09$\times$ faster across MHA and GQA shapes, at an error within 1.5% of the lowest BF16 error on six of the seven; reading every key, it is 1.14-1.30$\times$ faster at matched error. We validate on Qwen3-8B that a whole decode step is up to 1.46$\times$ faster while likelihood and long-context accuracy match those under BF16 kernels. The same principle makes the backward deterministic: CTAs round bounded partial gradients onto an integer grid declared before the reduction and add them in any order. FoldAttention thereby removes the determinism tax: its deterministic backward is up to 1.84$\times$ faster than deterministic FlashAttention-3/4 and 1.05$\times$ faster than the fastest nondeterministic kernel.

2. 2025

   Brain tumor segmentation based on deep learning, attention mechanisms, and energy-based uncertainty predictions

   Zachary Schwehr, Sriman Achanta

   Multimedia Tools and Applications, vol. 84, pp. 34229-34248

   • DOI
   • arXiv
   • Code

   Abstract

   Brain tumors are one of the deadliest forms of cancer with a mortality rate of over 80%. A quick and accurate diagnosis is crucial for increasing the chances of survival. However, in medical analysis, the manual annotation and segmentation of brain tumors are complicated. Multiple MRI modalities are typically analyzed as they provide unique information about the tumor regions. Although using several MRI modalities helps segment brain tumors, they tend to increase overfitting and computation in computer models. This paper proposed a region of interest detection algorithm that was implemented during data preprocessing to locate salient features and remove extraneous MRI data. This decreased the input size of the model, allowing for more aggressive data augmentations and deeper neural networks. Following the preprocessing of the MRI modalities, a fully convolutional autoencoder segmented the different brain MRIs using channel-wise attention and attention gates. Subsequently, test-time augmentations and an energy-based model were used for voxel-based uncertainty predictions. Experimentation was conducted on the BraTS 2019, 2020, and 2021 benchmarks, and the proposed models achieved a state-of-the-art segmentation performance with mean dice scores of 84.55, 88.52, and 90.82 on each respective dataset. Additionally, qualitative results were used to assess the segmentation models and uncertainty predictions. Extensive experimentation showed that the proposed framework with cropping and attention mechanisms achieved state-of-the-art performance when compared to its contemporaries.