Appendix A Extended Related Work A.1 MICA Comparison with Multivariate Transformer Architectures MICA differs from existing multivariate Transformer architectures in both mechanism and complexity. Moirai ( Woo et al., 2024 ) concatenates channel tokens and applies attention over the expanded token set without modifying the attention mechanism, incurring O ( P 2 C 2 ) \mathcal{O}(P^{2}C^{2}) O ( P 2 C 2 ) complexity that scales quadratically with channel count. MICA differs in that it directly modifies the attention backbone by augmenting scaled dot-product attention with a linear attention mechanism and mixing gate to model cross-channel dependencies. Timer-XL ( Liu et al., 2025 ) flattens the 2D multivariate structure into a 1D token sequence and applies attention over this expanded set, incurring O ( P 2 C 2 ) \mathcal{O}(P^{2}C^{2}) O ( P 2 C 2 ) complexity. Unlike Moirai ’s standard attention, Timer-XL employs a Kronecker product-based mask to preserve causal dependencies. In contrast, MICA avoids token expansion by compressing cross-channel information through linear attention and fusing it with temporal attention via a learnable gate, achieving O ( P 2 C + P C ) \mathcal{O}(P^{2}C+PC) O ( P 2 C + P C ) complexity. While iTransformer ( Liu et al., 2024 ) inverts the data representation to treat channels as tokens ( O ( P C 2 ) \mathcal{O}(PC^{2}) O ( P C 2 ) ), this approach eliminates temporal attention entirely. MICA preserves the standard temporal token structure while selectively integrating cross-channel information via gating. Crossformer ( Zhang & Yan, 2023 ) employs a two-stage attention design with dedicated cross-time and cross-dimension layers. MICA differs from Crossformer in several key aspects: 1) Execution: Crossformer ’s stages must be executed sequentially , where the output of the cross-time stage becomes the input to the cross-dimension stage, whereas MICA ’s local temporal attention and global cross-channel attention can operate in parallel . 2) Cross-channel compression: Crossformer introduces a router mechanism with a small fixed number of learnable vectors ( R ≪ C R\ll C R ≪ C ) that act as an information bottleneck. These router vectors compress cross-channel information by aggregating messages from all C C C channels, which are then queried by individual channels to retrieve the compressed representation. In contrast, MICA compresses cross-channel information directly via summation across the channel dimension to construct the memory matrix M \mathbf{M} M , which is subsequently queried. 3) Complexity scaling: Crossformer reduces complexity from O ( P C 2 ) \mathcal{O}(PC^{2}) O ( P C 2 ) to O ( P R C ) \mathcal{O}(PRC) O ( P R C ) only when R ≪ C R\ll C R ≪ C , requiring a constraint on the number of router vectors. MICA achieves O ( P 2 C + P C ) \mathcal{O}(P^{2}C+PC) O ( P 2 C + P C ) complexity while processing all C C C channels directly without requiring R ≪ C R\ll C R ≪ C for efficiency. 4) Information fusion: Crossformer combines temporal and cross-dimension attention outputs through addition, whereas MICA employs a learnable gate to compute an adaptive weighted combination of temporal and cross-dimension attention outputs. Chronos-2 ( Ansari et al., 2025 ) alternates between time and group attention layers that operate sequentially within each transformer block, achieving O ( P 2 C + P C 2 ) \mathcal{O}(P^{2}C+PC^{2}) O ( P 2 C + P C 2 ) complexity. The time layer aggregates across patches within individual series, while the group layer aggregates across all series at each patch index. In contrast, MICA computes local temporal and global cross-channel attention as independent operations that can be parallelized. MICA ’s cross-channel mechanism architecturally differs from group attention: Rather than group attention’s full attention over all C C C channels, MICA employs linear attention to compress cross-channel information via summation, achieving O ( P 2 C + P C ) \mathcal{O}(P^{2}C+PC) O ( P 2 C + P C ) complexity. MICA uses a learnable gate to adaptively fuse local and global representations rather than fixed alternation. MICA ’s approach shares conceptual similarities with SOFTS ( Han et al., 2024 ) , which also employs gating to fuse local (channel-specific) and global (cross-channel) representations. However, SOFTS operates within an MLP-based architecture and applies channel mixing at the encoder level, whereas MICA integrates cross-channel interaction directly within the transformer attention mechanism. Critically, SOFTS uses an MLP gate to combine local and global representations derived from separate processing pathways, while MICA ’s channel-aware gates compress and modulate cross-channel information in a single pass without repeated concatenation operations. A.2 Extended Multivariate Architectures Convolution-based models. ModernTCN ( Luo & Wang, 2024 ) captures cross-channel dependencies through depthwise-separable convolutions, which decouple spatial and channel-wise filtering to achieve linear complexity with channel count and sequences length O ( P C ) \mathcal{O}(PC) O ( P C ) . TimesNet ( Wu et al., 2023 ) takes a different approach by reshaping 1D time series into 2D tensors and applying 2D convolutions to model both intra- and inter-period variation with complexity of O ( P log ( P ) C ) \mathcal{O}(P\log(P)C) O ( P log ( P ) C ) . State-Space models. We note that TimePro ( Ma et al., 2025 ) reports a complexity of O ( P C ) \mathcal{O}(PC) O ( P C ) for the core HyperMamba component; however, the full model includes two linear projections with complexity O ( P 2 C ) \mathcal{O}(P^{2}C) O ( P 2 C ) , which is quadratic in the number of patches P P P . As P P P scales with sequence length, the overall complexity is more precisely O ( P 2 C ) \mathcal{O}(P^{2}C) O ( P 2 C ) when considering all components. Graph-based models. StemGNN ( Cao et al., 2020 ) models cross-channel dependencies by applying Graph Fourier Transform (GFT) jointly with Discrete Fourier Transform (DFT) to capture inter-series correlations and temporal dependencies together in the spectral domain; however, the eigendecomposition required by GFT incurs cubic cost in channel count, yielding an overall complexity of O ( C 3 + P C 2 + P l o g ( P ) C ) \mathcal{O}(C^{3}+PC^{2}+Plog(P)C) O ( C 3 + P C 2 + P l o g ( P ) C ) . TimeFilter ( Hu et al., 2025 ) constructs a patch-specific spatial-temporal graph that models temporal, inter-channel, and cross-channel dependencies while filtering out irrelevant edges per patch achieving a complexity of O ( P 2 C + P C 2 ) \mathcal{O}(P^{2}C+PC^{2}) O ( P 2 C + P C 2 ) , which is quadratic in both P P P and C C C . A.3 Efficient Attention Architectures Building on the development of linear attention ( Katharopoulos et al., 2020 ) , efficient attention mechanisms have been proposed to address the quadratic complexity of standard attention with respect to sequence length in sequence modeling. Performers ( Choromanski et al., 2021 ) construct an unbiased approximation of the softmax kernel using exponential random feature maps with orthogonal random vector sampling, ensuring non-negative attention weights and achieving O ( L R ) \mathcal{O}(LR) O ( L R ) complexity, where R R R is the number of random features. Linear scaling over sequence length L L L holds only when R ≪ L R\ll L R ≪ L , and while increasing R R R reduces approximation variance, it reintroduces quadratic-like costs as R → L R\rightarrow L R → L . Random Feature Attention ( RFA ) ( Peng et al., 2021 ) introduces a variation using sinusoidal random Fourier features ( Rahimi & Recht, 2007 ) with i.i.d. random vector sampling to construct an unbiased approximation of the softmax kernel, retaining the same O ( L R ) \mathcal{O}(LR) O ( L R ) complexity and R ≪ L R\ll L R ≪ L requirement. In place of Performers ’ orthogonal sampling, RFA instead reduces approximation variance via normalization of queries and keys. Performers ( Choromanski et al., 2021 ) address this by instead using positive orthogonal random features specifically designed to yield an unbiased and lower-variance approximation of the softmax kernel, achieving O ( L R ) \mathcal{O}(LR) O ( L R ) complexity where L L L is the context length and R R R is the number of random features. While this scales linearly in L L L when R R R is fixed, R R R introduces an additional hyperparameter and linear scaling only holds when R ≪ L R\ll L R ≪ L . cosFormer ( Qin et al., 2022 ) instead replaces softmax with a ReLU activation to ensure non-negativity, combined with a cosine-based re-weighting mechanism to recover the locality bias of softmax, achieving true O ( L ) \mathcal{O}(L) O ( L ) complexity without additional hyperparameters. However, adapting cosFormer for cross-channel attention in multivariate forecasting is not well-motivated: the cosine re-weighting encodes a notion of token proximity that is meaningful along the time dimension but meaningless across channels, which have no inherent ordering. RetNet ( Sun et al., 2023 ) processes tokens sequentially via a recurrent state update with an exponential decay factor that causes earlier tokens to contribute less to later ones, however this efficiency of O ( 1 ) \mathcal{O}(1) O ( 1 ) constant complexity with respect to sequence length L L L is exclusive to inference as training retains the same O ( L 2 ) \mathcal{O}(L^{2}) O ( L 2 ) quadratic complexity as standard attention. Similar to cosFormer , adapting this mechanism for cross-channel aggregation in multivariate forecasting is not well-motivated: channels have no inherent ordering, yet would be processed arbitrarily and systematically downweighted by the decay factor, imposing spurious structure on channels that have no inherent ordering. In contrast, Infini-Attention ( Munkhdalai et al., 2024 ) aggregates tokens via summation, which is order-invariant and achieves true O ( L ) \mathcal{O}(L) O ( L ) complexity without approximation error or additional hyperparameters beyond the gate parameter β \beta β . These properties make Infini-Attention a natural candidate for cross-channel aggregation in multivariate time series, motivating our adaptation of this architectural design for cross-channel mixing and compression. Additionally, dedicated channel-independent (CI) forecasting models have been developed to address the quadratic complexity of attention mechanisms, including Informer ( Zhou et al., 2021 ) , which achieves O ( L log L ) \mathcal{O}(L\log L) O ( L log L ) complexity with input length L L L ; Autoformer ( Wu et al., 2021 ) , which achieves O ( L log L ) \mathcal{O}(L\log L) O ( L log L ) complexity; Fedformer ( Zhou et al., 2022 ) , which achieves O ( L ) \mathcal{O}(L) O ( L ) complexity; and PatchTST ( Nie et al., 2023 ) , which achieves O ( P 2 ) \mathcal{O}(P^{2}) O ( P 2 ) complexity by segmenting the input into fixed-length patches of length P P P . While the MICA design is compatible with any channel-independent Transformer architectures, we select PatchTST as our base model given it outperforms Informer , Autoformer , and Fedformer on standard benchmarks ( Nie et al., 2023 ) and patch-based tokenization has become the dominant paradigm in modern Time Series Foundation Models, including Chronos-2 , TimesFM , Moirai , Timer-XL , and MOMENT . A.4 Infini-Attention Infini-Attention was originally proposed to address the long-context bottleneck for transformers with text data, given that processing is quadratic in sequence length ( Munkhdalai et al., 2024 ) . This attention mechanism extends the standard transformer architecture by incorporating a compressive memory to capture information from infinitely long contexts in an approach that is bounded in memory and computation. The key innovation is maintaining a separate compressive memory state that accumulates information from tokens in past context windows in addition to standard attention computed for the current context window. For each Transformer layer, Infini-Attention computes scaled dot-product attention for the current context window, i i i , A local = softmax ( Q K ⊤ d k ) V . \displaystyle\textbf{A}_{\text{local}}=\text{softmax}\left(\frac{\textbf{Q}\textbf{K}^{\top}}{\sqrt{d_{k}}}\right)\textbf{V}. A local = softmax ( d k Q K ⊤ ) V . (12) It also maintains a memory matrix M ∈ R B × N × d k × d v \mathbf{M}\in\mathbb{R}^{B\times N\times d_{k}\times d_{v}} M ∈ R B × N × d k × d v and a normalization vector z i ∈ R B × N × d k × 1 \mathbf{z}_{i}\in\mathbb{R}^{B\times N\times d_{k}\times 1} z i ∈ R B × N × d k × 1 , where B B B is the window batch size, N N N is the number of attention heads, d k d_{k} d k is the per-head hidden dimension of the key vector, and d v d_{v} d v is the per-head hidden dimension of the value vector. Initially, M \mathbf{M} M is initialized with zeros. Given query, key, and value matrices Q , K , V \mathbf{Q},\mathbf{K},\mathbf{V} Q , K , V for the current context window, i i i , the memory is updated using a linear attention mechanism: M i \displaystyle\mathbf{M}_{i} M i ← M i − 1 + ϕ ( K ) ⊤ V , \displaystyle\leftarrow\mathbf{M}_{i-1}+\phi(\mathbf{K})^{\top}\mathbf{V}, ← M i − 1 + ϕ ( K ) ⊤ V , (13) z i \displaystyle\mathbf{z}_{i} z i ← z i − 1 + ∑ l = 1 L ϕ ( K l ) , \displaystyle\leftarrow\mathbf{z}_{i-1}+\sum^{L}_{l=1}\phi(\mathbf{K}_{l}), ← z i − 1 + l = 1 ∑ L ϕ ( K l ) , (14) where the outer product K ⊤ V \mathbf{K}^{\top}\mathbf{V} K ⊤ V compresses key-value associations into the memory matrix, and z \mathbf{z} z tracks the normalization term for proper retrieval. Here, ϕ \phi ϕ corresponds to the nonlinear activation ELU ( X ) + 1 \text{ELU}(\mathbf{X})+1 ELU ( X ) + 1 . To retrieve information from the compressive memory, the query vectors attend to the memory matrix: A global = ϕ ( Q ) M i − 1 ϕ ( Q ) z i − 1 , \mathbf{A}_{\text{global}}=\frac{\phi(\mathbf{Q})\mathbf{M}_{i-1}}{\phi(\mathbf{Q})\mathbf{z}_{i-1}}, A global = ϕ ( Q ) z i − 1 ϕ ( Q ) M i − 1 , (15) where element-wise division by ϕ ( Q ) z i − 1 \phi(\mathbf{Q})\mathbf{z}_{i-1} ϕ ( Q ) z i − 1 normalizes the retrieved values. This linear attention mechanism enables efficient retrieval with O ( P d k d v ) \mathcal{O}(Pd_{k}d_{v}) O ( P d k d v ) complexity. A learned linear gating mechanism is used to control the influence of both local attention output and global attention output, A mixed = σ ( β ) ⊙ A global + ( 1 − σ ( β ) ) ⊙ A local , \mathbf{A}_{\text{mixed}}=\sigma(\beta)\odot\mathbf{A}_{\text{global}}+(1-\sigma(\beta))\odot\mathbf{A}_{\text{local}}, A mixed = σ ( β ) ⊙ A global + ( 1 − σ ( β )) ⊙ A local , (16) where β \beta β is a learned scalar parameter, σ \sigma σ is the sigmoid function, and ⊙ \odot ⊙ denotes element-wise multiplication. This gating allows the model to dynamically balance between attending to recent local context and distant global context stored in memory. Appendix B Dataset Details Table 3: Summary of forecasting datasets used in our empirical study. Dataset Domain Frequency # Unique Series IDs # Targets Min Length Max Length Horizon Val Size Test Size COVID Deaths ( Godahewa et al., 2021 ) Healthcare Daily 266 1 212 212 30 30 30 Simglucose ( Xie, 2018 ) Healthcare 5min 30 1 25920 25920 6 2592 2592 Iowa IHOP SMEX02 Windspeed ( Iowa State University, 2008 ; Iowa State University, 2005 ) Weather 5min 34 1 23040 23040 24 2304 2304 Iowa PLOWS Windspeed ( Iowa State University, 2010 ) Weather 5min 41 1 37440 37440 24 3744 3744 Jena Weather ( Aksu et al., 2024 ; Wu et al., 2021 ) Weather Hourly 1 21 8784 8784 48 48 912 Daily 1 21 366 366 30 30 60 M-DENSE ( Aksu et al., 2024 ; de Medrano & Aznarte, 2020 ) Transportation Hourly 30 1 17520 17520 48 48 960 Daily 30 1 730 730 30 30 90 Loop-Seattle ( Cui et al., 2018 ; Cui et al., 2019 ) Transportation Daily 323 1 365 365 30 30 60 ETT1 ( Aksu et al., 2024 ; Zhou et al., 2021 ) Energy Hourly 1 7 17420 17420 48 48 960 Daily 1 7 725 725 30 30 90 Weekly 1 7 103 103 8 8 16 ETT2 ( Aksu et al., 2024 ; Zhou et al., 2021 ) Energy Hourly 1 7 17420 17420 48 48 960 Daily 1 7 725 725 30 30 90 Weekly 1 7 103 103 8 8 16 Solar ( Lai et al., 2018 ) Weather Hourly 137 1 8760 8760 48 48 912 Daily 137 1 365 365 30 30 60 Weekly 137 1 52 52 8 8 8 Simglucose. The Simglucose dataset is from an open-source Python implementation of the FDA-approved UVa/Padova Simulator (2008 version), which is contained in the ‘simglucose’ GitHub repository ( Xie, 2018 ; Visentin et al., 2016 ) . It provides clinical and biological parameters for 30 patients (10 adults, 10 adolescents, and 10 children) with Type-I diabetes. Clinical parameters include information on age and weight, and biological parameters include information on insulin-glucose kinetics, such as absorption constants. We generate 90 days of data for each patient. In a similar approach to ( Potosnak et al., 2025a ) , the meal schedule used to generate simulated data was based on the Harrison-Benedict equation ( Arthur Harris & Gano Benedict, 21919 ) with approximately 3 meals per day and no additional snacks. Height data was not provided in the patient parameters, and so was estimated at 140cm, 170cm, and 175cm for adults, adolescents, and children, respectively ( Potosnak et al., 2025a ) . The ‘Dexcom’ continuous glucose monitor (CGM) option was used to obtain glucose readings at 5-minute intervals. We hypothesize that the shared system of ODEs governing insulin-glucose dynamics across patients may enable cross-patient multivariate modeling to capture common physiological patterns, potentially improving forecasting performance across patients. Iowa Windspeed. The IHOP and SMEX02 datasets contain 1-minute resolution surface meteorological data from 34 Automated Weather Observing System (AWOS) stations in Iowa from May 13, 2002 through June 25, 2002 and June 1, 2002 through July 31, 2002, respectively ( Iowa State University, 2008 ; Iowa State University, 2005 ) . We combine these datasets given their overlapping stations and time periods. The Iowa PLOWS dataset contains 1-minute resolution data from 41 AWOS stations from November 1, 2009 to March 10, 2010 ( Iowa State University, 2010 ) and is treated separately due to different stations and time period. Both datasets are downsampled to 5-minute resolution for wind speed forecasting. These open-source datasets provide well-suited test cases for multivariate modeling across geographically distributed stations, where accurate wind forecasting requires capturing cross-channel spatial dependencies from upwind measurement stations. Gift-Eval. We include several datasets from the Gift-Eval repository including: ETT1, ETT2, Jena Weather, COVID Deaths, Loop-Seattle, Solar, and M-Dense across different temporal resolutions ( Aksu et al., 2024 ) . We exclude Loop-Seattle (hourly) due to memory issues for baselines, such as Timer-XL , given the 80GB memory constraint of the A100 GPU. Our dataset selection is guided by the goal of representing several datasets across diverse domains of energy, healthcare, weather, and transportation while satisfying one of two key multivariate settings: (1) modeling multiple variates for a single entity (e.g., air temperature, atmospheric pressure, humidity, and wind direction at a single weather station such as in the case with the Jena Weather dataset), or (2) modeling a single variate across multiple entities (e.g., wind speed measurements across different weather stations such as in the case with the Iowa datasets). Following the GIFT-Eval protocol, we construct our test sets based on the product of the forecast horizon and the number of evaluation windows specified by the benchmark. To ensure efficient hyperparameter tuning, we set the validation set size equal to the forecast horizon for each dataset.
Appendix C Training Methodology and Hyperparameters In this Appendix section, we expand on the training methodology outlined in Section 4 . We train both MLP -based and Transformer-based models, maintaining consistent hyperparameters where applicable to enable controlled comparison of channel mixing strategies across architectures. Table 4 details the hyperparameters used for training configuration, common Transformer-based model hyperparameters, and common MLP -based model hyperparameters. Parameters not specified in the table are set to the defaults of the original implementations in the NeuralForecast library ( Olivares et al., 2022b ) . For MICA models incorporating an MLP mixing gate ( MOMENT and PatchTST ), we perform hyperparameter tuning using Optuna with 5 trials per model. The optimal hyperparameters are selected based on minimum validation loss, computed using the mean absolute error (MAE) loss function defined in Equation 2 . Table 4: Hyperparameters Hyperparameter Values Training configuration Single GPU SGD Batch Size * N N N channels Windows batch size 64 Maximum training steps S m a x S_{max} S ma x 12,000 Validation check steps 500 Initial learning rate 1e-3 Learning rate decay 0.5 Learning rate step size 4,000 Early stop patience steps 20 Random seed {1, 2, 3, 4, 5} Common Transformer-based model settings Input size 2 ∗ H 2*H 2 ∗ H Number of encoder layers 4 Hidden size 256 Number of Heads 4 Feed Forward hidden size 1024 Dropout 0.0 Projection/Head dropout 0.0 d k d_{k} d k 32 d v d_{v} d v 32 Scaler type Standard RevIN yes, without affine (standardization) Patch length (if applicable) 8 Stride (if applicable) 8 Positional Encoding (if applicable) sincos Common MLP -based model settings Input size 2 ∗ H 2*H 2 ∗ H Number of blocks 4 Hidden size 256 Dropout 0.0 Scaler type Standard RevIN yes, without affine (standardization) MICA MLP Gate Input size 2 ∗ H 2*H 2 ∗ H Hidden size {128, 256, 384, 512} Number of layers {2, 3, 4} Dropout rate {0.0, 0.1, 0.2} Appendix E A Memory-Efficient Enhancement to MICA: Algorithm Sketch We present an algorithmic sketch for a hardware-aware version of MICA designed to handle high-dimensional multivariate settings and long temporal contexts. The main idea lies in the fusion of local and global attention mechanisms into a specialized two-pass kernel. In the first pass ( Global Reduction ), a global compressive memory state M all \mathbf{M}_{\text{all}} M all is computed by aggregating keys and values across all channels in a single linear sweep with complexity O ( P C ) \mathcal{O}(PC) O ( P C ) . In the second pass ( Fused Channel-Wise Attention ), each channel is processed via a theoretical IO-aware kernel. Within this stage, local temporal patterns are computed through a tiled FlashAttention mechanism, while a local linear correction term for strict channel exclusion is simultaneously determined. By reusing the loaded keys and values from SRAM for both operations, redundant HBM access is eliminated, and the linear correction is computed with negligible overhead. Finally, an on-chip convex combination of the local and global outputs is performed using the learned mixing gate β \beta β , whereby only the final representations are written to memory. Through this sketched approach, spatial and temporal complexity are decoupled, providing a path for linear scaling with channel count C C C and optimized quadratic scaling with sequence length L L L . The following algorithm pertain to the channel inclusion MICA variant. Algorithm 1 Optimized Fused MICA (Flash-Local + Linear-Global) 0: Q , K , V ∈ R B × C × N × P × d \mathbf{Q,K,V}\in\mathbb{R}^{B\times C\times N\times P\times d} Q , K , V ∈ R B × C × N × P × d , mixing gate β \beta β 0: Y ∈ R B × C × N × P × d \mathbf{Y}\in\mathbb{R}^{B\times C\times N\times P\times d} Y ∈ R B × C × N × P × d 1: Hyperparameters: Block sizes B r , B c B_{r},B_{c} B r , B c for tiling; ϕ ( X ) = ELU ( X ) + 1 \phi(\mathbf{X})=\text{ELU}(\mathbf{X})+1 ϕ ( X ) = ELU ( X ) + 1 . 2: Pass 1: Global Memory Reduction (Kernel 1) 3: for batch b b b , head n n n in parallel do 4: Initialize M all ← 0 d × d , z all ← 0 d \mathbf{M}_{\text{all}}\leftarrow\mathbf{0}^{d\times d},\mathbf{z}_{\text{all}}\leftarrow\mathbf{0}^{d} M all ← 0 d × d , z all ← 0 d 5: for channel c = 1 c=1 c = 1 to C C C do 6: Load K ( c ) , V ( c ) \mathbf{K}^{(c)},\mathbf{V}^{(c)} K ( c ) , V ( c ) from HBM 7: K ′ ← ϕ ( K ( c ) ) \mathbf{K}^{\prime}\leftarrow\phi(\mathbf{K}^{(c)}) K ′ ← ϕ ( K ( c ) ) 8: M all ← M all + ( K ′ ) ⊤ V ( c ) \mathbf{M}_{\text{all}}\leftarrow\mathbf{M}_{\text{all}}+(\mathbf{K}^{\prime})^{\top}\mathbf{V}^{(c)} M all ← M all + ( K ′ ) ⊤ V ( c ) {Global State Accumulation} 9: z all ← z all + ∑ p K p ′ \mathbf{z}_{\text{all}}\leftarrow\mathbf{z}_{\text{all}}+\sum_{p}\mathbf{K}^{\prime}_{p} z all ← z all + ∑ p K p ′ 10: end for 11: Store M all , z all \mathbf{M}_{\text{all}},\mathbf{z}_{\text{all}} M all , z all to HBM 12: end for 13: Pass 2: Fused Channel-Wise Attention (Kernel 2) 14: for channel c = 1 c=1 c = 1 to C C C in parallel do 15: Load M all , z all \mathbf{M}_{\text{all}},\mathbf{z}_{\text{all}} M all , z all into SRAM 16: for query block i = 1 i=1 i = 1 to P / B r P/B_{r} P / B r do 17: Initialize Local Acc: O quad ( i ) ← 0 , ℓ ( i ) ← 0 , m ( i ) ← − ∞ \mathbf{O}_{\text{quad}}^{(i)}\leftarrow\mathbf{0},\ell^{(i)}\leftarrow\mathbf{0},m^{(i)}\leftarrow-\infty O quad ( i ) ← 0 , ℓ ( i ) ← 0 , m ( i ) ← − ∞ 18: for key-value block j = 1 j=1 j = 1 to P / B c P/B_{c} P / B c do 19: Load K [ j ] ( c ) , V [ j ] ( c ) \mathbf{K}^{(c)}_{[j]},\mathbf{V}^{(c)}_{[j]} K [ j ] ( c ) , V [ j ] ( c ) from channel c c c to SRAM 20: // Quadratic Local Attention (FlashAttention Logic) 21: Compute S i j = Q [ i ] ( c ) ( K [ j ] ( c ) ) ⊤ ⋅ scale \mathbf{S}_{ij}=\mathbf{Q}^{(c)}_{[i]}(\mathbf{K}^{(c)}_{[j]})^{\top}\cdot\text{scale} S ij = Q [ i ] ( c ) ( K [ j ] ( c ) ) ⊤ ⋅ scale 22: Update softmax stats m ( i ) , ℓ ( i ) m^{(i)},\ell^{(i)} m ( i ) , ℓ ( i ) and O quad ( i ) \mathbf{O}_{\text{quad}}^{(i)} O quad ( i ) using online softmax reduction 23: end for 24: O global ( i ) = ( ϕ ( Q [ i ] ( c ) ) M all ) / ( ϕ ( Q [ i ] ( c ) ) z all + ϵ ) \mathbf{O}_{\text{global}}^{(i)}=(\phi(\mathbf{Q}^{(c)}_{[i]})\mathbf{M}_{\text{all}})/(\phi(\mathbf{Q}^{(c)}_{[i]})\mathbf{z}_{\text{all}}+\epsilon) O global ( i ) = ( ϕ ( Q [ i ] ( c ) ) M all ) / ( ϕ ( Q [ i ] ( c ) ) z all + ϵ ) 25: Mixing: Y [ i ] ( c ) = σ ( β ) ⋅ O global ( i ) + ( 1 − σ ( β ) ) ⋅ O quad ( i ) \mathbf{Y}^{(c)}_{[i]}=\sigma(\beta)\cdot\mathbf{O}_{\text{global}}^{(i)}+(1-\sigma(\beta))\cdot\mathbf{O}_{\text{quad}}^{(i)} Y [ i ] ( c ) = σ ( β ) ⋅ O global ( i ) + ( 1 − σ ( β )) ⋅ O quad ( i ) 26: Write Y [ i ] ( c ) \mathbf{Y}^{(c)}_{[i]} Y [ i ] ( c ) to HBM 27: end for 28: end for Online Softmax Reduction: The update in line X maintains running statistics for numerically stable softmax computation across blocks. Specifically, for each block j j j , update: m new = max ( m old , max ( S i j ) ) m^{\text{new}}=\max(m^{\text{old}},\max(\mathbf{S}_{ij})) m new = max ( m old , max ( S ij )) and ℓ new = exp ( m old − m new ) ⋅ ℓ old + ∑ exp ( S i j − m new ) \ell^{\text{new}}=\exp(m^{\text{old}}-m^{\text{new}})\cdot\ell^{\text{old}}+\sum\exp(\mathbf{S}_{ij}-m^{\text{new}}) ℓ new = exp ( m old − m new ) ⋅ ℓ old + ∑ exp ( S ij − m new ) , following the FlashAttention approach ( Dao et al., 2022 ) .