[acceleration] Model-agnostic acceleration plugin layer (torch.compile / attention backend / feature caching)

.agents/knowledge/architecture.md

@@ -92,7 +92,7 @@ Stage 6: Policy Optimization
 
 ## Registry System
 
-All three registries map string keys → lazy import paths. Resolution: registry lookup → fallback to direct Python path → dynamic import. See `trainers/registry.py`, `models/registry.py`, `rewards/registry.py` for implementation.
+All four registries map string keys → lazy import paths. Resolution: registry lookup → fallback to direct Python path → dynamic import. See `trainers/registry.py`, `models/registry.py`, `rewards/registry.py`, `acceleration/registry.py` for implementation.
 
 ### Registered Components
 
@@ -147,13 +147,23 @@ All three registries map string keys → lazy import paths. Resolution: registry
 | `hpsv2` | `HPSv2RewardModel` | Pointwise |
 | `qwen_image_bench` | `QwenImageBenchRewardModel` | Pointwise |
 
+**Accelerators** (`acceleration/registry.py`):
+| Key | Class | Safety | Stage | Notes |
+|-----|-------|--------|-------|-------|
+| `attention_backend` | `AttentionBackendAccelerator` | lossless | both | Sets the diffusers attention backend on every transformer (requires a `backend` param). Listed as a `shared` entry (before `torch_compile`); this is the single code path for backend selection (the old `BaseAdapter._set_attention_backend` and the `model.attn_backend` knob were both removed — a config still setting `model.attn_backend` fails fast). Bagel forces flash_attention_2 at load and does not use it. |
+| `torch_compile` | `CompileAccelerator` | lossy | both | `torch.compile` of the shared transformer (regional/full); applied in-place after `post_init` so checkpoint keys / param identity stay stable. Marked `lossy` because it is applied symmetrically but is **not bit-exact across rollout vs training** (Inductor's grad/no-grad graph split → intermittent ~1e-5 on-policy residual, within `clip_range`); still allowed on coupled algos, validator warns. |
+| `diffusers_cache` | `DiffusersCacheAccelerator` | lossy | rollout | Diffusers `CacheMixin` feature caching (first_block/faster/pyramid/taylorseer/magcache). |
+
+Configured via the `acceleration:` block (`hparams/acceleration_args.py`): two ordered lists of `{name, params}` entries — `shared` (persistent `stage='both'` accelerators applied to rollout and training) and `rollout` (Stage-3 only). **List order is application order**: `shared` entries run their `setup()` in order (so `attention_backend` must precede `torch_compile`), and `rollout` entries nest their `rollout_context()` in order. The `acceleration/validator.py` enforces that a **lossy `rollout`** accelerator runs only on `decoupled`/`distillation` trainers (each trainer declares a `paradigm`), preserving train-inference consistency (constraint #7). A **lossy `stage='both'` accelerator in the `shared` slot** (e.g. `torch_compile`, applied symmetrically but not bit-exact across stages) is allowed on any paradigm but the validator **warns** on coupled trainers that the on-policy ratio will be ≈1, not exactly 1. Off by default.
+
 ---
 
 ## Extension Points
 
 - **New model adapter**: `guidance/new_model.md`, skill `/ff-new-model`, conventions `topics/adapter_conventions.md`
 - **New reward model**: `guidance/rewards.md`, skill `/ff-new-reward`
 - **New algorithm**: `guidance/algorithms.md`, skill `/ff-new-algorithm`
+- **New accelerator**: subclass `acceleration/abc.py::BaseAccelerator` (declare `safety`/`stage`), register in `acceleration/registry.py`
 
 ---
 

.gitignore

@@ -139,3 +139,8 @@ dataset/**/*.mov
 *.zip
 *.gz
 *.tar
+
+# ====================
+# Cursor scratch plans (design/roadmap docs; keep .cursor/rules tracked)
+# ====================
+.cursor/plans/

AGENTS.md

@@ -63,6 +63,7 @@ See `.agents/knowledge/architecture.md` "Module Dependency Graph" for full detai
 | `guidance/algorithms.md` | All 9 algorithms (GRPO, GRPO-Guard, DPPO, DPO, DGPO, DiffusionNFT, AWM, CRD, DiffusionOPD) deep dive |
 | `guidance/rewards.md` | Reward system design, custom model creation |
 | `guidance/new_model.md` | Step-by-step model adapter integration |
+| `guidance/acceleration.md` | Acceleration plugin layer (compile, attention backend, feature caching) |
 
 ## Available Skills
 

README.md

@@ -15,12 +15,14 @@ git submodule update --init
 pip install -e ./diffusers
 ```
 
-* **[2026-02-01]** Support for multiple **Attention Backends**! You can now optimize memory and speed by setting the `attn_backend` parameter in your config:
+* **[2026-02-01]** Support for multiple **Attention Backends**! Attention-backend selection now lives in the unified `acceleration:` block (the old `model.attn_backend` knob was removed), where it can be combined with `torch.compile` and feature caching — applied in list order:
 ```yaml
-  model:
-      attn_backend: "flash" # Options: "native", "xformers", "flash_hub", "_flash_3_hub", "_flash_3_varlen_hub"
+  acceleration:
+    shared:
+      - name: attention_backend
+        params: { backend: "flash" } # Options: "native", "xformers", "flash_hub", "_flash_3_hub", "_flash_3_varlen_hub"
 ```
-This experimental feature leverages `diffusers`'s `transformer.set_attention_backend`. Check the [official diffusers documentation](https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends#available-backends) for all available options.
+This experimental feature leverages `diffusers`'s `transformer.set_attention_backend`. Check the [official diffusers documentation](https://huggingface.co/docs/diffusers/main/en/optimization/attention_backends#available-backends) for all available options. See [`guidance/acceleration.md`](guidance/acceleration.md) for the full acceleration layer.
 > We recommend installing the `kernels` package (`pip install kernels`) and using `flash_hub`, `flash_varlen_hub`, `_flash_3_hub`, or `_flash_3_varlen_hub` to avoid the complexity and potential incompatibility of installing Flash-Attention directly.
 
 # 📕 Table of Contents

examples/awm/lora/flux1/default.yaml

@@ -30,7 +30,18 @@ model:
   model_type: "flux1"
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_hub'
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout
+# feature-caching is allowed.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
+#   rollout:                         # rollout-only (Stage 3); lossy -> decoupled/distillation only
+#     - name: diffusers_cache        # Options: diffusers_cache
+#       params: { policy: first_block, threshold: 0.08 }
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})

examples/awm/lora/flux2_klein_base/default.yaml

@@ -30,7 +30,18 @@ model:
   model_type: "flux2-klein"
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_hub' # Attention backend for training.
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout
+# feature-caching is allowed.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
+#   rollout:                         # rollout-only (Stage 3); lossy -> decoupled/distillation only
+#     - name: diffusers_cache        # Options: diffusers_cache
+#       params: { policy: first_block, threshold: 0.08 }
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})

examples/awm/lora/sd3_5/default.yaml

@@ -31,7 +31,18 @@ model:
   model_type: "sd3-5"
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_hub' # Attention backend for training.
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout
+# feature-caching is allowed.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
+#   rollout:                         # rollout-only (Stage 3); lossy -> decoupled/distillation only
+#     - name: diffusers_cache        # Options: diffusers_cache
+#       params: { policy: first_block, threshold: 0.08 }
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})

examples/crd/lora/sd3_5/default.yaml

@@ -32,7 +32,18 @@ model:
   model_type: "sd3-5"
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_hub' # Attention backend for training.
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. This is a decoupled trainer, so lossy rollout
+# feature-caching is allowed.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
+#   rollout:                         # rollout-only (Stage 3); lossy -> decoupled/distillation only
+#     - name: diffusers_cache        # Options: diffusers_cache
+#       params: { policy: first_block, threshold: 0.08 }
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{timestamp})

examples/grpo/full/qwen_image/default.yaml

@@ -29,7 +29,15 @@ model:
   model_type: "qwen-image"  # Options: flux1, flux1-kontext, flux2, qwenimage, qwenimage-edit
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_varlen_hub' # Attention backend for Qwen-Image Series, which uses masked attention with variable sequence length.
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe;
+# persistent `shared` accelerators run in both rollout and training.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_varlen_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})

examples/grpo/full/qwen_image_edit_plus/default.yaml

@@ -28,7 +28,15 @@ model:
   model_type: "qwen-image-edit-plus"
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_varlen_hub' # Attention backend for Qwen-Image Series, which uses masked attention with variable sequence length.
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe;
+# persistent `shared` accelerators run in both rollout and training.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_varlen_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})

examples/grpo/lora/bagel/i2i.yaml

@@ -38,7 +38,8 @@ model:
   model_type: "bagel"  # Model adapter key; see src/flow_factory/models/registry.py for options
   resume_path: null # Path to load previous checkpoint/lora adapter
   resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  attn_backend: 'flash' # Options: auto, flash, sdpa
+  # Bagel forces flash_attention_2 at load (pip install -e ".[bagel]"); it has no diffusers
+  # attention backend, so it takes no acceleration `attention_backend` entry.
 
 # Training Configuration
 train:

examples/grpo/lora/ltx2/t2av.yaml

@@ -33,7 +33,15 @@ model:
   model_type: "ltx2_t2av"
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null  # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_hub'  # Attention backend.
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe;
+# persistent `shared` accelerators run in both rollout and training.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})

examples/grpo/lora/ltx2/t2av_pickscore.yaml

@@ -31,7 +31,15 @@ model:
   model_type: "ltx2_t2av"
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null  # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_hub'  # Attention backend.
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe;
+# persistent `shared` accelerators run in both rollout and training.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})

examples/grpo/lora/qwen_image/default.yaml

@@ -30,7 +30,15 @@ model:
   model_type: "qwen-image"
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_varlen_hub' # Attention backend for Qwen-Image Series, which uses masked attention with variable sequence length.
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe;
+# persistent `shared` accelerators run in both rollout and training.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_varlen_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})

examples/grpo/lora/qwen_image_edit_plus/default.yaml

@@ -31,7 +31,15 @@ model:
   model_type: "qwen-image-edit-plus"
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_varlen_hub' # Attention backend for Qwen-Image Series, which uses masked attention with variable sequence length.
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe;
+# persistent `shared` accelerators run in both rollout and training.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_varlen_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})

examples/grpo/lora/sd3_5/default.yaml

@@ -33,7 +33,15 @@ model:
   model_type: "sd3-5"
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_hub' # Attention backend for training.
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe;
+# persistent `shared` accelerators run in both rollout and training.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})

examples/grpo/lora/sd3_5/geneval.yaml

@@ -30,7 +30,15 @@ model:
   model_type: "sd3-5"
   resume_path: null # Local path or HF repo id (e.g. 'owner/repo[/subdir][@rev]') for previous checkpoint/lora adapter
   resume_type: null # Options: lora, full, state. Null to auto-detect based on `finetune_type`
-  # attn_backend: '_flash_3_hub' # Attention backend for training.
+# Optional acceleration plugins (off by default). Applied in list order.
+# See guidance/acceleration.md. GRPO is coupled, so lossy rollout caching is unsafe;
+# persistent `shared` accelerators run in both rollout and training.
+# acceleration:
+#   shared:                          # persistent stage='both'; applied in list order
+#     - name: attention_backend      # set the attention backend first...
+#       params: { backend: _flash_3_hub }
+#     - name: torch_compile          # ...so the compiled graph captures it
+#       params: { mode: regional }   # regional (compile_repeated_blocks) | full
 
 log:
   run_name: null  # Run name (auto: {model_type}_{finetune_type}_{trainer_type}_{timestamp})