feat(kunlun): add paged_caching, paged_attention and paged_attention_…prefill - #1546
Open
rubik-hua wants to merge 1 commit into
Open
feat(kunlun): add paged_caching, paged_attention and paged_attention_…prefill#1546rubik-hua wants to merge 1 commit into
rubik-hua wants to merge 1 commit into
Conversation
…prefill
Enable --enable-paged-attn end to end on Kunlun XPU. Validated on P800:
all three infiniop suites pass and Qwen3-0.6B generates with the
paged-attn backend.
- paged_caching: flatten (token, head, K/V) row copies across all
clusters/cores; each row moves via GM2LM/LM2GM DMAs with no cross-core
sync; negative slot_mapping entries (padding) are skipped.
- paged_attention (decode): one (seq, kv_head) task per cluster; GQA heads
share KV tiles staged in shared memory so paged KV is read once per kv
head; online softmax keeps per-head (m, l, acc) state across tiles with
fp32 accumulation; supports ALiBi and I32/I64/U32 indices.
- paged_attention_prefill (varlen, causal): one (seq, q-tile, kv_head,
head-chunk) task per cluster over packed query tokens; KV tiles beyond a
row's causal limit (cache_len + u) are skipped; each (head, query-token)
row's online-softmax state lives in its owning core's local memory, so
state updates need no shared traffic or atomics.
All three register INFINI_DEVICE_KUNLUN in the dispatch switches under
ENABLE_KUNLUN_API and validate index-dtype consistency, q/cache
contiguity, and 64-bit task/row counters. Launch defaults match the P800
topology (12 clusters x 64 cores; launching more cores than physically
present silently drops work) and are overridable via
INFINIOP_KUNLUN_{PA,PC,PAP}_{CLUSTERS,CORES}.
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.
Enable --enable-paged-attn end to end on Kunlun XPU. Validated on P800:
all three infiniop suites pass and Qwen3-0.6B generates with the
paged-attn backend.
paged_caching: flatten (token, head, K/V) row copies across all clusters/cores; each row moves via GM2LM/LM2GM DMAs with no cross-core sync; negative slot_mapping entries (padding) are skipped.
paged_attention (decode): one (seq, kv_head) task per cluster; GQA heads share KV tiles staged in shared memory so paged KV is read once per kv head; online softmax keeps per-head (m, l, acc) state across tiles with fp32 accumulation; supports ALiBi and I32/I64/U32 indices.
paged_attention_prefill (varlen, causal): one (seq, q-tile, kv_head, head-chunk) task per cluster over packed query tokens; KV tiles beyond a row's causal limit (cache_len + u) are skipped; each (head, query-token) row's online-softmax state lives in its owning core's local memory, so state updates need no shared traffic or atomics.
All three register INFINI_DEVICE_KUNLUN in the dispatch switches under
ENABLE_KUNLUN_API and validate index-dtype consistency, q/cache
contiguity, and 64-bit task/row counters. Launch defaults match the P800
topology (12 clusters x 64 cores; launching more cores than physically
present silently drops work) and are overridable via
INFINIOP_KUNLUN_{PA,PC,PAP}_{CLUSTERS,CORES}.
验证测试
python examples/test_infer.py --device kunlun --model=/mnt/geogpt-doc-new/default/infinilm-models/DeepSeek-R1-Distill-Qwen-7B --enable-paged-attn --prompt "who are you"

python examples/test_infer.py --device kunlun --model=/mnt/geogpt-doc-new/default/infinilm-models/Qwen3-0.6B --enable-paged-attn --prompt "who are you"
