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_ou
t 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),
shape=(1, chunk_size, 1),
padding_mode=ct.PaddingMode.ZERO,
).reshape((chunk_size,))
d_last = ct.extract(d_tile, index=(chunk_size - 1,), shape=(1,))
b_dt_d = b_tile * (dt_tile * ct.exp2(d_last - d_tile))[:, None]
hid_tile = ct.exp2(d_last) * hid_tile
hid_tile = ct.mma(x_tile.transpose(), b_dt_d.astype(b.dtype), hid_tile)
After the final chunk, hid_tile
contains the final hidden state, which is stored. During inference this value can be fed back as the next call’s init_state
; during training it is usually discarded.
ct.store(
final_state,
index=(batch, head, head_tile_idx, state_tile_idx),
tile=hid_tile.reshape((1, 1, BLOCK_HEAD, BLOCK_STATE)),
)
The sequential dependence stays inside a single grid point (program instance), while the independent (batch, head, head_tile, state_tile)
slices run in parallel across the grid.
Kernel 3: compute_out
compute_out
is the final kernel in the forward pass. It takes the per-chunk initial hidden states produced by fwd_pass_state
and computes the token outputs, implementing the output equation derived earlier:
\(Y = \exp(G)Ch_{kT} + \left(\exp(G – G^{\mathsf{T}})\cdot (CB^{\mathsf{T}})\cdot M\right)X + DX\)
A direct PyTorch implementation would compute these three terms separately, materializing the full [B, n_chunks, chunk_size, chunk_size]
interaction matrix C @ B.T
in global memory. compute_out
fuses them into one tiled kernel and keeps that matrix in registers.
Each kernel launch owns one (batch, chunk, head, head_tile)
output tile and works in two phases. Phase 1 loops over state tiles, using ct.mma
to build the interaction matrix cb_tile = C @ B.T
and accumulate the carried-state term C @ H.T
at the same time.
chunk_idx, head_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
cb_tile = ct.full((chunk_size, chunk_size), 0.0, dtype=ct.float32)
out_tile = ct.full((chunk_size, BLOCK_HEAD), 0.0, dtype=ct.float32)
# Phase 1: build CB = C @ B.T and accumulate the carried-state output
for state_tile_idx in range(state_size // BLOCK_STATE):
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))
c_tile = ct.load(
c, 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))
h_tile = ct.load(
hiddens, index=(batch, head, chunk_idx, head_tile_idx, state_tile_idx),
shape=(1, 1, 1, BLOCK_HEAD, BLOCK_STATE),
).reshape((BLOCK_HEAD, BLOCK_STATE))
cb_tile = ct.mma(c_tile, b_tile.transpose(), cb_tile)
out_tile = ct.mma(c_tile, h_tile.transpose().astype(c_tile.dtype), out_tile)
Phase 2 applies the decay and causal mask to cb_tile
(ct.arange
and ct.where
build the lower-triangular mask M
, so no mask tensor is ever materialized) and then adds all of the terms to get the final output tile. Both accumulators start as ct.full
register tiles, and the final ct.mma
casts to ct.tfloat32
to run through the Tensor Cores at near-FP32 accuracy.
# Phase 2: decay + causal mask then final computation
d_tile = ct.load(decay_cumsum, index=(batch, chunk_idx, head),
shape=(1, chunk_size, 1), padding_mode=ct.PaddingMode.ZERO).reshape((chunk_size,))
dt_tile = ct.load(dt, index=(batch, chunk_idx, head),
shape=(1, chunk_size, 1), padding_mode=ct.PaddingMode.ZERO).reshape((chunk_size,))
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))
d_val = ct.load(d, index=(head,), shape=(1,))
cb_tile = cb_tile * ct.exp2(d_tile[:, None] - d_tile[None, :]) # pairwise decay
idx = ct.arange(chunk_size, dtype=ct.int32)
cb_tile = ct.where(idx[:, None] >= idx[None, :], cb_tile, 0.0) # causal mask
out_tile = out_tile * ct.exp2(d_tile[:, None]) # decay carried state
dt_x = dt_tile[:, None] * x_tile
out_tile = ct.mma(cb_tile.astype(ct.tfloat32), dt_x.astype(ct.tfloat32), out_tile)
out_tile = out_tile + d_val * x_tile # D * x skip
ct.store(out, index=(batch, chunk_idx, head, head_tile_idx),
tile=out_tile.reshape((1, chunk_size, 1, BLOCK_HEAD)).astype(out.dtype))
As mentioned, while some numerical unstable exponentials are computed in the ct.exp2(d_tile[:, None] - d_tile[None, :])
line (some of those are positive numbers), immediately mask them out so potential overflows don’t end up poisoning the computation.
Mamba-2 performance optimization
With the Mamba-2 workload defined in the cuTile-Python DSL, several key performance optimizations are introduced:
- Latency hint: cuTile allows you to annotate
ct.load
andct.store
operations with latency hints to model DRAM traffic. The compiler leverages these hints to optimize instruction scheduling and pipeline buffer depth. - Swap A/B: On the NVIDIA Blackwell fifth-generation Tensor Core, MMA operands are asymmetric: operand A can live in Tensor Memory, but operand B must live in Shared Memory. For an input-fusion pattern like the SIMT-op + Dot-op in
fwd_pass_state
, the identity \(B^{\mathsf{T}}A^{\mathsf{T}} = (AB)^{\mathsf{T}}\) is used to route the fused operand (b_dt_d
) to operand A, relieving shared-memory contention. - GMEM layout: The default strided memory access for
decay_cumsum
anddt
lowers to cache-unfriendly LDG instructions. Transposing to a[batch, head, token]
layout ensures contiguous access, enabling the use of TMA. - cuTile-Python autotuner: Since performance bottlenecks shift with input shape, no single optimization strategy fits every model configuration. These optimizations are integrated alongside cuTile parameters (
block_size
, occupancy) into a unified search space, leveraging cuTileexhaustive_search
to automatically find the optimal setup for each input shape.
Mamba-2 forward pass performance benchmarks
This implementation of the Mamba-2 forward pass kernel is available through TileGym. A complete example is shown below.
import torch
import tilegym
from tests import common
from tests.ops.test_mamba2 import Mamba2ProblemSpec
tilegym.set_backend("cutile")
num_tokens, state_size = 16384, 128
torch.manual_seed(0)
spec = Mamba2ProblemSpec(batch_size=4, num_tokens=num_tokens, num_groups=1,
num_heads=8, state_size=state_size, head_size=64,
chunk_size=128, dtype=torch.bfloat16, state_dtype=torch.float32)
fwd_inputs = spec.get_sample_inputs(with_reference=False, with_backward=False)
fwd = common.benchmark(
lambda: tilegym.ops.mamba2_chunk_forward(**fwd_inputs),
use_cupti=True,
)
print(f"Forward: {fwd['forward']['median'] * 1000:.2f} us")
The performance of the implemented Mamba-2 forward pass was evaluated against an open source implementation of the Mamba-2 kernel written in Triton, under the following software and hardware configurations:
Software
- mamba_ssm 2.3.2.post1 (NVIDIA Dynamo-Triton reference), Triton 3.7.1
- cuTile:
cuda-tile
1.5.0 - TileIR assembler (
tileiras
) 13.4 - PyTorch 2.13.0, CUDA 13.3 (Driver and Runtime API, NVCC V13.3.33), Driver 595.58.03
- cuda-python / cuda-bindings 13.3.1
- Python 3.12.3, Ubuntu 24.04.4 LTS, kernel 6.8.0-107
Hardware
- Single NVIDIA B300 SXM6, compute capability 10.3 (NVIDIA Blackwell Ultra)
- 1920 MHz max SM clock, 1100 W power limit, 268 GiB HBM
The results show that the optimized cuTile forward-pass kernels are 1.14x to 1.82x faster, with speedups growing as sequence lengths and state size increases. Note that these are kernel-level speedups only and end-to-end training results will also depend on the backward pass implementations.
Get started with cuTile
cuTile enables you to write clear, readable, and highly performant kernels in Python. Stay tuned for a follow-up post describing how the backward pass is implemented in cuTile.
Get started with cuTile and explore the documentation, GitHub repos, and sample code.