Pull requests / #2555

#2555 sycl: accelerate IQ4_XS MMVQ on Battlemage

open · @xiangyuT · 0 commentaires · Sur GitHub

Multi-GPUAMD / HIPNVIDIA / CUDAModels & quantsWindowsLinux

Description

## Summary

Reduce IQ4_XS × Q8_1 decode work on Intel Battlemage by sharing loads across adjacent quantization groups and token columns, and replacing per-value codebook lookups with a bit-select tree.

## What changed

- `sycl/src/kernels/cuda/native_mmvq.dp.cpp`: add ESIMD INT8 XMX kernels that reuse IQ4 headers and packed weights across two Q8_1 blocks and up to four token columns. Selected shapes also load the adjacent Q8_1 blocks and their scales together.
- `sycl/include/strata/kernels/iq4_lut.hpp`: add an exact packed IQ4 bit-select lookup, used by the remaining eligible wide-kernel shapes.
- Dispatch uses upstream `intel_gpu_gen` / `IntelGpuGen::Battlemage`. The optimized defaults cover K=2560/6144 and T=1–6, subject to alignment, row-count and shape guards; pair reuse and Q8 splitting cover T≤4. Other devices and unsupported shapes retain the original dispatch. CUDA/HIP implementations and public documentation are unchanged.

## Dispatch controls

| Control | Unset (default) or `1` | `0` |
|---|---|---|
| `STRATA_SYCL_IQ4_PAIR_REUSE` | Enable pair reuse and allow its Q8-split subpath. | Bypass pair reuse and Q8 splitting; the independent wide-kernel lookup optimization may still run. |
| **Child of pair reuse:** `STRATA_SYCL_IQ4_Q8_SPLIT` | Load adjacent Q8_1 blocks and scales together on selected shapes within the pair-reuse branch. | Use the pair-reuse kernel without Q8 splitting. |
| `STRATA_SYCL_IQ4_LOOKUP` | Use the bit-select lookup in eligible wide kernels when the pair-reuse branch is not selected. | Use the original wide-kernel lookup; pair reuse and Q8 splitting are unaffected. |

There is no single master switch for this PR. To select the original dispatch, set both `STRATA_SYCL_IQ4_PAIR_REUSE=0` and `STRATA_SYCL_IQ4_LOOKUP=0`. The Q8-split control has no effect when pair reuse is disabled.

All controls retain the device, alignment and shape guards described above. Set them before starting the engine and restart it when changing them: their values are cached on first use. They support diagnosis and fallback; normal operation needs no environment-variable setting.

## Validation

One Arc Pro B70 32 GB, Linux/xe, Xeon w5-3435X; `intel/omix:0.4.0-devel-ubuntu24.04`, oneAPI DPC++ 2026.1.1, compute runtime 26.31.39395.13, Level Zero loader 1.32.0. Container limit: 96 GiB RAM / 16 CPU quota. Builds and GPU tests ran inside this container.

Build:

```sh
cmake -S sycl -B build -G Ninja -DCMAKE_BUILD_TYPE=Release \
  -DCMAKE_C_COMPILER=icx -DCMAKE_CXX_COMPILER=icpx \
  -DSTRATA_SYCL_AOT=bmg-g31 -DSTRATA_SYCL_PARITY=ON
cmake --build build --target strata mmvq_bench
```

At submitted revision `e5c0c9c69d38c9469c185ba67c0f7d61e0ecc564`:

- IQ4 output comparisons and guard checks pass for default and fallback dispatch. Profiling confirms the expected kernel paths.
- Fresh serving runs with 4096, 8192 and 32768 input tokens each produce 128 output token IDs identical to upstream, with default dispatch and with the optimizations disabled (6 comparisons). These profiled runs are excluded from latency measurements.
- Offline device-image compilation of the affected unit passes for BMG G21, DG2 and PVC; this does not establish runtime correctness or performance on those devices.

Additional checks at measured revision `9683cda6c60139976058d0f2abe90c10c1aff426`: ten locally defined requests produce identical upstream token IDs at each max-context setting, 32768 and 131072 (20/20). They cover arithmetic, Python, prose, Chinese, JSON, Unicode and 4096/8192-token synthetic code-review inputs. A frozen 4096-token CLI input with `--pcie-frac 0` also produces all 248320 first-token logits byte-identically to upstream, with all values finite.

## End-to-end performance

Timing comparison: candidate `9683cda6c601` versus upstream `296b59d2f73c`, on the B70 environment described above. The current head (`e5c0c9c69d38`) has been checked for correctness but has not been remeasured for endpoint latency.

TTFT is the time from the request to the first nonempty SSE token. TPOT is the time between the first and last nonempty SSE tokens divided by generated tokens minus one. Negative changes mean lower latency.

| Input / output tokens | TTFT baseline → default (s) | Change | TPOT baseline → default (ms/token) | Change |
|---|---:|---:|---:|---:|
| 4096 / 128 | 3.474050 → 3.474139 | +0.003% | 14.322062 → 13.491105 | -5.802% |
| 8192 / 256 | 7.027478 → 7.030339 | +0.041% | 13.606768 → 12.818306 | -5.795% |
| 32768 / 256 | 29.693462 → 29.710838 | +0.059% | 14.078578 → 13.289964 | -5.606% |

<details>
<summary>Test configuration and sampling</summary>

| Item | Setting |
|---|---|
| Timing baseline / frontend | `296b59d2f73c8f04a805df29c2fa97655adace78` |
| Timed candidate | `9683cda6c60139976058d0f2abe90c10c1aff426`; this PR only |
| Model | Qwen3.8-Flash-Next-GSQ-RCO, IQ2_XS; revision `ed59f92082b1e93c0e96d60a8b11aab089b52f09` |
| Sampling | Greedy, seed 42, thinking off |
| KV and context | INT8 KV; max context 131072; resident KV 32768 |
| Expert cache | `--expert-cache 17772` experts; reserve 1536 MiB; mirror budget 16384 MiB |
| Prefill and prompt cache | Auto prefill; no prefill borrowing; prompt cache off |
| Speculation and adaptation | spec4 / min-p0.5 / suffix3; adapt-every8 / decay1 |

Observed capacity in every arm: 18609 resident expert slots and 5967 mirrored experts (8.02 GiB). The 16384 MiB setting is a budget, not the actual mirror allocation.

Existing upstream settings held equal in both arms: `STRATA_VERIFY_DEVICE_PLAN=1`, `STRATA_VERIFY_NO_HOST=1`, `STRATA_STAGER_THREADS=12`, `STRATA_MIRROR_MIB=16384`, `STRATA_DECODE_TIMING=1`; `ONEAPI_DEVICE_SELECTOR=level_zero:gpu`, `SYCL_CACHE_PERSISTENT=0`.

Each arm starts a fresh process. One short request and three full-input requests are excluded as warmups; three requests are then measured.

| Input / output tokens | Arm order | Measured requests: baseline / candidate |
|---|---|---:|
| 4096 / 128 | Five baseline/candidate pairs, alternating order | 15 / 15 |
| 8192 / 256 | Baseline / candidate / baseline | 6 / 3 |
| 32768 / 256 | Baseline / candidate / baseline | 6 / 3 |

The 4K percentage averages the five pair-mean ratios. The 8K/32K percentage averages the three request-wise ratios against the matching before/after baseline means.

</details>

<details>
<summary>Baseline variability and actual speculative work</summary>

**4K**

Pair changes, TTFT: -0.075%, +0.000%, +0.054%, +0.042%, -0.009%; TPOT: -5.762%, -5.829%, -5.812%, -5.819%, -5.787%.
Baseline first-to-last change: TTFT -0.003%, TPOT +0.009%; block-mean spans 0.081% / 0.052%.
Output token IDs match across all arms, including warmups. Mean windows/offered/accepted per measured request: baseline 54.000/111.000/74.000 (15 requests); candidate 54.000/111.000/74.000 (15 requests).

**8K**

Baseline drift: TTFT -0.053%, TPOT -0.017%.
Output token IDs match across all arms, including warmups. Mean windows/offered/accepted per measured request: baseline 102.667/209.000/153.667 (6 requests); candidate 102.667/209.000/153.667 (3 requests).

**32K**

Baseline drift: TTFT +0.009%, TPOT -0.002%.
Output token IDs match across all arms, including warmups. Mean windows/offered/accepted per measured request: baseline 103.333/202.333/154.000 (6 requests); candidate 103.333/201.667/154.000 (3 requests).

Speculation settings are fixed; offered counts include suffix drafts and actual work may differ because suffix selection is time-adaptive. The measurements are warmed and post-readiness; the host page cache was not flushed, so no disk-cold TTFT result is claimed.

</details>

## Kernel measurements

Independent microbenchmarks use synthetic weights and activations at matrix dimensions used by the model. K and N are the input and output dimensions; T is the tested number of token columns. Timings were collected at `9683cda6c60139976058d0f2abe90c10c1aff426`.

The baseline selects the original dispatch in the same binary with the IQ4 optimizations disabled; the candidate uses default dispatch. Each baseline/candidate/baseline arm has three measured samples after a ≥2-second/1000-call warmup; samples last ≥3 seconds/2000 calls. Times include host submission and the final queue wait. These are isolated kernel-call timings; model-level results are reported in the E2E table above.

| Operation | K / N / T | Baseline (µs) | Default (µs) | Change | Drift |
|---|---|---:|---:|---:|---:|
| GDN QKV projection | 2560/10240/1 | 39.476 | 24.147 | -38.831% | +0.006% |
| GDN QKV projection | 2560/10240/2 | 43.162 | 25.796 | -40.235% | +0.010% |
| GDN QKV projection | 2560/10240/3 | 47.355 | 26.392 | -44.268% | -0.005% |
| GDN QKV projection | 2560/10240/4 | 52.056 | 26.666 | -48.774% | +0.007% |
| GDN QKV projection | 2560/10240/5 | 57.041 | 49.766 | -12.754% | -0.005% |
| GDN QKV projection | 2560/10240/6 | 65.613 | 60.329 | -8.053% | +0.006% |
| GDN output projection | 6144/2560/1 | 26.958 | 19.187 | -28.826% | +0.002% |
| GDN output projection | 6144/2560/3 | 30.833 | 19.917 | -35.404% | -0.007% |
| GDN output projection | 6144/2560/4 | 34.426 | 20.518 | -40.401% | -0.008% |
| GDN output projection | 6144/2560/6 | 50.411 | 45.301 | -10.136% | -0.020% |
| MTP draft head | 2560/106299/1 | 420.032 | 272.352 | -35.159% | +0.164% |
| Shared expert gate/up | 2560/640/3 | 5.502 | 3.662 | -33.438% | -0.043% |

## Limitations

Runtime and performance validation is limited to B70. Other GPUs, Windows and multi-GPU operation were not runtime-tested. The patch changes only the two SYCL implementation files listed above; no test-source changes are included.

Sur le site

Liens install, modèles, releases.