Proposal: `xpu` backend for `sol_attn` (Intel Arc / Xe-HPG), VSA subset first

FastH3 (FastVideo's 4-step VSA distillation of MiniMax-H3) is arriving in ComfyUI through Comfy-Org/ComfyUI#15958, with the sparse attention going through `sol_attn` (#117, released in v0.2.32). Today that leaves every non-CUDA/HIP device with no way to run a VSA-trained checkpoint at all, and I would like to propose an `xpu` backend that closes the gap for Intel Arc, in two stages. Before writing code I want to check the shape of it with you.

**Why an eager fallback does not cover it.** `backends/eager/sol_attn.py` is registered for all devices, but it forms the full `(B, H, T, T)` fp32 score matrix (`s_tok = qh @ kh.transpose(-1, -2)`). At MiniMax-H3 sizes (56 heads, T in the tens of thousands) that is 160 GB and up, so it is a test oracle, not a runtime path. The test node attached to #117 gates on `q.device.type == "cuda"` and calls `backends.cuda.sol_attn_chunked` directly, and its VSA-mode override is deliberately dense-only, so on any other device a VSA checkpoint silently runs dense attention without the gate branch and produces noise. That is the correct choice for the node; it just means non-CUDA users cannot get there from here.

**What the target hardware has, measured on one Arc A770 (Xe-HPG, 16 GB).** From Triton's device query: `has_subgroup_matrix_multiply_accumulate: True` (XMX present), `has_subgroup_2d_block_io: False`, `local_mem_size: 65536`. The consequences for kernel choice, all measured rather than assumed:

| path | `2x16x4056x128` bf16, forward | note |
|---|---|---|
| oneDNN fused SDPA (what `F.scaled_dot_product_attention` runs on XPU) | 12.4 ms, ~22 TFLOPS | shipping config; ~1.6x faster with the config proposed in uxlfoundation/oneDNN#5808 |
| Triton FA2, best of 6 configs | 6454 ms | correct, but the lowering has no 2D block loads to lean on |
| `flex_attention` + `torch.compile` | out of resources | |

So a Triton port of the sparse kernel is not a route on this architecture, and porting the CUDA `sol_attn` kernels (CuTe / `mma`, INT8 Sage-style) is not either. The one performant attention kernel on Xe-HPG is oneDNN's, and it already runs at head_dim 128 in bf16.

**Stage 1, the proposal I would put up first: a torch-level `backends/xpu/sol_attn.py` that never materialises `T x T`.** Same signature and semantics as the eager reference, registered with `devices={"xpu"}`, initially for the VSA subset only (`topk_ratio > 0`, `tail=False`, `block_len`, `coarse_gate`; the tau / pooled-tail mode can follow). Sketch:

1. Block means and the routing scores are computed at block granularity: `colmean` is linear in the query rows, so `centroid @ kc^T` gives the eager reference's `(B, H, N, N)` routing matrix exactly (up to rounding) without the `(B, H, T, N)` intermediate. Top-k per query block, diagonal and sink blocks forced exact, as in eager.
2. The fine stage iterates over chunks of query blocks: gather the selected key/value blocks into contiguous `(chunk, k*64, D)` tensors and run `F.scaled_dot_product_attention`, which dispatches to oneDNN's fused kernel. `tail=False` makes this exactly the VSA fine stage; `block_len` becomes a key mask on the padded rows.
3. The gated coarse branch, `gate * softmax(q_mean k_mean^T) v_mean`, is a small batched matmul on the block means and is added per block.

Cost: FLOPs are the sparse ones, but the gather pays memory traffic proportional to `T * keep_ratio * D` per query block, so I expect attention time in the neighbourhood of dense, not below it, at the sizes H3 uses. The point of stage 1 is that VSA checkpoints run and produce video on XPU; FastH3's 4-step schedule is the larger part of its speedup and needs nothing from the kernel. I would validate it against the eager oracle at small `T` with the existing `tests/test_sol_attn.py` invariants (B>1, ragged tails, `block_len`, `coarse_gate`, strided inputs).

One XPU-specific wrinkle I already know about and would handle inside the backend: `torch.topk` on XPU (torch 2.13) crashes above 2048 elements in bf16 and fp32 while fp16 is fine, so the block top-k would run in fp16 or via sort.

**Stage 2, outside this repo: a block-index variant of oneDNN's SDPA microkernel.** Looking at `src/gpu/intel/sdpa/micro.cl` on oneDNN `main`: each workgroup owns one query tile and walks the keys in a flat loop, `for (int k0 = 0; k0 < k0end; k0 += ugemm_kq_wg_tile_m)`, and the causal path already clamps `k0end` per query tile. Block-sparse attention is that loop driven by an index array instead of a counter, with online softmax, XMX use and SLM layout untouched. For the Xe-HPG head_size-128 records the query tile is 64 rows and the key tile 64 (#5808 config) or 128 (shipping), which lines up with VSA's 64-token cubes. That is the path to a real sparsity win on this hardware; it is a oneDNN contribution and a longer conversation, and stage 1 gives it a correctness oracle. I mention it so the backend's interface is designed with a fused kernel behind it in mind rather than as a permanent gather path.

**Registry changes stage 1 would need**, all small: a `backends/xpu` package imported when `torch.xpu.is_available()` (the same pattern as `hip` in `comfy_kitchen/__init__.py:41-47`), `"xpu"` at the front of the priority list on such a host, a column in the README capability matrix, and `FunctionConstraints` entries with `default_devices={"xpu"}` (the dataclass default is `{"cuda", "cpu"}`). Nothing compiled, so the pure-Python wheel carries it.

Questions, each answerable in a line:

1. Would you take a backend that is torch-level only as a first PR, or do you want `backends/` to mean compiled kernels? If the latter I would put stage 1 under `eager` behind a device check instead.
2. Is registry-level `sol_attn` staying a supported public entry point next to `sol_attn_chunked` (#150), so a backend implementing only `sol_attn` is useful to the H3 integration?
3. Is the VSA recipe surface in `sol_attn` (`topk_ratio`, `tail=False`, `block_len`, `coarse_gate`) considered stable enough to implement against now?
4. Are the `tests/test_sol_attn.py` cases device-parametrisable, or should the xpu backend ship its own oracle comparison?

Happy to sign off DCO and put the stage 1 PR up once the shape is agreed. If measurements of the oneDNN kernel on Arc would help decide any of this, I can attach them in whatever form is convenient.

This issue was created with the help of an LLM.

Start by reading backends/eager/sol_attn.py and tests/test_sol_attn.py, then inspect the backend import pattern in comfy_kitchen/__init__.py:41-47, the README capability matrix, and FunctionConstraints. The stage 1 result should provide an XPU sol_attn entry point, registry support, and passing comparisons against the eager invariants for B>1, ragged tails, block_len, coarse_gate, and strided inputs.

Written by the indexing model from the issue text.