Skip to content

So you wanna CUDA Graph#

GPUs continue to get faster and their hunger for kernel launches driven by the CPU is seemingly insatiable. A natural solution to this problem is cuda graphs. The post will walk through some common footguns both in kernel design and user invocation for ragged kernels, how to alleviate them and where there is still room to improve today - especially at the framework level.

TLDR:#

This tutorial will describe some quarks that arise when using cuda-graphs with ragged attention kernels. It proposes some solutions and fun patterns for dealing with the slowdowns that can arise, while ultimately suggesting codesigning a solution based off of your training dataset.

What are ragged kernels?#

Most sequence data comes from an underlying distribution and does not have one exact length - unless the data is boring. A batch of documents, conversations, or videos therefore has a different number of tokens per sample. "But Tensors are regular! What am I to do" You say, the normal solution truncates or pads every sample to a fixed length.

A ragged representation instead packs only the real tokens into one contiguous buffer:

The metadata preserves the logical boundaries: sequence i occupies tokens[cu_seqlens[i]:cu_seqlens[i + 1]]. Operations like attention use this information to prevent tokens from interacting when they don't belong to the same document. Token-parallel operations like projections do not care about these boundaries, and operate directly on the packed [L, D] tensor. Everybody wins!

What is a CUDA Graph#

Its a graph. Duh, but really - the common PyTorch usage and definition is to record all gpu work issued to a capture stream and its dependent streams, then replay that captured dependency graph later. Why would you want to do this? Because if you know exactly what you want to launch you dont need to redo the Python and dispatcher work that invoked those kernels 1 by 1. Instead you can launch the full Graph in one go and remove a majority if not all the CPU overhead.

Hello World#

Lets setup a small test program, since this is attn_gym we will focus on a simple proxy attention module. However, the learnings will apply to both varlen attention and the new linear-attention variants we are adding to the gym. I will start with varlen for now in the examples.

from collections.abc import Callable, Iterator
from dataclasses import dataclass
from enum import Enum
from itertools import pairwise
from pathlib import Path
from typing import Literal, TypeVar

import torch
import typer
from torch import nn
from torch.cuda.graph_annotations import mark_kernels
from torch.nn.attention.varlen import varlen_attn

Tensor = torch.Tensor
GraphOutput = TypeVar("GraphOutput")
TraceFormat = Literal["chrome_json", "track_event"]
TRACE_PATH = Path(__file__).resolve().parents[1]


class VarLenAttention(nn.Module):
    def __init__(self, dim: int, num_heads: int) -> None:
        super().__init__()

        self.dim = dim
        self.num_heads = num_heads
        self.head_dim = dim // num_heads

        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)

    def forward(
        self,
        x: Tensor,
        cu_seqlens: Tensor,
        max_seqlen: int,
        *,
        mask_inactive_capacity: bool = False,
    ) -> Tensor:
        tokens = x.shape[0]
        qkv = self.qkv(x).view(tokens, 3, self.num_heads, self.head_dim)
        active = None
        if mask_inactive_capacity:
            active = torch.arange(tokens, device=x.device, dtype=torch.int32) < cu_seqlens[-1]
            qkv = torch.where(active[:, None, None, None], qkv, qkv.detach())
        q, k, v = qkv.unbind(dim=1)
        attn = varlen_attn(
            q,
            k,
            v,
            cu_seqlens,
            cu_seqlens,
            max_seqlen,
            max_seqlen,
        )
        if active is not None:
            attn = torch.where(active[:, None, None], attn, 0)
        return self.out_proj(attn.reshape(tokens, self.dim))


def get_zipf_tokens(
    total_tokens: int, n_seqs: int, device: torch.device | str = "cuda"
) -> tuple[Tensor, int]:
    """Distribute tokens with randomized Zipf-ranked sequence lengths."""
    if total_tokens < n_seqs:
        raise ValueError("total_tokens must provide at least one token per sequence")

    ranks = torch.randperm(n_seqs).add(1)
    weights = ranks.to(torch.float64).pow(-1.2)
    scaled_lengths = weights * ((total_tokens - n_seqs) / weights.sum())
    extra_lengths = scaled_lengths.floor().to(torch.int32)
    remainder = total_tokens - n_seqs - int(extra_lengths.sum())
    if remainder:
        fractional_order = (scaled_lengths - extra_lengths).argsort(descending=True)
        extra_lengths[fractional_order[:remainder]] += 1

    lengths = extra_lengths.add(1)
    cu_seqlens = torch.cat(
        (torch.zeros(1, dtype=torch.int32), lengths.cumsum(0, dtype=torch.int32))
    ).to(device)
    return cu_seqlens, int(lengths.max())


def hello_world() -> Tensor:
    from transformer_nuggets.utils.benchmark import profiler

    total_tokens = 4096
    dim = 512
    num_heads = 8
    n_seqs = 32

    cu_seqlens, max_seqlen = get_zipf_tokens(total_tokens, n_seqs)
    model = VarLenAttention(dim, num_heads).cuda().to(torch.bfloat16)
    inputs = torch.randn(total_tokens, dim, device="cuda", dtype=torch.bfloat16)
    # Warmup then run
    for _ in range(5):
        out = model(inputs, cu_seqlens, max_seqlen)
        torch.autograd.grad(out, model.parameters(), torch.ones_like(out), retain_graph=True)
    torch.cuda.synchronize()
    with profiler(TRACE_PATH / "docs/assets/traces/hello_world_no_cuda_graphs"):  # (1)!
        out = model(inputs, cu_seqlens, max_seqlen)
        torch.autograd.grad(out, model.parameters(), torch.ones_like(out), retain_graph=True)
  1. simple wrapper around the standard pytorch profiler

run it and see how where the very expensive gpu is spending its time:

This is somewhat a contrived example since we are using small model dim and token counts, regardless "We Bought the Whole GPU, So We're Damn Well Going to Use the Whole GPU" - if we can.

Graph Time#

PyTorch's api pretty closely mirror the lower level gpu apis for CG (CUDA Graphs, im done writing those two words).

def capture_graph(
    function: Callable[[], GraphOutput],
    warmup: int = 3,
    enable_annotations: bool = False,
) -> tuple[torch.cuda.CUDAGraph, GraphOutput]:
    stream = torch.cuda.Stream()  # (1)!
    stream.wait_stream(torch.cuda.current_stream())
    with torch.cuda.stream(stream):
        for _ in range(warmup):
            function()
    torch.cuda.current_stream().wait_stream(stream)  # (2)!

    graph = torch.cuda.CUDAGraph()
    with torch.cuda.graph(
        graph,
        stream=stream,
        enable_annotations=enable_annotations,
    ):
        output = function()
    torch.cuda.current_stream().wait_stream(stream)
    return graph, output
  1. We use a side stream to isolate the exact work we want to record.
  2. We warmup to initialize any lazy CUDA state such as library handles and then wait on it before the real capture.
def hello_world_graph() -> Tensor:
    from transformer_nuggets.utils.benchmark import profiler

    total_tokens = 4096
    dim = 512
    num_heads = 8
    n_seqs = 32

    cu_seqlens, max_seqlen = get_zipf_tokens(total_tokens, n_seqs)
    model = VarLenAttention(dim, num_heads).cuda().to(torch.bfloat16)
    inputs = torch.randn(total_tokens, dim, device="cuda", dtype=torch.bfloat16)
    parameters = tuple(model.parameters())
    grad_output = torch.ones_like(inputs)

    def forward_backward() -> Tensor:
        output = model(inputs, cu_seqlens, max_seqlen)
        torch.autograd.grad(output, parameters, grad_output, retain_graph=True)
        return output

    graph, output = capture_graph(forward_backward)
    torch.cuda.synchronize()
    with profiler(
        TRACE_PATH / "docs/assets/traces/hello_world_with_cuda_graphs",
        warmup=1,
    ) as active_profiler:
        graph.replay()
        active_profiler.step()
        graph.replay()  # (1)!
        active_profiler.step()
    return output
  1. its replay time!

We can see from the trace that the launch gaps are dramatically smaller and the GPU stream is much denser. Note that the 46us cudaGraphLaunch shown here is inflated by profiling. In general, when I look for CPU overhead I use the PyTorch profiler, but when I want to measure it I use Python's built-in timer.

Terminology time#

Max Tokens T

The physical token capacity of the static input buffers. The captured graph always sees tensors with T token rows. When we capture the worst case, T = T_max.

On replay, Active Tokens L is the amount of real token data. It is available on device as cu_seqlens[-1], and must satisfy L <= T.

Max Sequences N

The number of sequence slots supported by the captured graph. This fixes the shape of cu_seqlens to [N + 1].

On replay, Active Sequences M may be smaller than N. The unused tail is represented by repeating the active token endpoint: [0, ..., L, L, L].

The graph therefore keeps the physical [T, D] and [N + 1] shapes fixed. Each replay changes only the logical workload: L <= T and M <= N.

A More Realistic Training Loop#

Packed batch loader 
@dataclass(frozen=True, slots=True)
class PackedBatch:
    """A fixed-capacity CPU batch with variable active tokens and sequences."""

    input_ids: Tensor
    labels: Tensor
    loss_mask: Tensor
    cu_seqlens: Tensor
    active_tokens: int
    active_sequences: int
    max_seqlen: int


def packed_batch_loader(
    num_batches: int,
    token_capacity: int,
    vocab_size: int,
    sequence_capacity: int,
) -> Iterator[PackedBatch]:
    active_token_counts = torch.linspace(
        token_capacity,
        max(sequence_capacity, token_capacity // 8),
        num_batches,
        dtype=torch.int64,
    )
    for active_tokens_tensor in active_token_counts:
        active_tokens = int(active_tokens_tensor)
        active_sequences = max(
            1,
            round(sequence_capacity * active_tokens / token_capacity),
        )
        active_cu_seqlens, actual_max_seqlen = get_zipf_tokens(
            active_tokens,
            active_sequences,
            device="cpu",
        )

        cu_seqlens = torch.full(
            (sequence_capacity + 1,),
            active_tokens,
            dtype=torch.int32,
        )
        cu_seqlens[: active_sequences + 1].copy_(active_cu_seqlens)
        loss_mask = torch.zeros(token_capacity, dtype=torch.bool)
        for start, end in pairwise(active_cu_seqlens):
            response_start = int(start) + (int(end) - int(start)) // 2
            loss_mask[response_start : int(end)] = True

        yield PackedBatch(
            input_ids=torch.randint(vocab_size, (token_capacity,), pin_memory=True),
            labels=torch.randint(vocab_size, (token_capacity,), pin_memory=True),
            loss_mask=loss_mask.pin_memory(),
            cu_seqlens=cu_seqlens.pin_memory(),
            active_tokens=active_tokens,
            active_sequences=active_sequences,
            max_seqlen=actual_max_seqlen,
        )

def hello_world_training_loop(
    *,
    enable_graph_annotations: bool = False,
    trace_path: Path | None = None,
    trace_format: TraceFormat = "track_event",
    fix_overlapping_events: bool = True,
) -> Tensor:
    token_capacity = 4096  # (1)!
    dim = 4096
    num_heads = 32
    sequence_capacity = 32  # (2)!
    max_seqlen = token_capacity  # (3)!
    num_batches = 4

    batches = packed_batch_loader(
        num_batches,
        token_capacity,
        dim,
        sequence_capacity,
    )
    first_batch = next(batches)  # (4)!
    static_input_ids = torch.empty_like(first_batch.input_ids, device="cuda")
    static_labels = torch.empty_like(first_batch.labels, device="cuda")
    static_loss_mask = torch.empty_like(first_batch.loss_mask, device="cuda")
    static_cu_seqlens = torch.empty_like(first_batch.cu_seqlens, device="cuda")
    static_input_ids.copy_(first_batch.input_ids, non_blocking=True)
    static_labels.copy_(first_batch.labels, non_blocking=True)
    static_loss_mask.copy_(first_batch.loss_mask, non_blocking=True)
    static_cu_seqlens.copy_(first_batch.cu_seqlens, non_blocking=True)
    torch.cuda.synchronize()

    embedding = nn.Embedding(dim, dim, device="cuda", dtype=torch.bfloat16)
    model = VarLenAttention(dim, num_heads).cuda().to(torch.bfloat16)
    parameters = (*embedding.parameters(), *model.parameters())
    optimizer = torch.optim.Adam(parameters, lr=1e-3, fused=True)

    def forward_backward() -> tuple[Tensor, tuple[Tensor, ...]]:
        with mark_kernels("embedding"):
            inputs = embedding(static_input_ids)
        with mark_kernels("attention"):  # (5)!
            output = model(
                inputs,
                static_cu_seqlens,
                max_seqlen,
                mask_inactive_capacity=True,
            )
        with mark_kernels("loss"):
            token_losses = torch.nn.functional.cross_entropy(
                output.float(), static_labels, reduction="none"
            )
            loss = torch.where(static_loss_mask, token_losses, 0).sum() / static_loss_mask.sum()
        with mark_kernels("backward", backward=False):
            grads = torch.autograd.grad(loss, parameters)  # (6)!
        return loss, grads

    def eager_optimizer_step(graph_grads: tuple[Tensor, ...]) -> None:
        for parameter, grad in zip(parameters, graph_grads, strict=True):
            parameter.grad = grad  # (8)!
        optimizer.step()  # (9)!

    trace_profiler = training_loop_profiler(
        trace_path=trace_path,
        trace_format=trace_format,
        fix_overlapping_events=fix_overlapping_events,
    )
    graph, (loss, graph_grads) = capture_graph(
        forward_backward,
        enable_annotations=enable_graph_annotations,
    )
    graph.replay()  # (7)!
    eager_optimizer_step(graph_grads)
    torch.cuda.synchronize()
    with trace_profiler:
        for _ in range(num_batches - 1):
            with torch.profiler.record_function("data_loading"):
                batch = next(batches)
            batch_shape = (
                f"L={batch.active_tokens}, M={batch.active_sequences}, "
                f"max_seqlen={batch.max_seqlen}"
            )
            with torch.profiler.record_function(f"copy_to_static[{batch_shape}]"):
                static_input_ids.copy_(batch.input_ids, non_blocking=True)
                static_labels.copy_(batch.labels, non_blocking=True)
                static_loss_mask.copy_(batch.loss_mask, non_blocking=True)
                static_cu_seqlens.copy_(batch.cu_seqlens, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)  # (10)!
            with torch.profiler.record_function("fwd_bwd_replay"):
                graph.replay()
            with torch.profiler.record_function("optimizer_step"):
                eager_optimizer_step(graph_grads)
        torch.cuda.synchronize()
    return loss
  1. token_capacity is a policy typically chosen by the data pipeline and whatever your global batchsize you found acceptable for your model arch. Every replay must satisfy L <= T.
  2. sequence_capacity is another data-pipeline quantity. If the only guarantee is that a nonempty sequence contains at least one token, then the active count satisfies M <= T and choosing N = T covers the absolute worst case - a batch of single token docs. Obviously thats a shitty dataset so this should be gleaned from real data.
  3. Varlen attention also requires a static upper bound on the maximum sequence length across every replay. If one sequence can consume the full token budget, the safe bound is max_seqlen = T. I will argue later on that the max_seqlen arg is bad, and that there are established patterns to avoid paying for it with minimal perf hit.
  4. This synthetic loader yields the max capacity-sized batch first. Capture therefore sees the largest physical token buffer and all N sequence slots; later batches only change the values copied into those same static tensors.
  5. checkout our fancy new use of mark_kernels, more in the next section
  6. torch.autograd.grad makes each parameter gradient an explicit output of the fwd+bwd graph. Just like loss, these tensors keep the same graph-pool addresses across replay.
  7. The first replay produces the graph outputs for the first batch before we hand its gradients to eager Adam.
  8. parameter.grad = grad attaches those graph outputs to the ordinary optimizer interface
  9. The fused Adam step is updating inplace
  10. set_to_none=True only removes the parameter.grad references. The graph_grads tuple still owns the graph outputs, replay writes the next gradients into the same storage and we attach them again before optimizer.step().

Parameters already have stable storage, above we allocate static input buffers once, copy each new batch into them and keep the graph outputs alive so the allocator cannot reuse their storage. TorchTitan provides utilities that hide the input copies; this example performs them directly.

CG memory management is a big topic. A little to big for this tutorial. For a taste, larger distributed systems can use more advanced storage handoffs instead of retaining and directly consuming the original graph-output tensors.

What about intermediaries? This module produce intermediate in projections prior to calling attention. How do we ensure that this output always lands in exactly the right the slot? Even if we could ensure that how do we keep this tensor alive between graph replays??

What about intermediates?#

The graph does not keep every intermediate Python Tensor alive. During capture the caching allocator gives intermediates addresses from a graph-private pool and the kernels record those exact addresses. On replay Python and the allocator do not recreate the intermediates; CUDA just launches the recorded kernels reading and writing the same addresses. The private pool remains alive until every graph using it and every live tensor created during capture is gone.

Small Digression into profiling#

Eager code is great. It works seamlessly with the stock pytorch profiler - providing plenty of info for an intrepid user. One very handy feature is record_function that lets users annotate regions of your code and have these labels land on the perfetto trace. If you have gigs to spare you might even use with_stacks=True and get a microscopic view of the world. This is not the same for cuda-graphs, historically. Below is an example of what the stock profiler might produce.

In the stock-profiler half of the merged trace below you can still see which graph launch owns what kernels, their graph and node ids, launch metadata and dependency arrows. This is already useful, but the graph still looks like a wall of kernel names. We can do better!

At PyTorch we believe observability is only becoming more important. What used to take ages for an individual developer to digest can take an agent seconds, provided we give it structured traces instead of screenshots and a wall of kernel names;

1. There are new CUDA Graph annotation APIs. Python does not run the individual operators again when a graph is replayed, so record_function cannot recover the internal regions after capture. mark_kernels records metadata while the graph is being captured, and enable_annotations=True keeps the mapping from graph nodes back to those labels. With this info stored on the graph nodes we can write postprocessors like: transformer-nuggets post processing which joins that metadata to the replayed kernels and reconstructs our labels -> embedding, attention, loss and backward.

2. CUPTI's Hardware Event System (HES) is a separate Blackwell feature for collecting kernel timestamps with lower per-node overhead. The comparison below intentionally uses the regular PyTorch profiler on both sides; the only difference is whether graph annotations are enabled.

If you think wow these are some well annotated cuda-graph traces now you know why - sidequest done!

The anatomy of a ragged kernel#

A ragged kernel needs to map a regular GPU grid onto sequences with different lengths. Turns out there are a few ways to do this.

Assume that for every token we have some parallel work to do. Attention forward is one example: every input token needs an output. We could launch one CTA per token, no that is dumb and not how gpus work. Instead we tile the tokens into chunks. For the rest of this section assume one CTA handles one chunk of \(C\) tokens. Real kernels may divide the work further, but this gives us something concrete to schedule.

Suppose a packed tensor contains N sequences and sequence \(i\) contains \(s_i\) tokens. A simple mapping uses one grid axis for the sequence and another for its chunk:

grid(sequence_i, chunk_j)

start = chunk_j * C
if start >= seqlen[sequence_i]:
    # wrong sequence bro
    return
process(sequence_i, start, C)

The chunk axis must be large enough for the longest sequence. Since this is baked into the grid it also must be static, which means we need to launch \(\left\lceil S_{\max} / C \right\rceil\) chunks for each of the N sequences:

\[ G_{\mathrm{rect}} = N \left\lceil \frac{S_{\max}}{C} \right\rceil, \qquad S_{\max} = \max_i s_i, \]

This covers every sequence, but it overcounts. The minimal covering is:

\[ G_{\mathrm{active}} = \sum_{i=0}^{N-1} \left\lceil \frac{s_i}{C} \right\rceil \]

The problem is that the longest sequence stretches out our required grid. CTAs assigned to chunks outside shorter sequences immediately return, which sounds cheap but it can add up.

FlashAttention 2 used this same rectangular grid launch, requiring the host to keep track of max_seqlen_q.

Perhaps there is another way#

We are adding new attention variants to the gym and will use KDA as our first case study.

1. Scheduling without max_seqlen#

If you look at our chunk_kda function you will notice it does not accept max_seqlen at all. But how!?

The rectangular scheduler uses max_seqlen to decide how many chunk slots every sequence should receive. What if we stop treating our grid as rectangular but instead flatten all of the sequence-local chunks into one minimal logical work list. A CUDA Graph-safe kernel can cover that list with either a static capacity grid or bounded persistent workers, then use offsets to recover which sequence owns each chunk. This sounds kinda similar to what we originally did for ragged tensors in memory!

For sequence lengths \(s_i\), we build chunk_offsets:

\[ \mathtt{chunk\_offsets}[0] = 0, \qquad \mathtt{chunk\_offsets}[i+1] = \mathtt{chunk\_offsets}[i] + \left\lceil \frac{s_i}{C} \right\rceil. \]

This is just another prefix sum over sequences saying how many chunks are needed to cover each. chunk_offsets[i] is where sequence \(i\)'s chunks begin in the flat list, and

\[ \mathtt{chunk\_offsets[-1]} = G_{\mathrm{active}} \]
Dont we already have a prefix sum -> cu_seqlens#

The ceiling must be applied separately to each sequence because chunking resets at every sequence boundary. For \(C = 64\):

cu_seqlens [0, 65, 128]
difference
lengths [65, 63]
ceil each / 64
chunks [2, 1]
prefix sum
chunk_offsets [0, 2, 3]
ceil_div(cu_seqlens, 64) [0, 2, 2] wrong boundary

The flat grid gives each CTA one integer, flat_chunk, but the CTA still needs to know which sequence it belongs to and where that chunk begins in the packed tensor. With our handy new prefix sums we can do the following:

sequence = search_first_greater(chunk_offsets, flat_chunk) - 1
local_chunk = flat_chunk - chunk_offsets[sequence]
token_start = cu_seqlens[sequence] + local_chunk * C

For the example above, chunk_offsets = [0, 2, 3]. Flat chunks 0 and 1 belong to sequence 0; flat chunk 2 belongs to sequence 1 and is its local chunk 0, beginning at packed-token offset 65. Since our sums are monotonic we can binary_search our way to the right location.

There is 1 more missing piece to the puzzle though; what is the host launch size? Instead of getting it from max_seqlen, we can derive an upper bound from the static token capacity T, sequence capacity N and chunk size \(C\):

Let M be the number of nonempty sequences in this replay. We know \(M \le N\) because there are only N sequence slots, and \(M \le T\) because every nonempty sequence needs at least one token. Therefore \(M \le \min(T,N)\).

For each nonempty sequence, \(\lceil s_i/C \rceil = 1 + \lfloor (s_i-1)/C \rfloor\). Starting from the real chunk count:

\[ \begin{aligned} G_{\mathrm{active}} &= \sum_{i=0}^{M-1} \left\lceil \frac{s_i}{C} \right\rceil \\ &= M + \sum_{i=0}^{M-1} \left\lfloor \frac{s_i-1}{C} \right\rfloor \\ &\le M + \left\lfloor \frac{\sum_{i=0}^{M-1} s_i-M}{C} \right\rfloor \\ &\le M + \left\lfloor \frac{T-M}{C} \right\rfloor. \end{aligned} \]

Using the largest possible active sequence count, \(M_{\max} = \min(T,N)\), gives the host capacity:

\[ G_{\mathrm{cap}}(T,N,C) = M_{\max} + \left\lfloor \frac{T-M_{\max}}{C} \right\rfloor. \]

The per-sequence tile-prefix and flat-decoding pattern is not a new idea, see:

2. Graph overcapture: bound capacity-sized launches#

Time to bring CUDA Graphs back into this party. We just saw how to not pay worst case sequence peformance with an updated grid. What about worst case num_tokens? During capture we have to use T_max tokens - for the worst case.

Although we have tightend the bounds - a graph captured for T_max keeps launching chunk_capacity(T_max, N, C) even when a replay contains only L << T_max active tokens.

What we need is a yet another type of Grid Schedule. Building off of the previous lets add a PERSISTENT schedule.

Here is the flat-task Triton form. CuTeDSL K3/K4 keep their pair/head grid axes and stride only the chunk axis, but the idea is the same:

capacity_tasks = chunk_capacity(T_max, N, chunk) * subtasks
num_workers = min(capacity_tasks, num_sms * ctas_per_sm)
launch grid(num_workers)

kernel(worker):
    active_tasks = chunk_offsets[-1] * subtasks
    for task in range(worker, active_tasks, num_workers):
        process(task)

Each worker reads the active task count from the device and strides over the real task list. The launch stays fixed but is capped at the resident-worker count; the task-loop work scales with the active list instead of worst case tokens. This does not mean persistence always wins—the static early-exit grid is still better when captured and active work are close.

Show me the numbers!#

There are really two separate ways our captured graph can guess too high:

  1. Too many sequences. Keep the real token count fixed at L =8192, capture room for N =256 sequences and change how many real sequences M make up those tokens.
  2. Too many tokens. Capture T_max =16384 tokens, keep M =64 and replay progressively fewer real tokens L.

For every complete implementation I also capture an ideal graph for the actual (L, M, max_seqlen) shape with dashed lines. Comparing the worst-case graph against the ideal graph shows how much forward-plus-first-order-backward replay time we lose to extra scheduling capacity. The FA4 persistent-forward series uses exact points for its slowdown denominator but does not draw a separate dashed trace.

Holding L fixed while changing M also changes the tokens per document and therefore the chunk and tile counts. So its not a perfect isolation of M's effect. The third constant sequence length tab tries to isolate this -> every active sequence has 128 tokens while also moving the actual vs capacity ratios M/N=L/T in lock step. It uses max_seqlen_q=8192 for FA2/FA4 (absolute worst case).

Captured versus exact CUDA Graphs under fragmentation and token overcapture

Note: the third graph was done on a GB200 vs the B200 - at different power limits.

Does this hold E2E#

I have been cheating a little. Every operation above is a ragged kernel that natively supports cu_seqlens. A real attention module also has token paralell QKV and output projections, normalization, masks and a bunch of ordinary PyTorch operations.

Most of those operations only see a physical [T_max, D] tensor. Therefore they processes T_max rows even when only L are real. At T_max =8192 and L =512, that is 16 times as many rows!

Correctness tangent#

Another sidequest but this one is important because I have seen many an issue that boils down to not masking the padding tokens correclty. The inactive padding token suffix of your ragged tensor is not harmless. Its easy to remeber these tokens in the fwd; i.e. loss mask, you cant forget about the grad. A op may leave inactive outputs or input gradients undefined. That is kind of the point from the above secitons; dont do wasted work and only zero out the padded tokens when you need to. Ohh the joys of floating points; NaN * 0 = NaN.

The training example therefore uses two different masks:

        # Keep the endpoint on-device: a captured graph rebuilds this one mask on
        # replay, then every value mask and gradient barrier below reuses it.
        active_mask = (  # (1)!
            active_token_mask(hidden_states, cu_seqlens)
            if self.mask_inactive_capacity and cu_seqlens is not None
            else None
        )
        # A zero `grad` cannot neutralize a NaN activation in a weight reduction.
        hidden_states = mask_inactive_tokens(hidden_states, active_mask)  # (2)!
        hidden_states_compute, qkv = self.qkv_projection(hidden_states)
        # The short-convolution dInput suffix is undefined; keep it out of qkv_proj dW.
        qkv = mask_inactive_token_gradients(qkv, active_mask)  # (3)!
  1. Build one reusable device mask from cu_seqlens[-1]. Because this happens inside capture, replay rebuilds it from the current device endpoint without a host read.
  2. Value masking writes the inactive rows to zero before an ordinary operation can read or save them.
  3. Leave the forward tensor unchanged, but zero its inactive-row gradients during backward. Placing the barrier after qkv_projection prevents undefined inactive gradients from the short convolution from contaminating the projection's weight-gradient reduction.

The mental model is: mask where inactive rows can participate in a mix or reduction, or where inactive values or gradients become undefined.

A projection Y = XW is token-parallel in the forward, but its weight gradient reduces over every physical row:

\[ \nabla W = X^\mathsf{T}\nabla Y = \sum_t X_t^\mathsf{T}\nabla Y_t. \]

Zero out the padding before an operation where inactive rows can mix, contribute to a token-global statistic, or be saved for a backward reduction over rows. If a ragged op may produce undefined gradients for padded rows, zero them before they reach the previous layer. Here, mask_inactive_token_gradients stops the short convolution's invalid padded-row dInput before qkv_projection uses it to compute dW.

mini-TLDR: pure tokenwise operations with defined padded outputs and gradients do not need a mask at every boundary. The masks belong where rows first mix, reduce, or become undefined.

A real KDA module#

Putting whole KDAAttention module back together: QKV and output projections, normalization, short convolution, masking, the KDA core and backward. This Kimi-style shape uses T_max =4096, N =32, hidden size 2304 and 32 KDA heads.

KDAAttention scaling and its fixed-shape module floor

At full capacity chunk_kda core is 1.23 ms versus 3.43 ms for the complete module. At half tokens kda falls to 0.77 ms -> a 38% reduction. However, the complete module only falls to 2.91 ms, or 15%. Amdahl's law strikes again! The padding-aware part scales while most surrounding work still processes T_max rows :(

Below is a merged trace of full and half tokens. Feel free to inspect for yourself.

Could we use more than one graph?#

This MM tax seems harder to fight without rewriting all of PyTorch. What if we just captured a few physical token capacities and selected the smallest one that fits L?

capacity = min(candidate for candidate in capacities if L_host <= candidate)
graphs[capacity].replay()

Here L_host is the token count the data loader already knows; reading cu_seqlens[-1].item() only to dispatch would introduce a device-to-host sync. Compute-wise this is great, but we have to be careful about memory.

Pools and Pools of Memory#

As we know, replay needs the same exact virtual addresses it saw during capture. PyTorch protects those addresses by retaining graph allocations in a private allocator pool. This however causes memory fragmentation: the default allocator cannot borrow an inactive private block while the graph still owns it.

If one Naively captures every sub graph in its own pool then we get no memory resuse. But if we capture them on the same side stream and ensure their GPU executions never overlap, we can give them one shared pool. Our static inputs and outputs can all be slices from max-sized backing buffers!

static_input = torch.empty((T_max, D), device="cuda")
static_output = torch.empty_like(static_input)
pool = torch.cuda.graph_pool_handle()

for capacity in capacities:
    with torch.cuda.graph(graphs[capacity], pool=pool):
        static_output[:capacity].copy_(model(static_input[:capacity]))

All input prefix views share the input base address, and all output prefix views share a separate output base address. Graph-local intermediates can reuse the same memory because only one bucket executes at a time.

Measured on B200 from fresh-process allocator snapshots of the same two-projection BF16 core, T_max =1024 and D=4096; the shared-pool case also copies into an external output buffer.

Crossing graph boundaries#

Sharing a pool also shares lifetimes. A pool-backed output is temporary and can be snatched from under you: the next graph using that pool may overwrite it. This can be tricky in real programs, so Inductor has some great functionality to help here: CUDA Graph Trees manage compatible path generations and pending forward/backward lifetimes in one pool.

A common boundary break in CG for training is fwd_bwd -> optimizer. The optimizer graph expects every gradient at the address it captured. Copying graph-owned grads into persistent eager buffers works, but now some of your largest tensors exist twice! There are tricks that can release the storage between steps and recover it before optimizer replay but suffice it to say be careful and use the memory profiler to help you track all your allocations.

We have a nice little system setup now: a small set of CUDA Graphs handles the coarse changes in physical token count, while our device-driven ragged kernels handle the sequence count, lengths and remaining underfill within each graph. Together the demonstrated KDA/FA schedulers can scale across many input sizes without needing one graph for every exact shape!

Show me the data!#

Hopefully you are convinced from the above that padding tokens are expensive, and there are lower level ways to mitigate this cost. HOWEVER, like all great ML problems we have try to engineer our way out of this problem -> when perhaps we should have just looked at the data..

I sampled 20,000 source rows from Dolma 3's 5.93T-token mix, the OLMo 3 7B stage-1 pretraining mix. Using the OLMo 3 tokenizer, I split them into 8192 token capped segments. A trick to get much better packing is to let your dataloader look ahead and try to pack samples from this buffer of candidates. Thus a candidate buffer of 32 means each batch may inspect and select from at most 32 entries.e

The figure below sweeps local token budgets from 8K to 32K and shows how much of each budget this packing policy actually fills.

Packing efficiency by dataset and candidate-buffer size

For candidate-buffer size 32:

local T median M p99 M mean fill mean padding
8192 1 13 99.80% 16 tokens
16384 4 32 99.75% 40 tokens

For this Dolma segmentation and packing policy we probably dont even need multiple token-capacity graphs. M <=32 is guaranteed by the candidate policy, so N =32 is sufficient by construction; the loader also fills the token budget almost perfectly. The raw row overhead of one graph is about 0.2% at 8K and 0.25% at 16K.

Is this trend universal? No. An illustrative, not tokenizer-controlled comparison, look at FineWeb sample-10BT sample using tiktoken:o200k_base. At a 16K budget its mean underfill falls from 8.8% with candidate buffer 32 to 0.7% with buffer 64.

Conclusion#

This lil tutorial ended up more meandering then expected, but, in the end we have found 3 complimentary techincques for handiling ragedness in training:

  1. Measure the data, then pack it away in the dataloader. Choose T, N and the candidate-buffer policy from the real sequence distribution instead of assuming one packing policy works everywhere.
  2. Make the sequence-aware kernels read the real work dynamically on the GPU. The demonstrated KDA and FA schedulers can support multiple schedules. A runtime work-aware persistent grid can substantially reduce padding work when overcapture is large, while the static early-exit grid can still win near full capacity.
  3. Add graph buckets in a way that limits memory fragmentation. A core vLLM idea applied to training: capture a finite list of physical capacities, select the smallest T_bucket >= L, and share one global graph pool.

The final recipe is: measure the dataset, pick T, N and engineer prolbems out in the dataloader when you can, start with one graph, add capacity buckets only if the measured padding justifies them, let sequence-aware kernels read the actual work from the GPU and mask whatever padding is left.