[4/N] extend CUDA-specific graph API to pytorch's accelerator graph backend, to make graph supported on Intel GPU (XPU)

.claude/skills/add-mstar-model/SKILL.md

@@ -1,6 +1,6 @@
 ---
 name: add-mstar-model
-description: End-to-end workflow for adding a model to M* — reconnaissance on the reference implementation, decomposing it into walks and nodes, deciding resources vs submodules, registration, eager parity against an oracle, batching and CUDA graphs, human verification of media output, competitor benchmarking and performance tuning. Use when porting a reference or Hugging Face model, or adding a model package, registry entry, serving config or model-specific API adapter.
+description: End-to-end workflow for adding a model to M* — reconnaissance on the reference implementation, decomposing it into walks and nodes, deciding resources vs submodules, registration, eager parity against an oracle, batching and accelerator graphs, human verification of media output, competitor benchmarking and performance tuning. Use when porting a reference or Hugging Face model, or adding a model package, registry entry, serving config or model-specific API adapter.
 ---
 
 # Adding a model to M*
@@ -110,7 +110,7 @@ The reference is not ground truth either: measure both it and the port against a
 
 Make the weight loader raise if any parameter goes unfilled. A silently half-loaded model produces plausible output instead of an error.
 
-Record the oracle once and keep the parity suite; rerun it after every performance commit, and add a parity test per performance commit. Do not add CUDA-graph configs while chasing parity — capture makes every discrepancy harder to localize and the two kinds of debugging do not mix.
+Record the oracle once and keep the parity suite; rerun it after every performance commit, and add a parity test per performance commit. Do not add accelerator graph configs while chasing parity — capture makes every discrepancy harder to localize and the two kinds of debugging do not mix.
 
 The oracle and the server are two sides of an A/B, so they must never share a GPU at the same time: run them on different GPUs when nothing is being timed, and back to back when something is. Record every oracle input you expect to need in one session up front, so you don't stop the server to re-record later.
 
@@ -126,9 +126,9 @@ Each of these is a one-line convention that a reference implementation states so
 - **Key per-request submodule state by `fwd_info.rid_handle` in `prepare_inputs`**, not `request_id`: the forward and the engine's cleanup use the handle.
 - **The conductor merges each pass's `step_metadata` into the request's `metadata.kwargs`.** A model whose walks send *different* step metadata (multi-stage, multi-resolution) and also keeps request-level facts in kwargs under the same names has those facts overwritten by the previous walk's values: a stage-1 half-resolution `height` becomes the request's height for stage 2. Keep request-level facts under keys step metadata never uses (LTX-2.5 nests them under `kwargs["request"]`).
 
-## Phase 5 — Batching and CUDA graphs
+## Phase 5 — Batching and accelerator graphs
 
-Add `can_batch` / `forward_batched` and the graph declarations. If only part of the forward is capturable, use piecewise graphs (`PiecewiseBatchedConfig` / `PiecewisePackedConfig`) rather than abandoning capture for the node. `unpack_packed_output` and `postprocess` are the hooks for work on the tail end of a graphed forward.
+Add `can_batch` / `forward_batched` and declare capture recipes through `get_accelerator_graph_configs` / `get_piecewise_accelerator_graph_configs`. If only part of the forward is capturable, use piecewise graphs (`PiecewiseBatchedConfig` / `PiecewisePackedConfig`) rather than abandoning capture for the node. `unpack_packed_output` and `postprocess` are the hooks for work on the tail end of a graphed forward.
 
 Take the shape ledger from Phase 1 and give **every varying dimension exactly one treatment**: its own bucket, normalized before the graph, or eager behind an explicit gate. Declare varying dims in the graph config — engine code that infers them from tensor shapes is model-specific and gets rejected. Normalizing an input before a graph changes numerics, so the parity test comes before shipping it. Buckets are not free: each costs startup through a Dynamo re-trace, and batch buckets are geometric. Per-request facts the engine needs for bucketing travel via step metadata into the key-info hook; model-side branching on request shape does not reach the graph selector.
 

.claude/skills/add-mstar-model/references/model-shape-map.md

@@ -8,7 +8,7 @@ Adapted from the map drafted in PR [#253](https://github.com/mstar-project/mstar
 
 Answer these before looking at the table, and write the answers down:
 
-1. Which compute stages can be independent graph nodes, and which are better fused for a larger CUDA graph?
+1. Which compute stages can be independent graph nodes, and which are better fused for a larger accelerator graph?
 2. Which walks and loops traverse those nodes?
 3. Which tensors cross node, walk, iteration or partition boundaries?
 4. Which persistent state needs engine admission, stable storage, capacity, dependency ordering or cleanup?
@@ -25,7 +25,7 @@ Classify every persistent value from its *behaviour* before picking reference co
 | Audio tower feeding an AR language model | `higgs_audio` | Ragged/cacheless tower attention, then AR resources | Verify audio preprocessing and tower output before language decode |
 | Two AR stages connected as a live stream | `qwen3_omni`, `qwen3_tts` | Each stage declares its own resource keys; async partitions and stream buffers connect them. qwen3-omni also fuses talker + code predictor into one node for capture size | Validate each stage directly, then one-token chunk routing |
 | Conditional multimodal walks with CFG branches | `bagel` | Walk-specific branches share declared resources; step declarations encode forks and combined labels | Start with uncaptured single-request walks per modality |
-| Image/video flow transformer (text encoder → DiT loop → VAE) | `flux2_klein`, `z_image` on the DiT scaffold (`mstar/model/components/diffusion/`) | Start from the scaffold's `DenoiseLoopSubmodule`, which supplies equal-bucket batching, per-step CUDA graphs, ragged attention and the async overshoot veto; see the [diffusion-scaffold skill](../../diffusion-scaffold/SKILL.md) | Share initial noise with the reference and compare the velocity at every step |
+| Image/video flow transformer (text encoder → DiT loop → VAE) | `flux2_klein`, `z_image` on the DiT scaffold (`mstar/model/components/diffusion/`) | Start from the scaffold's `DenoiseLoopSubmodule`, which supplies equal-bucket batching, per-step accelerator graphs, ragged attention and the async overshoot veto; see the [diffusion-scaffold skill](../../diffusion-scaffold/SKILL.md) | Share initial noise with the reference and compare the velocity at every step |
 | Diffusion or flow loop over stateless compute, predating the scaffold | `cosmos3`, `wan22` | Loop state normally rides graph edges or request state; declare attention resources only where planning actually needs them. These know their iteration count at ingestion, so they exit early in `prepare_inputs` for the async worker's extra step | Share initial latents with the reference and compare every denoise step |
 | Vision-language-action policy | `pi05` | Vision encoder plus language KV resources and an action flow loop; no sampler when actions come from flow matching | Verify the vision/language prefix, then every action denoise step |
 | World-model rollout with growing AR context | `vjepa2` | One-shot predictor may need no resource at all; action-conditioned rollout declares KV and attention for persistent history | Compare each rollout iteration and the loop-back signal |

.claude/skills/benchmarking/SKILL.md

@@ -173,7 +173,7 @@ Comparing against vLLM or SGLang has its own ways to go wrong.
 
 ## Three traps that produce fake regressions
 
-**1. Kernel selection changes across server processes.** Compilation uses `mode="max-autotune-no-cudagraphs"` (`mstar/engine/cuda_graph_runner.py`). Inductor benchmarks candidate kernels at compile time and keeps the fastest by measured wall time — and those measurements are noisy, so two processes running *identical* code can select different kernels. That changes float accumulation order and perturbs logits in the low bits. Greedy stages are unaffected (argmax rarely flips), but any sampled stage can flip a token, and one flipped token changes the whole downstream generation. Observed on qwen3-omni: identical byte totals across repeats within one process, but mean audio duration 55.6s vs 59.0s across two processes of the same code. It moves throughput too, not only sampled output: the same Whisper code over four restarts gave 368–455 RTFx at c=32, while c=1 stayed within ~1%. So don't attribute a cross-process difference to the branch; sample several processes per side, and treat a high-concurrency delta under ~10% as noise until it survives several restarts. When the change can be switched at runtime, run both arms on one build rather than two. Setting `TORCHINDUCTOR_CACHE_DIR` to a shared persistent path should pin selection after the first compile (worth doing, unverified).
+**1. Kernel selection changes across server processes.** Compilation uses `mode="max-autotune-no-cudagraphs"` (`mstar/engine/accelerator_graph_runner.py`). Inductor benchmarks candidate kernels at compile time and keeps the fastest by measured wall time — and those measurements are noisy, so two processes running *identical* code can select different kernels. That changes float accumulation order and perturbs logits in the low bits. Greedy stages are unaffected (argmax rarely flips), but any sampled stage can flip a token, and one flipped token changes the whole downstream generation. Observed on qwen3-omni: identical byte totals across repeats within one process, but mean audio duration 55.6s vs 59.0s across two processes of the same code. It moves throughput too, not only sampled output: the same Whisper code over four restarts gave 368–455 RTFx at c=32, while c=1 stayed within ~1%. So don't attribute a cross-process difference to the branch; sample several processes per side, and treat a high-concurrency delta under ~10% as noise until it survives several restarts. When the change can be switched at runtime, run both arms on one build rather than two. Setting `TORCHINDUCTOR_CACHE_DIR` to a shared persistent path should pin selection after the first compile (worth doing, unverified).
 
 **2. The first request of a new shape pays compilation, not inference.** That cost lands entirely on one request and reads like a tail regression. Concrete case: in the VBench I2V set, `req7` is the first landscape image (1024x672; req0-6 are portrait) and takes ~23s where its neighbours take ~7s. `req9` is also landscape and is *not* slow, because req7 already paid. Re-running confirms it — req7 drops from 23.8s to 6.7s. When a p95 or p99 moves on a generation benchmark, check whether one request is carrying a compile, and compare compile cost and steady-state cost separately: a branch can be slower to compile and faster to run.
 

.claude/skills/diffusion-scaffold/SKILL.md

@@ -1,6 +1,6 @@
 ---
 name: diffusion-scaffold
-description: Building an image/video flow or diffusion model on M*'s DiT scaffold (mstar/model/components/diffusion/, DenoiseLoopSubmodule) — bucket keys, schedules and seeding, multi-stage recipes, ragged attention spans, compile and per-bucket CUDA graphs, and the async overshoot veto. Use when porting a DiT-style model, writing a DenoiseLoopSubmodule subclass, or changing the scaffold itself.
+description: Building an image/video flow or diffusion model on M*'s DiT scaffold (mstar/model/components/diffusion/, DenoiseLoopSubmodule) — bucket keys, schedules and seeding, multi-stage recipes, ragged attention spans, compile and per-bucket accelerator graphs, and the async overshoot veto. Use when porting a DiT-style model, writing a DenoiseLoopSubmodule subclass, or changing the scaffold itself.
 ---
 
 # The DiT scaffold
@@ -13,7 +13,7 @@ Read the [add-mstar-model skill](../add-mstar-model/SKILL.md) for the port as a
 
 ## The bucket key
 
-`bucket_key_for(fwd_info)` returns everything that has to match for two requests to share a forward. It is also the CUDA-graph bucket key and the key of the per-bucket layout cache (`build_layout`: rotary tables and the like, built once and never evicted because a captured graph reads them at fixed addresses).
+`bucket_key_for(fwd_info)` returns everything that has to match for two requests to share a forward. It is also the accelerator graph bucket key and the key of the per-bucket layout cache (`build_layout`: rotary tables and the like, built once and never evicted because a captured graph reads them at fixed addresses).
 
 - **It is not just shape.** The latent grid and the text length belong in it, but so does anything that changes the forward: whether guidance is on, the conditioning layout, which recipe stage is running.
 - **Too fine a key means no batching, too coarse means wrong output.** Two requests that differ in something the key leaves out are stacked into one forward and one of them is computed with the other's layout. A row-independence test (each row of a batched forward equals that request run alone) catches this.

.claude/skills/engine-resources/SKILL.md

@@ -1,6 +1,6 @@
 ---
 name: engine-resources
-description: Pitfalls when writing or changing an engine Resource under mstar/engine/resources/ — double-buffering against the async worker, pre-plan, padding rows inside CUDA-graph capture buckets, and attention-backend (FlashInfer, FlashAttention) quirks. Use before adding a new resource kind, changing how a resource stages buffers in plan(), or swapping an attention kernel.
+description: Pitfalls when writing or changing an engine Resource under mstar/engine/resources/ — double-buffering against the async worker, pre-plan, padding rows inside accelerator graph capture buckets, and attention-backend (FlashInfer, FlashAttention) quirks. Use before adding a new resource kind, changing how a resource stages buffers in plan(), or swapping an attention kernel.
 ---
 
 # Writing an engine resource
@@ -35,7 +35,7 @@ KV pages and recurrent (GDN/Mamba) state slots are two forms of the same per-req
 
 ## Capture buckets
 
-- Batch buckets are geometric (`DEFAULT_CAPTURE_BATCH_SIZES` in `cuda_graph_runner.py`: 1, 2, 4, …, 64), so a replay usually runs wider than the live batch.
+- Batch buckets are geometric (`DEFAULT_CAPTURE_BATCH_SIZES` in `accelerator_graph_runner.py`: 1, 2, 4, …, 64), so a replay usually runs wider than the live batch.
 - **A padded row must be well-formed, not zeroed.** The replay always runs the bucket's full width. Point padding at `SINK_PAGE` (see `KVPlanState.copy_`) and give it a length of at least 1 unless the target kernel has handling for zero-length rows.
 - **Take the step's width from the bucket when captured, and from the live batch otherwise.** Eager buffers are reused and sized to the widest batch seen. Padding an eager step out to the buffer width makes the index tensors describe query rows that `q` does not have, and the kernel reads off the end. That surfaces as an asynchronous illegal memory access far from the cause, so a cheap `q.shape[0] != expected` assertion at the call site pays for itself. A warmup at higher concurrency followed by a timed batch of 1 is what exposes this.