loading
loading
NVIDIA CUDA Tile (cuTile) is a Python-based DSL and parallel programming model for NVIDIA GPUs that automatically leverages advanced hardware capabilities to deliver peak performance for tile-level GPU programming. By abstracting thread- and warp-level scheduling behind high-level, tile-centric primitives, cuTile enables you to write GPU kernels in Python while compiler optimizations automatically target advanced hardware features like Tensor Cores.
As an early adopter of cuTile, machine learning engineers at Jane Street implemented Mamba-2, a state-of-the-art sequence model, in cuTile and partnered with NVIDIA to optimize the kernels. Jane Street is a deeply technical global trading firm, using deep learning to drive innovation in advanced quantitative research. This post walks you through the three custom kernels powering the Mamba-2 forward pass and the cuTile optimizations behind them.
Overview and prerequisites
To use cuTile, ensure your environment meets the requirements outlined in the cuTile Python Quickstart guide. Import the package as cuda.tile and expose it as ct throughout the kernel implementations featured in this post.
Sequence modeling computational trade-off
Sequence modeling architectures face a computational trade-off: traditional softmax attention trains rapidly using matrix multiplication accelerated by Tensor Cores, but at inference, compute and memory requirements increase linearly with sequence length. In contrast, recurrent neural networks (RNNs) suffer from slow, nonparallelizable training but require constant compute and fixed memory footprint regardless of sequence length.
Mamba-2 architecture
The Mamba-2 architecture comes from a family of models aiming for the best of both worlds. Similar to an RNN, it has a fixed size hidden state and a recurrent update that can be computed chunkwise in a computationally efficient manner using Tensor Cores.
The core recurrence of Mamba-2 is:
\(h_t = \exp(A \cdot dt_t) h_{t-1} + dt_t B_t x_t\)
Where:
- \(x_t\), \(dt_t\), \(B_t\), and \(c_t\) are the inputs to the Mamba-2 layer (obtained by projecting the inputs to the block, like queries, keys and values in an attention layer) at token \(t\)
- \(y_t\) is the output at token \(t\)
- \(h_t\) is the fixed-size hidden state carried across tokens. It’s not a vector as in RNNs, but a matrix (more numbers and more representational capacity)
- \(A\) and \(D\) are learned model parameters
Intuitively:
- \(dt_t\) (one scalar per head) represents how much of the state should be updated. All values in \(dt\) are negative Because the previous state is multiplied \(\exp(Adt_t)\), the state stays the same when \(dt_t = 0\)
- \(x_t\) (a vector of size H) plays a role similar to the values in attention
- \(B_t\) (a vector of size S) plays a role similar to the keys in attention
- \(c_t\) (a vector of size S) is how the state is mixed to produce the outputs and plays a role similar to the queries in attention
The state \(h_t\) is a matrix of size S,H (\(B_t x_t^\top\) is an outer product). To make computation easier to write, use \(a_t = Adt_t\) and call \(b_t\) the product \(dt_t B_t\). The equation is then simply:
\(h_t = \exp(a_t)h_{t-1} + b_t x_t^{\mathsf{T}}, y_t = c_t^{\mathsf{T}}h_{t-1} + Dx_t\)
This recurrent formulation is ideal for sequential inference. However, training and prefill require something faster than a for-loop. While the linear state update allows for a parallel scan algorithm, this approach requires materializing an enormous hidden state tensor making it primarily memory-bound and unable to utilize Tensor Cores.
An alternative approach leverages Tensor Cores by computing any future hidden state directly through matrix multiplication. Ignoring the \(\exp(a_t)\) term for simplification, going from \(h_t\) to \(h_{t+T}\) is just:
\(h_{t+T} = h_t + b_{t+1}x_{t+1}^{\mathsf{T}} + b_{t+2}x_{t+2}^{\mathsf{T}} + \cdots + b_{t+T}x_{t+T}^{\mathsf{T}}\)
If building a matrix \(B\) of shape (T, S) by stacking all of the \(b_{t+i}\) and a matrix \(X\) of shape (T, H) by stacking all of the \(x_{t+i}\) then this is just \(h_{t+T} = h_t + BX^T\):
With this jump formula, you can efficiently compute the hidden state at various chunk “checkpoints” in the sequence. The next step is to find a way to compute outputs \(y_1\), \(y_2\), \(y_3\), \(y_4\) given only \(h_0\). This would be easy if computing \(h_1\), \(h_2\), \(h_3\), \(h_4\), … but doing so would defeat the point of the jump formula.
Still removing the decay term \(\exp(a_t)\) for now, the goal is to compute:
\(y_{t+i} = c_{t+i}^{\mathsf{T}}h_{t+i} + Dx_{t+i} = c_{t+i}^{\mathsf{T}}\left(h_t + b_{t+1}x_{t+1}^{\mathsf{T}} + \cdots + b_{t+i}x_{t+i}^{\mathsf{T}} \right) + Dx_{t+i}\)
Where the previous formula is used to expand the hidden state. The skip connection term \(Dx_{t+i}\) is easy to reason about. For the others, it looks like this:
Again, grouping a chunk of \(x_t\), \(b_t\), \(c_t\) and \(y_t\) in matrices enables re-expressing this with Tensor Core-friendly matrix products (Figure 3) where the triangle in the \(\mathit{C}\) by \(\mathit{B}\) product is a causal mask (like in causal softmax attention).
When computing the output \(y_{t+i}\) you have some product of \(c_{t+i}\) by all of the \(b_j\) for \(j \leq t + i\) so when computing the full matrix product, you need to mask out the ones in the future (above the diagonal). In math this becomes (adding back the skip connection):
\(Y = Ch_t + \left((CB^{\mathsf{T}})\cdot M\right)X + DX\)
With \(M\) the traditional causal mask.
Note that this looks very similar to attention (without the softmax) because when the decay term is removed, Mamba-2 is just linear attention (with some renaming). Adding back the decay makes things a little bit trickier. The expansion of the recurrence will look like:
\(h_(t+T)=exp(a_(t+1)+a_(t+2)+⋯+a_(t+T) ) h_t+exp(a_(t+2)+⋯+a_(t+T) ) b_(t+1) x_(t+1)+exp(a_(t+3)+⋯+a_(t+T) ) b_(t+2) x_(t+2)+⋯+b_(t+T) x_(t+T)\)
Reformulating in log-space allows for replacing sequential multiplications with cumulative sums. This remains a matmul if the rows of \(B\) (or \(X\)) are pre-multiplied by the exponentials. Keeping track of the cumulative sum of \(a_t\) in a tensor \(G\), this results in the following new equation:
\(h_{t+T} = \exp(G_{last})h_t + \left(\exp(G_{last} – G)\cdot B\right)X^{\mathsf{T}}\)
Plugging the new hidden state jump into the output computation results in the following:
\(Y = \exp(G)Ch_{kT} + \left(\exp(G – G^{\mathsf{T}})\cdot (CB^{\mathsf{T}})\cdot M\right)X + DX\)
where the matrix \(\exp(G – G^{\mathsf{T}})\) contains at \((i, j)\) the term \(\exp(g_i – g_j)\). Note that all of the terms where \(i < j\) will be masked by the causal mask, so only keep terms where \(g_t\) is the sum of more of the \(a_t\) than \(g_j\). By construction, all \(a_t\) values are negative, which makes \(g_i – g_j\) negative. Any exponentials of positive numbers are never kept, guaranteeing numerical stability.
This will translate into three kernels:
- chunk_cumsum: Compute the cumulative sums of the \(a_t\) inside each chunk.
- fwd_pass_state: Compute the hidden state at the beginning of each chunk using the jump formula above.
- compute_out: Compute the outputs for each chunk using the initial hidden state for the chunk computed in the previous step.
Kernel 1: Computing cumulative sums
The cumulative sums of the \(a_t = Adt_t\) (the \(G\) in the previous math) are needed in both Steps 2 and 3. To avoid recomputing several times, this first kernel will compute \(a_t\) once and store it. Note that this is not a global cumulative sum, just the cumulative sum within each chunk.
The inputs to this kernel are a per-head scalar decay a and a per-token timestep dt. The effective decay at token t for head h is a[h] * dt[t, h], and the kernel accumulates that product along the chunk dimension. These cumulative sums are all in log-space and the exponential of them is always taken (or some linear combination of them). Since it’s much cheaper to compute exp2 instead of exp, the result is scaled by log2(e).
This kernel introduces the cuTile primitives that all three kernels use: ct.bid, ct.load / ct.store, and ct.cumsum.
The full kernel is short:
@ct.kernel
def chunk_cumsum_kernel(
a: torch.Tensor, # [num_heads]
dt: torch.Tensor, # [batch_size, num_tokens, num_heads]
output: torch.Tensor, # [batch_size, num_tokens, num_heads]
chunk_size: ct.Constant[int],
BLOCK_H: ct.Constant[int],
) -> None:
chunk_idx, head_tile_idx, batch = ct.bid(0), ct.bid(1), ct.bid(2)
dt_tile = ct.load(dt, index=(batch, chunk_idx, head_tile_idx),
shape=(1, chunk_size, BLOCK_H),
padding_mode=ct.PaddingMode.ZERO).astype(ct.float32)
a_tile = ct.load(a, index=(head_tile_idx,),
shape=(BLOCK_H,),
padding_mode=ct.PaddingMode.ZERO).astype(ct.float32)
d_tile = ct.cumsum(a_tile * dt_tile, axis=1) * LOG2E
ct.store(output, index=(batch, chunk_idx, head_tile_idx), tile=d_tile.astype(output.dtype))
The kernel is launched with a grid of (num_chunks, num_head_tiles, batch_size), the first line determines where you are in the grid using ct.bid. The broadcast product of a_tile (shape [BLOCK_H]) and dt_tile (shape [1, chunk_size, BLOCK_H]) gives the effective per-token, per-head decay over the tile. ct.cumsum(..., axis=1) then scans along the intra-chunk token dimension, and the result is stored in the layout expected by fwd_pass_state and compute_out.
Kernel 2: fwd_pass_state
fwd_pass_state computes the initial hidden state of each chunk and the final hidden state after the last chunk. It does not directly compute the output. Instead, it prepares the carried state that compute_out will later project into outputs. At its core, the kernel is a sequential loop over chunks as described earlier. Each kernel launch loads one (batch, head, head_tile, state_tile) slice of the hidden state, as shown below. Within that slice, hid_tile has a shape [BLOCK_HEAD, BLOCK_STATE] and will contain a tile of the initial hidden state for each chunk.
head_tile_idx, state_tile_idx, batch_head = ct.bid(0), ct.bid(1), ct.bid(2)
batch = batch_head // x.shape[2]
head = batch_head % x.shape[2]
group = head // num_heads_per_group
hid_tile = ct.load(
init_state,
index=(batch, head, head_tile_idx, state_tile_idx),
shape=(1, 1, BLOCK_HEAD, BLOCK_STATE),
).reshape((BLOCK_HEAD, BLOCK_STATE))
In each iteration of the loop, first store the hidden state:
num_chunks = ct.cdiv(b.shape[1], chunk_size)
for chunk_idx in range(num_chunks):
ct.store(
hiddens,
index=(batch, head, chunk_idx, head_tile_idx, state_tile_idx),
tile=hid_tile.reshape((1, 1, 1, BLOCK_HEAD, BLOCK_STATE)).astype(hiddens.dtype),
)
The rest of the loop loads one chunk of x, b, dt, and decay_cumsum to compute the next chunk’s initial hidden state. As shown previously, the math is:
\(h_{t+T} = \exp(G_{last})h_t + \left(\exp(G_{last} – G)\cdot B\right)X^{\mathsf{T}}\)
ct.extract enables taking the last token from the decay tile loaded (to compute \(G_last\)). Use exp2 instead of exp for performance reasons, as shown previously. ct.mma is the API that will execute the matrix-multiply add while leveraging the Tensor Cores.
x_tile = ct.load(
x, index=(batch, chunk_idx, head, head_tile_idx),
shape=(1, chunk_size, 1, BLOCK_HEAD),
padding_mode=ct.PaddingMode.ZERO,
).reshape((chunk_size, BLOCK_HEAD))
b_tile = ct.load(
b, index=(batch, chunk_idx, group, state_tile_idx),
shape=(1, chunk_size, 1, BLOCK_STATE),
padding_mode=ct.PaddingMode.ZERO,
).reshape((chunk_size, BLOCK_STATE))
dt_tile = ct.load(
dt, index=(batch, chunk_idx, head),
shape=(1, chunk_size, 1),
padding_mode=ct.PaddingMode.ZERO,
).reshape((chunk_size,))
d_tile = ct.load(
decay_cumsum, index=(batch, chunk_idx, head),
document fetched 2026-10-24 11:20 · canonical https://developer.nvidia.com/blog/accelerating-mamba-2-with-high-performance-forward-pass-kernels-in-nvidia-cuda-tile