Skip to content

vLLM: integrate FP16 packed QSA extension - #651

Open
xiangyuT wants to merge 2 commits into
intel:upgrade/v0.26.0from
xiangyuT:dev/qwen38-qsa-fp16-handoff-20260828
Open

vLLM: integrate FP16 packed QSA extension#651
xiangyuT wants to merge 2 commits into
intel:upgrade/v0.26.0from
xiangyuT:dev/qwen38-qsa-fp16-handoff-20260828

Conversation

@xiangyuT

@xiangyuT xiangyuT commented Aug 28, 2026

Copy link
Copy Markdown
Contributor

Summary

  • add custom_esimd_kernels_vllm.qsa_ops as a standalone SYCL extension for the fixed Qwen3.8 TP8-rank sparse paged-attention contract
  • retain the validated FP16 exact-packed T4/SG32 attention kernel and add explicit R=1 coverage
  • add a fused FP16 qsa_select_paged_tokens(...) XPU kernel that reads paged compressed keys, computes four-head ReLU(Q*K) scores, performs streaming top-512 selection, expands blocks to 2,048 tokens, appends the three-token tail, and writes caller-owned [R,2051] int32 output without per-row .item()/CPU synchronization
  • build attention and selection into the same qsa_ops extension from the regular setup.py/setup_sycl.py paths and the focused setup_qsa_only.py iteration path
  • add focused build, correctness, latency, B/C/B, and targeted-profiler coverage

The initial imported attention source SHA256 is f1e3ad106cff2c01232e9ac90565f1decbe4804af97fbbf8967de065e145d1ad. The current PR head is 87a31a5e340d3b73c89a7ec80a850536f6659b61.

Exact contracts

Attention:

  • FP16 Q/output [R,3,256]
  • FP16 K/V exact-packed stride [262144,512,256,1]
  • V storage_offset=256
  • logical indices [R,2051] int32, including holes and duplicates
  • caller-owned output; inputs remain unchanged

Selection:

  • FP16 Q [R,4,128]
  • FP16 compressed-key cache [P,128,1,128]
  • int32 page table/token-to-request/sequence lengths; int64 query positions
  • fixed token_topk=2048, compress_ratio=4, caller-owned [R,2051] int32 output padded with -1
  • descending score order with ascending logical-index order for exact ties

Validation

  • python -m pytest -q tests/test_qsa_build_config.py: 3 passed
  • TORCH_XPU_ARCH_LIST=bmg-g31 MAX_JOBS=1 python setup_qsa_only.py build_ext --inplace: passed; built DSO SHA256 6a4c25705f4ebc19dc22c6e5b89f8016b89f2d099cecc547639feea95ed71cc3
  • XPU1-only attention + selection suite: 21 passed
    • attention: R=[1,2] across empty, holes/duplicates/cross-page, valid width 32, and full width 2051, plus fail-closed dtype/layout checks
    • selection: empty, short tail, full page, cross-page tail, duplicate physical-page mapping, exact-score ties, P8192 saturated width 2051, and multi-request, plus fail-closed dtype/output checks
  • git diff --check: passed

R=1 attention latency uses five excluded warm-up units, 10 samples, and 128 calls per measurement unit:

valid width mean latency
32 0.014193 ms
1024 0.355889 ms
2051 0.710274 ms

R=1 selection B/C/B uses five excluded warm-up units, 10 pairs, and 64 calls per measurement unit:

case Torch reference fused speedup wins baseline drift
cross-page tail, 129 visible blocks / 519 valid tokens 0.834475 ms 0.053155 ms 15.70x 10/10 -0.064%
P8192 saturated, 2303 visible blocks / 2051 valid tokens 0.862417 ms 0.319464 ms 2.70x 10/10 +0.271%

A one-call targeted profiler capture for the P8192 R=1 case observed 109 XPU events and 1,221 CPU events for the Torch reference, including two device-to-host copies. The fused path observed one XPU event, QsaSelectPagedTokensKernel, and 209 CPU profiler events. This is a rank-local targeted trace, not a vLLM E2E trace.

Scope and remaining gates

The current public vllm_for_multi_arc.patch does not yet contain the Qwen3.8 QSA caller, so this PR exposes the formal extension ABI for subsequent integration. It does not claim TP8 serving/E2E performance.

A full bdist wheel, BMG-g21 build, vLLM source-overlay E2E, TP8 GSM8K strict gate, public package validation, and focused image validation were intentionally not run on this host because the required E2E environment is unavailable. The package version remains 0.1.0.

Follow-up TODO

  • consider removing the single-purpose qsa_build.py helper and inlining the QSA SyclExtension definition into the existing setup scripts to match the current repository style
  • keep the compiled qsa_ops extension standalone unless a later change deliberately splits the kernel from its pybind module and integrates the binding into the core extension
  • preserve one authoritative set of validated FP16 exact-packed build macros when performing that refactor

Signed-off-by: Xiangyu Tian <xiangyu.tian@intel.com>
Signed-off-by: Xiangyu Tian <xiangyu.tian@intel.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant