Eliminating Dynamic Operations for CUDA Graph Capture: The `select_anchors` Fix

In the ongoing effort to stabilize and accelerate a multi-GPU speculative decoding training pipeline, a single message from the assistant captures a pivotal moment of architectural insight. At message index 10328, the assistant identifies and begins to eliminate two remaining blockers in the dflash_model.py code that prevent CUDA graph capture—operations that are invisible to the casual reader but catastrophic for performance optimization. This message, though brief in its visible output (a single apply_patch call), represents the culmination of a long chain of reasoning about what it truly means to make a PyTorch training loop "capturable" by CUDA graphs.

The Context: A Pipeline Under Siege

To understand why this message matters, one must appreciate the engineering nightmare that preceded it. The DFlash training pipeline is a custom multi-GPU system that trains a speculative decoding "drafter" model against a larger "target" model. The pipeline uses Python threading to manage multiple target GPUs feeding hidden states to a drafter GPU, all within a single process. Throughout segment 56, this pipeline had been plagued by a cascade of failures: missing CUDA extensions caused the target model's GatedDeltaNet layers to fall back to slow PyTorch kernels; torch.compile(flex_attention) crashed from multi-threaded FX tracing race conditions; attempts to replace flex_attention with per-block batched SDPA failed due to memory overhead; and a per-thread execution lock proved insufficient to isolate compilation state.

The user's frustration was palpable: throughput remained stuck at ~12K tok/s with volatile GPU memory and low utilization. The root cause, as the assistant had diagnosed in the preceding messages ([msg 10321]), was that the drafter forward+backward path allocated and freed differently-sized tensors every step due to variable total_seq_len. This prevented CUDA graph replay, caused allocator churn, and created GIL contention across 12+ threads.

The proposed solution was ambitious: pad all hidden state batches to a fixed token_budget of 49,152 tokens, preallocate persistent GPU buffers, and replace all dynamic operations with fixed-shape equivalents. This would enable torch.compile(mode="reduce-overhead") to capture CUDA graphs that could be replayed without any allocation overhead.

The Message: Identifying Hidden Dynamism

The assistant's reasoning in message 10328 reveals a crucial realization:

For actual CUDA capture, there are two more blockers in dflash_model.py beyond padding: select_anchors() currently uses nonzero, randperm(valid_indices.numel()), and lengths.tolist(), all dynamic/non-capturable. I'm replacing anchor selection with fixed-shape random top-k over a fixed [token_budget] score vector, and replacing document-id construction with fixed-shape vectorized masks.

This is the moment where the assistant moves beyond the obvious (padding tensors to fixed sizes) to the subtle (eliminating operations whose output shapes depend on data). The padding fix alone, applied in messages [msg 10324] through [msg 10327], would have made the input tensors fixed-size, but the internal operations of the model would still produce variable-sized tensors, breaking graph capture.

The specific operations identified are:

nonzero: This operation returns the indices of non-zero elements, and its output size depends on how many non-zero elements exist. In select_anchors, it was used to find valid anchor positions based on the loss mask. With variable-length inputs, the number of valid positions changes every batch, producing a tensor of unpredictable size.
randperm(valid_indices.numel()): After finding valid indices, the code generated a random permutation whose size depended on the number of valid indices. Again, a dynamic output size.
lengths.tolist(): Converting a tensor to a Python list breaks the computational graph and introduces a Python-level dependency that CUDA graphs cannot capture. These operations are invisible to someone reading the training loop at a high level—they're just implementation details of how anchors are selected. But to torch.compile's inductor backend and to CUDA graph capture, they are graph breaks that force dynamic allocation and prevent the entire forward+backward from being recorded as a single replayable graph.

The Fix: Fixed-Shape Random Top-K

The assistant's solution is elegant: replace the dynamic anchor selection with a fixed-shape random top-k operation over a score vector of size [token_budget]. Instead of first finding valid positions and then randomly selecting among them, the new approach:

Creates a fixed-size score vector of length token_budget (always 49,152)
Masks out invalid positions (padding tokens) by setting their scores to a very negative value
Uses topk to select exactly num_anchors positions from this fixed-size vector The key insight is that topk over a fixed-size tensor always produces outputs of the same shape—[num_anchors]—regardless of how many valid positions exist. The masking ensures that invalid positions are never selected (their scores are pushed to -inf), while the output tensor shape remains constant. Similarly, document-id construction is replaced with fixed-shape vectorized operations that produce tensors of predetermined size, with padding positions getting sentinel values (like doc_id=-1) that the attention mask correctly excludes.

The Reasoning Process: Incremental Discovery

What makes this message fascinating is the reasoning trajectory visible in the assistant's thinking. The message begins with the assistant evaluating the capture complexity:

I'm considering the capture process, and it seems like the block mask needs to be set, even though the flex kernel relies on dynamic inputs. To achieve full capture, fixed BlockMask addresses might be necessary, which adds complexity.

This shows the assistant working through the implications of CUDA graph capture at a deep level. The flex_attention kernel, which uses a BlockMask to implement block-sparse attention with document boundaries, normally receives dynamically-computed block masks. For full graph capture, even the block mask's memory addresses must be fixed—the mask content can change, but the tensors holding it must be preallocated and reused via copy_().

The assistant then pivots to the concrete implementation:

I'm thinking of executing the plan incrementally, starting with patching the select_anchors and then creating the anchor block mask modification.

This reveals a deliberate, incremental strategy. Rather than attempting all changes at once and debugging the resulting chaos, the assistant identifies the minimal set of changes needed and applies them one at a time. The select_anchors fix is the first of these because it's a self-contained function with clear inputs and outputs—a safe place to start.

Assumptions and Knowledge Required

To understand this message, one needs significant background knowledge:

CUDA graph capture semantics: That torch.compile(mode="reduce-overhead") works by recording CUDA operations into a graph that can be replayed, and that this requires all tensor shapes and memory addresses to be fixed across replays.
The difference between dynamic and static operations: Why nonzero and randperm produce variable-sized outputs while topk over a fixed-size input produces fixed-size outputs.
The DFlash architecture: That select_anchors is called during the drafter forward pass to choose which token positions will serve as "anchor" points for the block-sparse attention pattern, and that these anchors determine which KV pairs are attended to.
The token budget concept: That the pipeline packs multiple documents into a single batch up to a maximum of 49,152 tokens, and that this maximum is now being used as the fixed size for all padded tensors. The assistant also makes a critical assumption: that replacing dynamic anchor selection with fixed-shape top-k produces equivalent training behavior. This is reasonable because the original randperm over valid indices is mathematically equivalent to selecting num_anchors random positions from the set of valid positions, and the top-k approach with random scores achieves the same distribution. However, this assumption would need to be validated by checking that the loss and accuracy metrics remain unchanged after the patch.

Output Knowledge Created

This message creates several important outputs:

A patched select_anchors function that is CUDA-graph-compatible, replacing the dynamic implementation with a fixed-shape one.
A documented understanding that padding inputs alone is insufficient for CUDA graph capture—internal operations must also be audited for dynamic outputs.
A template for further fixes: The same pattern (identify dynamic ops, replace with fixed-shape equivalents using masking) can be applied to other parts of the codebase.
Evidence of the incremental strategy: The assistant is deliberately working through blockers one at a time, starting with the most self-contained changes.

The Broader Significance

This message represents a microcosm of the entire segment's struggle. The DFlash training pipeline is pushing the boundaries of what PyTorch's compilation stack can do in a multi-threaded, multi-GPU environment. Every layer of abstraction—Python threading, the CUDA caching allocator, torch.compile, flex_attention, CUDAGraph Trees—introduces potential failure modes. The assistant is methodically working through them, discovering that each "fix" reveals new blockers.

The select_anchors fix is particularly instructive because it shows how CUDA graph capture changes the way one must think about algorithm implementation. Operations that are perfectly natural in eager mode (find valid indices, shuffle them, pick the first N) become graph breaks in compiled mode. The fix requires rethinking the algorithm in terms of fixed-shape operations, using masking as a universal technique to handle variable data within fixed-size tensors.

This is a lesson that extends far beyond this specific training pipeline. As more ML workloads adopt torch.compile and CUDA graph capture for performance, developers will need to internalize these patterns: preallocate everything, use masking instead of conditional logic, and ensure every operation produces outputs of deterministic shape. Message 10328 is a concrete example of this mindset shift in action.

Conclusion

Message 10328 is deceptively small—a single patch to a single function—but it encapsulates a profound engineering insight: that CUDA graph capture is not just about padding inputs to fixed sizes, but about auditing every operation in the computational graph for dynamic behavior. The assistant's identification of nonzero, randperm, and tolist() as blockers, and its replacement of them with fixed-shape alternatives, demonstrates a deep understanding of how PyTorch's compilation stack works and what it requires for optimal performance. This message is a turning point in the segment: the assistant has now identified all the major sources of dynamism in the drafter forward path and is systematically eliminating them, paving the way for the CUDA graph capture that will ultimately stabilize the training pipeline.

Eliminating Dynamic Operations for CUDA Graph Capture: The select_anchors Fix