vLLM: integrate FP16 packed QSA extension - #651
Open
xiangyuT wants to merge 2 commits into
Open
Conversation
Signed-off-by: Xiangyu Tian <xiangyu.tian@intel.com>
Signed-off-by: Xiangyu Tian <xiangyu.tian@intel.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
custom_esimd_kernels_vllm.qsa_opsas a standalone SYCL extension for the fixed Qwen3.8 TP8-rank sparse paged-attention contractR=1coverageqsa_select_paged_tokens(...)XPU kernel that reads paged compressed keys, computes four-headReLU(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 synchronizationqsa_opsextension from the regularsetup.py/setup_sycl.pypaths and the focusedsetup_qsa_only.pyiteration pathThe initial imported attention source SHA256 is
f1e3ad106cff2c01232e9ac90565f1decbe4804af97fbbf8967de065e145d1ad. The current PR head is87a31a5e340d3b73c89a7ec80a850536f6659b61.Exact contracts
Attention:
[R,3,256][262144,512,256,1]storage_offset=256[R,2051]int32, including holes and duplicatesSelection:
[R,4,128][P,128,1,128]token_topk=2048,compress_ratio=4, caller-owned[R,2051]int32 output padded with-1Validation
python -m pytest -q tests/test_qsa_build_config.py: 3 passedTORCH_XPU_ARCH_LIST=bmg-g31 MAX_JOBS=1 python setup_qsa_only.py build_ext --inplace: passed; built DSO SHA2566a4c25705f4ebc19dc22c6e5b89f8016b89f2d099cecc547639feea95ed71cc3R=[1,2]across empty, holes/duplicates/cross-page, valid width 32, and full width 2051, plus fail-closed dtype/layout checksgit diff --check: passedR=1 attention latency uses five excluded warm-up units, 10 samples, and 128 calls per measurement unit:
R=1 selection B/C/B uses five excluded warm-up units, 10 pairs, and 64 calls per measurement unit:
A one-call targeted profiler capture for the P8192
R=1case 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.patchdoes 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
qsa_build.pyhelper and inlining the QSASyclExtensiondefinition into the existing setup scripts to match the current repository styleqsa_opsextension standalone unless a later change deliberately splits the kernel from its pybind module and integrates the binding into the core extension