Skip to content

feat(kunlun): add paged_caching, paged_attention and paged_attention_…prefill - #1546

Open
rubik-hua wants to merge 1 commit into
InfiniTensor:mainfrom
rubik-hua:paged-attn
Open

feat(kunlun): add paged_caching, paged_attention and paged_attention_…prefill#1546
rubik-hua wants to merge 1 commit into
InfiniTensor:mainfrom
rubik-hua:paged-attn

Conversation

@rubik-hua

Copy link
Copy Markdown

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"
image

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"
image

…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}.
@rubik-hua
rubik-hua requested a review from a team September 3, 2026 13:12
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