---
title: Blockwise Flash Kernel for TFLA
url: https://www.emergentmind.com/topics/blockwise-flash-kernel
type: topic
---

# Blockwise Flash Kernel for TFLA

Tiled Flash Linear Attention (TFLA), commonly referred to as the Blockwise Flash Kernel, is a two-level, sequence-parallel GPU kernel algorithm engineered for efficient computation in linear recurrent neural networks (RNNs) such as matrix-memory LSTMs (mLSTM/xLSTM). TFLA extends the chunkwise-parallel principles of Flash Linear Attention (FLA) by introducing an additional level of fine-grained tiling within each chunk, addressing bottlenecks in arithmetic intensity and memory I/O that hamper the scaling of long-context sequence models. This kernel enables both high arithmetic intensity and arbitrary large chunk sizes, resulting in significant performance improvements over prior state-of-the-art kernels, including Flash Attention, Linear Attention, and Mamba, on modern accelerators [2503.14376].

## 1. Mathematical Formulation

TFLA generalizes to any linear RNN with gating, but is illustrated concretely for the mLSTM/xLSTM cell. At each timestep $t$, the cell maintains:
- Hidden state $h_t \in \mathbb{R}^d$
- Matrix memory $M_t \in \mathbb{R}^{d_q \times d}$
- Per-timestep scalar gates: $i_t$ (input), $f_t$ (forget), with optional output gate $o_t$

For the standard "mLSTMexp" recurrence with exponential input gate, the update equations are:
\[
m_t = \max\left(\log \sigma(\tilde f_t) + m_{t-1},\; \tilde i_t\right)
\]
\[
f_t = \exp\left(\log \sigma(\tilde f_t) + m_{t-1} - m_t\right), \quad i_t = \exp\left(\tilde i_t - m_t\right)
\]
\[
M_t = f_t \cdot M_{t-1} + i_t \cdot (k_t v_t^\top)
\]
\[
\tilde h_t = M_t^\top (q_t/\sqrt{d_q}), \quad h_t = \sigma(\tilde o_t) \odot \mathrm{NORM}(\tilde h_t)
\]
where $\sigma$ denotes the sigmoid, and $\mathrm{NORM}(\cdot)$ is either RMS- or Layer-Norm. The running max-state $m_t$ stabilizes the exp-input gate, ensuring numerical safety analogous to the softmax.

A simplified variant, "mLSTMsig," replaces the unbounded exponential gates with bounded sigmoidal gates,
\[
f_t = \sigma(\tilde f_t), \qquad i_t = \sigma(\tilde i_t)
\]
\[
M_t = f_t \cdot M_{t-1} + i_t \cdot (k_t v_t^\top)
\]
\[
h_t = \sigma(\tilde o_t) \odot \mathrm{NORM}(M_t^\top (q_t/\sqrt{d_q}))
\]
where, by construction, no extra max-state or normalizer is required to prevent overflow.

## 2. Blockwise Tiling and Kernel Workflow

The TFLA kernel partitions the input sequence of length $T$ into $N_c = \lceil T/C \rceil$ chunks, each of length $C$. Within each chunk, the $C \times C$ attention-style matrix is further divided into $(C/T)^2$ tiles of size $T \times T$. The procedure in each chunk comprises:

1. **Recurrent Pass:** Materializing the recurrent state $M_{k-1}$ from the previous chunk.
2. **Parallel Tiled Computation:** For every $T \times T$ tile within the chunk, the kernel accumulates $QK^{\top}$, constructs cumulative-forget and input exponents, applies a numerically-safe softmax in blockwise fashion, fuses the result, and performs the final matmul with the value matrix block $V$.
3. **Output Rescaling:** Intra-chunk outputs are rescaled to align with the scale of the inter-chunk contribution.
4. **Inter-chunk Contribution:** Outer-chunk state contributes via $H_\mathrm{inter} = \overline{Q} \cdot M_{k-1}$.
5. **Output Combination:** The final chunk output is a weighted sum of intra- and inter-chunk outputs, normalized for stability.

Pseudocode implementing this tiling and fusion appears explicitly in the kernel’s specification [2503.14376].

## 3. GPU Implementation Details

TFLA exploits the modern GPU architecture using a three-dimensional thread-block grid:
- $N_\mathrm{batch} \times N_\mathrm{head} \times N_\mathrm{chunk}$ for coarse parallelism
- $C/T$ tiles for the “rows” of the intra-chunk computation
- $d_h/T$ tiles along the head/value dimension

Inside each thread-block (of size $T\times T$), the kernel:
- Loads $Q$, $K$, $V$ blocks into registers or shared memory
- Accumulates $QK^\top$ and computes softmax statistics with maximal data reuse
- Fuses the elementwise multiplications and the final value matrix multiplication (via tensor cores; e.g., Triton’s `tl.dot` or CUDA’s WMMA)
- Synchronizes only within tiles using `__syncthreads()`, avoiding the need for global barrier synchronization

On-chip SRAM buffers ($\sim T^2$ floats per tile) are used for blockwise accumulators, reducing global memory traffic and maximizing arithmetic intensity by (a) reusing $Q$, $K$ reads, (b) fusing the safe-softmax with block accumulation, and (c) streaming $V$ efficiently.

## 4. Complexity and Memory Analysis

For a chunk of length $C$ and tile length $T$:
- **FLOPs/chunk (mLSTMsig):**
  - Recurrent pass: $O(d_{qk} \cdot d_h + C \cdot d_{qk} \cdot d_h)$
  - Intra-chunk: $\sum_{\mathrm{tiles}} [2T d_{qk} d_h + 3T^2 + O(T^2)] = O(C d_{qk} d_h + C^2)$
  - Inter-chunk matmul: $O(C d_{qk} d_h)$
  - **Total:** $O(C d_{qk} d_h + C^2)$

- **Memory I/O/chunk:**
  - Recurrent: reads $C(d_{qk}+d_h)$, writes $d_{qk} d_h$
  - Parallel: reads $C(2d_{qk}+d_h)$, writes $C d_h$
  - **Total:** $O(C d + d^2)$

Comparison with other kernels:
- Flash Attention requires $O(T^2d)$ FLOPs and $O(T^2)$ I/O.
- FLA requires $O(C d^2 + C^2)$ FLOPs but suffers from $O(T/C d^2)$ extra I/O due to intermediate state storage for each chunk.
- TFLA eliminates this memory bottleneck via intra-chunk tiling and maximally fused compute.

## 5. Performance Benchmarks

Empirical evaluations were conducted with TFLA-mLSTMexp and TFLA-mLSTMsig on NVIDIA H100 GPUs for long-context benchmarks ($\text{embedding-dim}=4096$, $\text{seq-len}=65536$, $\text{head-dim}=256$, $\text{batch}\times T=65536$):
- **Inference (forward only):**
  - TFLA-mLSTMsig is approximately 30% faster than TFLA-mLSTMexp.
  - TFLA-mLSTMsig is 20–40% faster than FlashAttention 3 and over 3× faster than Mamba 2.

- **Training (forward+backward):**
  - TFLA-mLSTMsig achieves a 2× speedup compared to Mamba 2.
  - TFLA-mLSTMsig matches or surpasses FlashAttention 3 performance for sequence lengths above 4k tokens.

Varying chunk size $C$ controls the trade-off between memory usage and runtime: smaller $C$ yields more stored states (higher memory, lower compute), while larger $C$ gives fewer stored states (lower memory, higher compute). On H100, optimal performance occurs for $C=128\ldots256$ and $T=32\ldots64$ [2503.14376].

## 6. Comparative Evaluation and Trade-offs

TFLA achieves several critical advancements:
- Support for arbitrarily large chunks without proportional increase in GPU state storage
- High arithmetic intensity by block-fusing softmax–matmul operations for peak tensor-core utilization
- A tunable chunk size parameter $C$ for direct control over DRAM-I/O versus compute-bound performance and peak memory
- Broad applicability as an efficient drop-in building block for long-context RNN architectures

A direct comparison is summarized below:

| Kernel            | Arithmetic Complexity | Memory I/O | Bottleneck                     |
|-------------------|----------------------|------------|-------------------------------|
| Flash Attention   | $O(T^2 d)$           | $O(T^2)$   | Quadratic time/memory           |
| Linear FLA        | $O(C d^2 + C^2)$     | $O(T/C d^2)$ | Intermediate state materialization |
| TFLA (Blockwise)  | $O(C d_{qk} d_h + C^2)$ | $O(C d + d^2)$ | None (maximal tiling)          |

*The table is strictly constructed from information present in [2503.14376].*

## 7. Context and Applications

TFLA advances the practical deployment of linear RNNs for long-context sequence modeling. By enabling high-levels of parallelism and efficient hardware utilization, it supports accelerated training and inference for language modeling tasks at scales previously prohibitive due to memory and throughput constraints. The kernel’s two-level tiling design and flexible configuration make it a key enabling technology for integrating mLSTM or xLSTM units into pipelines intended for tasks requiring efficient long-sequence processing, particularly where state-of-the-art GPU acceleration is available [2503.14376].

Source: https://www.emergentmind.com/topics/blockwise-flash-kernel