The 0.15 GB Breakthrough: Isolating a Fusion Bug in PyTorch's FlexAttention on Blackwell GPUs

In the trenches of large-scale ML training, debugging often resembles forensic science more than software engineering. The subject message of this article — message index 7894 in a sprawling opencode session — captures one such pivotal forensic moment. In this single bash command, an AI assistant runs a targeted Python test on a remote machine to determine whether torch.compile(flex_attention) produces a fused backward kernel when used with a custom attention mask (BlockMask) on NVIDIA Blackwell GPUs (sm_120). The result is unambiguous: Fwd: 0.08 GB, Bwd: 0.15 GB, FUSED - OK. This tiny output, just six lines, resolves a question that had consumed the preceding hours of debugging and unlocks the path forward for training a speculative decoding model called DFlash.

The Context: A Cascade of OOM Crashes

To understand why this message matters, we must step back into the broader session. The assistant was training DFlash — a speculative decoding drafter model — on a 4× RTX PRO 6000 Blackwell GPU node. The training pipeline had already survived six script bugs (incorrect config copying, missing sequence packing, absent noise augmentation, per-document anchor boundary violations, wrong position IDs, and lack of torch.compile). But once those were fixed, a new class of hardware-specific failures emerged: out-of-memory (OOM) crashes during the backward pass of the drafter's attention computation.

The crashes were maddeningly consistent. GPU 2 (the drafter GPU) would hit exactly 84.07 GB used with 10.90 GB free — always the same 15 GB allocation that corresponded to an unfused attention score matrix. The error trace pointed to sdpa_dense_backward inside PyTorch's flex_attention implementation. The assistant had already confirmed in an earlier test (message 7886) that uncompiled flex_attention used 17.85 GB for the backward pass, while compiled flex_attention used only 0.26 GB. The fused kernel, generated by Triton through torch.compile, avoided materializing the full QK^T score matrix — the very thing causing the OOM.

Yet despite adding a module-level _compiled_flex_attention = torch.compile(flex_attention) to the model code, the training still crashed with the same dense backward error. The assistant's reasoning in message 7893 reveals the confusion: "Still the same sdpa_dense_backward error! The module-level torch.compile(flex_attention) is not being used." The assistant speculated that the BlockMask — a custom attention mask object created by create_block_mask with a closure-based mask function — might be breaking the compiler's ability to trace through the flex_attention call, causing it to fall back to the eager-mode dense implementation.

The Subject Message: A Controlled Experiment

Message 7894 is the assistant's response to this hypothesis. Rather than continuing to speculate, the assistant designs a minimal, controlled experiment that isolates the exact variable in question: does torch.compile(flex_attention) produce a fused backward when a BlockMask is passed?

The test is executed via SSH on the remote Blackwell node:

import torch
from torch.nn.attention.flex_attention import flex_attention, create_block_mask

device = "cuda:2"
torch.manual_seed(42)

heads, kv_heads, head_dim = 32, 8, 128
Q_LEN, KV_LEN = 2048, 10240

q = torch.randn(1, heads, Q_LEN, head_dim, device=device, dtype=torch.bfloat16, requires_grad=True)
k = torch.randn(1, kv_heads, KV_LEN, head_dim, device=device, dtype=torch.bfloat16, requires_grad=True)
v = torch.randn(1, kv_heads, KV_LEN, head_dim, device=device, dtype=torch.bfloat16, requires_grad=True)

def causal_mask(b, h, q_idx, kv_idx):
    return q_idx >= kv_idx

block_mask = create_block_mask(causal_mask, B=None, H=None, Q_LEN=Q_LEN, KV_LEN=KV_LEN, device=device)

print("=== With BlockMask, compiled flex_attention ===")
compiled_fa = torch.compile(flex_attention)
torch.cuda.reset_peak_memory_stats(device)
out = compiled_fa(q, k, v, block_mask=block_mask, enable_gqa=True)
fwd = torch.cuda.max_memory_allocated(device)/1e9
torch.cuda.reset_peak_memory_stats(device)
out.sum().backward()
bwd = torch.cuda.max_memory_allocated(device)/1e9
print(f"Fwd: {fwd:.2f} GB, Bwd: {bwd:.2f} GB")
if bwd < 5:
    print("FUSED - OK")
else:
    print("DENSE - will OOM")

The output: Fwd: 0.08 GB, Bwd: 0.15 GB, FUSED - OK.

This is a decisive result. The compiled flex_attention works correctly with a BlockMask, producing a fused backward kernel that uses only 0.15 GB of GPU memory — orders of magnitude less than the 15+ GB the unfused version would require. The hypothesis that BlockMask breaks compilation is falsified.

Why This Message Matters: The Art of Debugging by Elimination

This message is a textbook example of the scientific method applied to systems debugging. The assistant had a hypothesis — "the BlockMask with closures is breaking compilation" — and designed an experiment to test it. The experiment is carefully controlled:

1. Isolation: The test uses standalone tensors, not the full model, eliminating confounding factors like gradient checkpointing, optimizer states, or other model components.
2. Representative inputs: The tensor dimensions (Q_LEN=2048, KV_LEN=10240) mirror the actual DFlash training scenario where 128 anchors × 16 block_size = 2048 query tokens and ~10240 KV tokens.
3. Real BlockMask: The test uses create_block_mask with a causal_mask function, which is structurally similar to the DFlash mask pattern (though simpler — the actual DFlash mask has document isolation and prefix constraints).
4. Memory measurement: By resetting peak memory stats before each phase, the test cleanly separates forward and backward memory usage.
5. Clear success criterion: Backward peak < 5 GB indicates a fused kernel; anything above means the dense fallback is being used. The result forces a shift in the assistant's mental model. If the compilation works in isolation, the bug must be in how the compiled function is integrated into the training loop. The subsequent messages (7895-7897) confirm this: the module-level _compiled_flex_attention is correctly defined and marked as compiled. The problem must be elsewhere — perhaps a stale import, a different flex_attention call path (e.g., from HuggingFace Transformers), or the compiled function being shadowed or redefined at runtime.

Input Knowledge Required

To fully understand this message, one needs knowledge of several interconnected domains:

• PyTorch's flex_attention API: Understanding that flex_attention is a higher-order operator that supports custom score_mod and mask_mod functions, and that it has two execution paths — a fused Triton kernel (when compiled) and a dense fallback (when not compiled).
• torch.compile and Triton: Knowing that torch.compile uses TorchDynamo to trace the computation graph and generate fused Triton kernels, and that the compilation is architecture-specific (sm_120 for Blackwell).
• The BlockMask abstraction: Understanding that create_block_mask creates a tiled representation of the attention mask that flex_attention uses to skip entire blocks, enabling efficient sparse attention.
• GPU memory hierarchy: Appreciating that materializing a full attention score matrix of shape [32 heads × 2048 Q_len × 10240 KV_len] in float32 would require ~2.7 GB, and that the backward pass needs additional storage for softmax intermediates, making the unfused path prohibitively expensive.
• The DFlash training architecture: Knowing that the drafter uses Grouped Query Attention (GQA) with 32 query heads and 8 KV heads, and that the attention spans a packed sequence of up to 8192 tokens with document boundaries.

Output Knowledge Created

This message produces several critical pieces of knowledge:

1. Definitive proof: torch.compile(flex_attention) with a BlockMask produces a fused backward kernel on Blackwell GPUs, using only 0.15 GB for the backward pass.
2. Narrowed bug scope: The compilation failure is not caused by the BlockMask itself, but by something in how the compiled function is integrated into the training loop — potentially a stale import, a shadowed variable, or an alternative flex_attention call path.
3. Reproducible test case: The exact Python snippet serves as a regression test that can be run on any Blackwell node to verify that the flex_attention compilation path works.
4. Confidence in the approach: The fused kernel approach is validated, justifying continued investment in the torch.compile strategy rather than pivoting to alternative attention implementations (like manual chunking or flash_attn).

Assumptions and Potential Pitfalls

The assistant makes several assumptions in this test:

• The causal_mask is representative: The test uses a simple causal mask (q_idx >= kv_idx), while the actual DFlash mask is more complex (combining causal masking with document isolation and prefix blocks). If the compiler handles simple masks but fails on complex ones, the test result would be misleading.
• Single-device test generalizes to multi-GPU: The test runs on a single GPU (cuda:2), while the training uses two GPU pairs with data parallelism. The compilation behavior might differ under distributed contexts.
• Standalone tensors generalize to model parameters: The test uses randomly initialized tensors with requires_grad=True, not actual model parameters wrapped in nn.Module instances. The compiler might behave differently with module parameters due to gradient scaling, parameter freezing, or other module-level optimizations.
• The 5 GB threshold is diagnostic: The assistant uses bwd < 5 as the criterion for "fused." This is reasonable (the unfused backward uses ~15 GB), but it's a heuristic, not a guarantee that the fused kernel is being used.

The Broader Narrative: A Debugging Odyssey

This message sits within a larger arc of debugging that spans multiple chunks and segments. The assistant had already:

• Fixed six training script bugs (chunk 0)
• Diagnosed FLA Triton autotuner crashes on sm_120
• Cleared corrupted Triton disk caches
• Implemented sequential warmup to avoid autotuner race conditions
• Upgraded Triton from 3.6.0 to 3.7.0
• Added lazy compilation for flex_attention to avoid cache corruption The flex_attention OOM was the next frontier. The assistant's reasoning in message 7893 shows a deep understanding of PyTorch's internals — correctly identifying that FlexAttentionAutogradFunction dispatches to either a fused Triton backward or a dense backward depending on whether the forward was traced by the compiler. The speculation about BlockMask closures breaking compilation was reasonable but ultimately incorrect, and message 7894 provides the experimental evidence to correct that assumption.

Conclusion: The Power of a Single Test

Message 7894 is a masterclass in targeted debugging. In just six lines of output, it eliminates an entire class of hypotheses and redirects the investigation toward the actual integration issue. The fused backward at 0.15 GB — versus 17.85 GB unfused — represents a 119× memory reduction, which is the difference between a training run that OOMs on the first step and one that completes all six epochs.

The subsequent messages in the session reveal the true culprit: not the BlockMask, but a race condition in FLA's CachedAutotuner that corrupts self.nargs when two GPU pairs concurrently trigger the same autotuner instance. The assistant eventually pivots to a structural fix — running target model forward passes sequentially across GPU pairs — that avoids the concurrent autotuner calls entirely. But that fix is only possible because message 7894 first cleared the path by proving that the compiled flex_attention itself was sound.

In the high-stakes world of training large language models on bleeding-edge hardware, where a single OOM crash can waste hours of GPU time, the ability to isolate variables and run decisive experiments is invaluable. This message exemplifies that discipline: a clean hypothesis, a controlled test, and a result that changes everything.