Reduce diffusion-expert launch overhead with bounded CUDA graphs

[brandonpelfrey]

Summary

This adds the same opt-in bounded CUDA-graph replay to Alpamayo R1 that is proposed for Alpamayo 1.5 PR #36 and Alpamayo 2 PR #12.

Only the repeated diffusion-expert forwards in trajectory sampling are captured. The autoregressive VLM rollout and training/autograd paths remain eager. The public API is consistent across the model repositories:

model.enable_diffusion_expert_cuda_graph(
    max_batch_size=...,
    max_graphs=4,
)
model.diffusion_expert_cuda_graph_stats

Keeping this capability on the model also lets downstream inference adapters and model subclasses enable the optimization without carrying a separate CUDA-graph implementation.

The Alpamayo 1.5 PR contains the corresponding end-to-end policy-workload results and pipeline-level analysis. This PR does not assume that the exact system-level improvement transfers unchanged to every Alpamayo R1 workload.

Variable batch sizes

Graphs are captured lazily for exact tensor and prompt-KV signatures. Different observed batch sizes therefore receive separate graphs automatically; callers do not need to pad to max_batch_size.

Up to max_graphs signatures are retained. Once that bound is reached, unseen signatures use eager execution. max_batch_size is the accepted upper bound, rather than a padding target.

Change

• Add enable_diffusion_expert_cuda_graph(max_batch_size=..., max_graphs=...).
• Capture and replay exact tensor and prompt-KV signatures lazily.
• Serialize a complete diffusion-sampling operation over graph-owned static buffers.
• Leave the caller's prompt cache unchanged.
• Bound retained signatures and use eager execution for unsupported or excess signatures.
• Preserve eager behavior for training, autograd-enabled execution, causal attention, and other unsupported calls.
• Support both standard and navigation classifier-free guidance paths.
• Document opt-in use, dynamic batches, statistics, and retained CUDA memory.

The feature is off by default and is enabled only after model placement and eval().

Validation

• CUDA graph tests: 7 passed
• Ruff and formatting checks passed.
• Python compilation and pytest collection passed.
• Coverage includes eager parity, dynamic batches and prompt lengths, bounded signature caching, eager fallback, prompt-cache preservation, replay across inference/no-grad modes, and concurrent CUDA work during thread-local capture.

Risk

Each captured signature retains static input and prompt-KV buffers, so enabling the feature increases CUDA memory use. max_graphs bounds retained state, and unsupported or excess signatures use the eager path. Default behavior is unchanged.