SYCL: pi0.5 on the Arc B390 iGPU (flash attention, fused nodes, INT4 DPAS GEMMs)

[hungho77]

Summary

π0.5 inference on the Intel Arc B390 iGPU (Core Ultra X7 358H, SYCL), for the SO101 11-task checkpoints in bf16, W8A8, W4A4 + o/down INT8 and W4A4. Denoise steps stay at 10 and n_action_steps is unchanged. Every shared-code change is opt-in, and only pi0.5 turns it on. Other models build the same graphs as before.

• Flash attention (sycl-tla) under --flash-attn
  • Prefix, vision and the action expert's cached-KV attention run on sycl-tla.
  • Split-KV for few query rows.
  • Patches for sycl-tla's partial-tile bugs.
  • tests/fmha_sycl_check.
• Vision
  • All views go through SigLIP in one cached graph.
  • Fused q/k/v as one GEMM over all views.
  • Projections are oneDNN matmuls with fused bias/GELU/residual.
  • Fused LayerNorm nodes.
  • Attention reads straight from the fused qkv.
• FoldQuant SYCL
  • Work-group-per-row activation kernel up to K = 16384.
  • INT4 activations stored one code per byte.
  • Fused q/k/v and gate/up sites, with the GEGLU fed straight into the down activation.
  • adaRMS in the act nodes and gated residuals in the o/down epilogues (expert).
• Fused nodes
  • Prefix: q/k/v GEMM, RoPE and F16 packing as one node.
  • Expert: q/k/v GEMM, RoPE and cached-KV attention as one node.
  • Expert attention over contiguous keys as two oneDNN GEMMs and a softmax.
• INT4 × INT4 DPAS GEMMs (CuTe, src/sycl/vla_sycl_s4gemm.cpp) for the W4A4 sites of the prefix and the expert.
  • Exact int32 output for q/k/v, o and down.
  • The gate/up kernel applies the GEGLU in its epilogue.
  • Tiles chosen by K.
  • VLA_FQ_NO_S4GEMM=1 falls back to oneDNN.
• vla-server --warmup N (default 1) before binding.

Latency

Intel PC, 2 × 224 views, 49-token prompt, --flash-attn, 10 denoise steps, p50 of 10 runs.

Model	main	this branch	speedup
SO101 W4A4	332.1 ms	101.8 ms	3.3x
SO101 W4A4 + o/down INT8	335.9 ms	109.0 ms	3.1x
SO101 W8A8	352.7 ms	131.3 ms	2.7x
SO101 bf16	372.2 ms	237.1 ms	1.6x
LIBERO W4A4	284.8 ms	98.7 ms	2.9x
LIBERO W4A4 + o/down INT8	286.9 ms	106.1 ms	2.7x
LIBERO W8A8	305.7 ms	128.0 ms	2.4x
LIBERO bf16	326.1 ms	212.4 ms	1.5x

Vision drops from 77.5 to 24 ms.

Accuracy

• VLA_FQ_CHECK=1: every FoldQuant node is byte-exact against the CPU reference.
  • SO101 W4A4 / W4A4 + o/down INT8 / W8A8, bf16, and LIBERO.
• Replay on 330 recorded banana frames with fixed noise, measured as action cosine vs bf16 on the same runtime:
  • this branch, this branch with every opt-out set, and main 65d16f5 give the same distribution for each model (W8A8 median 0.9999);
  • FoldQuant's PyTorch fake-quant reproduces the W4A4 drift (median 0.979 here, 0.981 in PyTorch), so that drift comes from the SO101 checkpoint's LLM W4A4 quantization, not from this runtime.

Test plan

• ctest on the Intel PC (SYCL), 11/11
• ctest on Jetson AGX Orin (CUDA), 14/14
• VLA_FQ_CHECK=1 per-node SYCL vs CPU check, 0 mismatches
• vla-bench latency table above
• Replay MAE and action cosine on recorded SO101 episodes