[GH-ISSUE #152] Proposal: `xpu` backend for `sol_attn` (Intel Arc / Xe-HPG), VSA subset first #48

New issue

Open

opened 2026-10-26 22:19:40 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 22:19:40 +00:00

Owner

Originally created by @AIVirtuoso on GitHub (Sep 3, 2026).
Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/152

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:

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.
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.
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:

Would you take a backend that is torch-level only as a first PR, or do you want backends/<name> to mean compiled kernels? If the latter I would put stage 1 under eager behind a device check instead.
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?
Is the VSA recipe surface in sol_attn (topk_ratio, tail=False, block_len, coarse_gate) considered stable enough to implement against now?
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.

Originally created by @AIVirtuoso on GitHub (Sep 3, 2026). Original GitHub issue: https://github.com/Comfy-Org/comfy-kitchen/issues/152 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/<name>` 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.

bsakaguchi referenced this issue

2026-10-26 22:20:24 +00:00

[PR #48] [MERGED] ci: add CLA Assistant workflow #102

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

Comfy-Org/comfy-kitchen#48

No description provided.

Rows
Columns

[GH-ISSUE #152] Proposal: xpu backend for sol_attn (Intel Arc / Xe-HPG), VSA subset first #48

[GH-ISSUE #152] Proposal: `xpu` backend for `sol_attn` (Intel Arc / Xe-HPG), VSA subset first #48