feat(ark): add INT4 S4 pre-packed Q*K kernel for SageAttention

[luoyu-intel]

feat(ark): add INT4 S4 pre-packed Q*K kernel for SageAttention

Summary

Replace SageAttention V1's INT8 Q*K path with standard signed INT4 DPAS on Intel Arc Pro B60 (Battlemage G21). This PR introduces a pre-packed INT4 Q/K path that achieves ~11.8% throughput improvement over the existing INT8 path on the benchmark shape, with exact numerical correctness (max_diff = 0).

Motivation

INT4 Q*K via XE_DPAS_TT on Intel GPU offers native 4-bit dot-product accumulation, halving the data bandwidth of the Q*K matmul compared to INT8 while maintaining signed precision. This is the first step toward a full INT4 attention pipeline; P*V remains FP16/BF16 in this PR.

What's Changed

File	Change
sycl_tla_sdpa.hpp	Typed make_gmem_ptr for subbyte iterators; MMA ops templated on ElementQ/ElementK; INT4 conditional tile shapes (K tile = 64); locked M=256 / Stages=2
CMakeLists.txt	Bake two INT4 sycl-tla header patches in-place (integer_subbyte.h drop default template arg; copy_traits_xe_2d.hpp raw_pointer_cast subbyte fix) so fresh builds compile without manual edits
__init__.py	New sage_s4 public API
ark.cpp	sage_s4 pybind entry point
sdpa.cpp	S4 launchers (launch_prefill_kernel_{f16,bf16}_{64,128}_sage_s4) + sdpa_impl_qks4_pvhalf dispatch
xe_sage_fwd_kernel.hpp	Kernel adjustments for INT4 Q*K
sycl_s8_wrapper.hpp	INT4 wrapper plumbing
sycl_tla_common.hpp	Shared utilities
sycl_tla_s8_gemm.hpp	GEMM tile adjustments
test/bench_sage_int4.py	New — pre-packed S4 vs S8 benchmark
test/test_sage_int4.py	New — INT4 correctness regression (6 cases)
.gitignore	Ignore out-of-source xbuild* build dirs

Technical Details

• INT4 packing: 2 signed 4-bit values per byte; CUTE subbyte iterators preserve logical INT4 indexing.
• Q*K MMA: XE_DPAS_TT<cute::gcd(SGTileQ, 8), int32_t, ElementQ, ElementK> (signed 4-bit dot product into int32 accumulator).
• P*V MMA: XE_DPAS_TT<cute::gcd(SGTileQ, 8), float, ElementO> — unchanged, FP16/BF16.
• Tile shapes (INT4, M=256): QK = Shape<_256, _64, _64> (K tile 64 vs INT8's 32); PV = Shape<_256, _32, _64>; Out = Shape<_256, _128>.
• Pipeline: 2 stages (both QK and PV pipelines).
• sycl-tla patches (applied via CMake during build, guarded by string-match check):
  1. cutlass/integer_subbyte.h: remove = true default from template <int Bits, bool Signed = true> → template <int Bits, bool Signed>.
  2. cute/atom/copy_traits_xe_2d.hpp: base_ptr((uint64_t) &*src.data()) → base_ptr(reinterpret_cast<uint64_t>(raw_pointer_cast(&*src.data()))) — required because subbyte_iterator has no implicit conversion to uint64_t.

Performance

Config	Shape B=1, H=40, S=10240, D=128
S4 (INT4 Q*K, FP16 P*V)	18.89 ms
S8 (INT8 Q*K, FP16 P*V)	21.41 ms
Speedup	~11.8% faster
max_diff vs S8	0 (exact)

Testing

• test_sage_int4.py: 6 passed (accuracy across shapes/dtypes)
• test_sdpa_parity.py: 9 passed (SAGE/SDPA regression — no existing kernel broken)
• bench_sage_int4.py: S4 vs S8 benchmark, prints s4_ms, s8_ms, speedup ratio, max_diff, mean_diff

Scope / Deferred

This PR covers pre-packed INT4 Q/K only. The following are intentionally deferred:

• INT4 auto-quantizer (runtime quantization of FP16/BF16 Q/K)
• INT4 P*V (4-bit probability×value matmul)
• Decode (small-batch) kernel
• NHD / arbitrary-stride layouts
• High-level automatic quantization API

Build & Run

# OneAPI 2026.0 + conda env (Python 3.13, torch 2.13+xpu)
source /path/to/oneapi/setvars.sh --force
conda activate t213

# Build (self-contained; clones sycl-tla, applies patches automatically)
cmake -S auto_round_extension/ark/auto_round_kernel -B xbuild-int4 \
  -DARK_XPU=ON -DARK_SYCL_TLA=ON -DCMAKE_CXX_COMPILER=icx
cmake --build xbuild-int4 --parallel 10

# Copy .so into package
cp xbuild-int4/auto_round_kernel_xpu.cpython-313-x86_64-linux-gnu.so \
   auto_round_extension/ark/auto_round_kernel/

# Run tests
PYTHONPATH=. python -m pytest auto_round_extension/ark/test/test_sage_int4.py -v
PYTHONPATH=. python -m pytest auto_round_extension/ark/test/test_sdpa_parity.py -v

# Benchmark
PYTHONPATH=. python auto_round_extension/ark/test/bench_sage_int4.py \
  --batch 1 --heads 40 --seq 10240 --head-dim 128 --block-size 64 \
  --warmup 20 --iterations 100