Add inference speed optimization

[YanjingLiLi]

This pull request introduces an optional performance optimization to the RFDiffusion3 inference engine by enabling selective compilation of key model submodules with torch.compile, and adds a new, memory-aware dense attention path for inference. These changes aim to reduce steady-state rollout time and improve inference efficiency, while maintaining compatibility with memory constraints and training requirements.

Inference engine enhancements:

• Added a new compile_model option to both the YAML config (rfdiffusion3.yaml) and the RFD3InferenceConfig/RFD3InferenceEngine classes, allowing users to opt-in to compiling hot diffusion submodules for up to 1.6x faster inference after a one-time warmup. [1] [2] [3] [4]

Attention layer improvements:

• Introduced a new dense attention path (dense_sdpa_pairbias_attention) for inference, which is automatically selected based on memory availability, device, and batch size. This path avoids materializing large gather tensors and leverages efficient CUDA and PyTorch SDPA routines. [1] [2]
• Added the use_dense_sdpa_pairbias decision function and logging to transparently select between dense and sparse attention paths, with environment variables to override behavior and detailed reporting for debugging.

Codebase and logging:

• Improved logging to report which attention path is taken, why, and under what shapes and memory conditions, aiding reproducibility and debugging.
• Minor import and code organization updates in attention.py to support the new features.

These changes are opt-in and maintain backward compatibility, ensuring that training and low-memory inference are unaffected unless explicitly enabled.

To enable both optimizations, use RFD3_DENSE_SDPA_ATTENTION=1 in the env and compile_model=True in the config.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3fa1e4f754

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Preserve duplicate neighbors in dense attention

For CUDA inference with L > k, this boolean mask collapses repeated entries in indices, but the original sparse softmax treats every occurrence as a separate key and therefore gives repeated keys additional weight. This occurs for inputs with more than three chains because get_sparse_attention_indices_with_inter_chain() pads with random indices and leaves many rows initialized with repeated zero indices, so the encoder and decoder can silently produce materially different diffusion outputs when the dense path is selected. Either retain multiplicity in the dense bias or reject the dense path when a row contains duplicates.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Reset compilation state when the model is discarded

When an engine with compile_model=True is reused after leaving its context manager, BaseInferenceEngine.__exit__() discards the trainer/model and resets initialized_, but compiled_ remains true. The next initialize() constructs a fresh, uncompiled model and then skips _compile_diffusion_submodules(), silently disabling the requested optimization for every subsequent run on that engine; the compilation flag needs to be reset whenever the underlying model is torn down.

Useful? React with 👍 / 👎.

[r-krishna]

lgtm, thanks Yanjing!