Skip to content

Fuse QKV prologue into Qwen3-14B paged attention #789

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Little-oil wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Fuse QKV prologue into Qwen3-14B paged attention #789

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
16c8d90
feat(qwen3): Stage 1 isolated rope_qkv CCEC kernel (compute byte-exac…
Jul 17, 2026
59d2d55
feat(qwen3): Stage 2 — fuse rope_qkv prologue into paged_attention_cce
Jul 17, 2026
f2c48ae
fix(qwen3): args-based indexing for fused rope_qkv cross-core sync
Jul 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
142 changes: 18 additions & 124 deletions models/qwen3/14b/decode_fwd.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -468,127 +468,8 @@ def _decode_layer( # noqa: PLR0913 — model signature is intrinsic
)
v_proj = pl.assemble(v_proj, v_acc, [0, n0], atomic=pl.AtomicType.Add)

# QK-norm and RoPE retain the main implementation's fused arithmetic,
# but run as a standalone producer for the external CANN attention.
with pl.spmd(
ROPE_CORES,
name_hint="rope_qkv",
deps=[q_proj_tid, k_proj_tid, v_proj_tid, rms_tid],
) as rope_tid:
rope_core = pl.get_block_idx()
q_red_pad = pl.full(
[1, (Q_HEAD_PAD - Q_PER_KV) * HEAD_DIM],
dtype=pl.FP32,
value=0.0,
)
k_red_pad = pl.full(
[1, (K_RED_ROWS - 1) * HEAD_DIM],
dtype=pl.FP32,
value=0.0,
)
for it in pl.pipeline(ROPE_ITEMS_PER_CORE, stage=2):
g_idx = rope_core + it * ROPE_CORES
if g_idx < NUM_KV_HEADS * user_batch:
ki = g_idx // user_batch
b = g_idx - ki * user_batch
ctx_len = pl.read(seq_lens, [b])
inv_rms_b = pl.read(inv_rms_states, [b, 0])
pos = ctx_len - 1
wr_slot = pl.cast(pl.tensor.read(slot_mapping, [b]), pl.INDEX)
wr_slot_block = wr_slot // BLOCK_SIZE
wr_slot_offset = wr_slot - wr_slot_block * BLOCK_SIZE
cos_lo = rope_cos[pos : pos + 1, 0:HALF_DIM]
cos_hi = rope_cos[pos : pos + 1, HALF_DIM:HEAD_DIM]
sin_lo = rope_sin[pos : pos + 1, 0:HALF_DIM]
sin_hi = rope_sin[pos : pos + 1, HALF_DIM:HEAD_DIM]

kv_col = ki * HEAD_DIM
k_raw = pl.mul(
pl.reshape(
pl.concat(
k_proj[b : b + 1, kv_col : kv_col + HEAD_DIM],
k_red_pad,
),
[K_RED_ROWS, HEAD_DIM],
),
inv_rms_b,
)
k_ss = pl.row_sum(pl.mul(k_raw, k_raw))
k_inv = pl.recip(pl.sqrt(pl.add(pl.mul(k_ss, HEAD_DIM_INV), EPS)))
k_normed = pl.row_expand_mul(
pl.col_expand_mul(k_raw, k_norm_w),
k_inv,
)
k_full = k_normed[0:1, :]
k_lo = k_full[:, 0:HALF_DIM]
k_hi = k_full[:, HALF_DIM:HEAD_DIM]
rot_lo = pl.sub(
pl.col_expand_mul(k_lo, cos_lo),
pl.col_expand_mul(k_hi, sin_lo),
)
rot_hi = pl.add(
pl.col_expand_mul(k_hi, cos_hi),
pl.col_expand_mul(k_lo, sin_hi),
)
cache_row = (
layer_cache_base
+ (wr_slot_block * BLOCK_SIZE + wr_slot_offset) * NUM_KV_HEADS
+ ki
)
k_cache = pl.assemble(
k_cache,
pl.cast(pl.concat(rot_lo, rot_hi), target_type=pl.BF16),
[cache_row, 0],
)
v_row_bf16 = pl.cast(
pl.mul(
v_proj[b : b + 1, ki * HEAD_DIM : (ki + 1) * HEAD_DIM],
inv_rms_b,
),
target_type=pl.BF16,
)
v_cache = pl.assemble(v_cache, v_row_bf16, [cache_row, 0])

q_base = ki * Q_PER_KV
q_raw = pl.mul(
pl.reshape(
pl.concat(
q_proj[
b : b + 1,
q_base * HEAD_DIM : (q_base + Q_PER_KV) * HEAD_DIM,
],
q_red_pad,
),
[Q_HEAD_PAD, HEAD_DIM],
),
inv_rms_b,
)
q_ss = pl.row_sum(pl.mul(q_raw, q_raw))
q_inv = pl.recip(pl.sqrt(pl.add(pl.mul(q_ss, HEAD_DIM_INV), EPS)))
q_heads = pl.row_expand_mul(
pl.col_expand_mul(q_raw, q_norm_w),
q_inv,
)
q_lo = q_heads[:, 0:HALF_DIM]
q_hi = q_heads[:, HALF_DIM:HEAD_DIM]
q_rot_lo = pl.sub(
pl.col_expand_mul(q_lo, cos_lo),
pl.col_expand_mul(q_hi, sin_lo),
)
q_rot_hi = pl.add(
pl.col_expand_mul(q_hi, cos_hi),
pl.col_expand_mul(q_lo, sin_hi),
)
q_row = b * NUM_HEADS + q_base
q_tnd_flat = pl.assemble(
q_tnd_flat,
pl.cast(
pl.concat(q_rot_lo, q_rot_hi)[0:Q_PER_KV, :],
target_type=pl.BF16,
),
[q_row, 0],
)

# QK-norm + RoPE + KV-cache writes are fused into the mixed-core attention extern
# (paged_attention_cce); q_tnd / k_cache / v_cache are written by its AIV prologue (InOut).
q_tnd = pl.reshape(q_tnd_flat, [BATCH, NUM_HEADS, HEAD_DIM])
attn_out_tnd = pl.reshape(attn_out, [BATCH, NUM_HEADS, HEAD_DIM])
attention_core_num = PA_DEFAULT_BLOCK_DIM
Expand All @@ -598,18 +479,31 @@ def _decode_layer( # noqa: PLR0913 — model signature is intrinsic
allow_early_resolve=True,
sync_start=True,
deps=[
rope_tid,
q_proj_tid,
k_proj_tid,
v_proj_tid,
rms_tid,
pa_tiling_tid,
attn_out_seed_tid,
mlp_out_seed_tid,
],
) as attn_done_tid:
attn_out_tnd = paged_attention_cce(
q_proj,
k_proj,
v_proj,
inv_rms_states,
q_norm_w,
k_norm_w,
seq_lens,
slot_mapping,
rope_cos,
rope_sin,
block_table,
attn_out_tnd,
q_tnd,
k_cache,
v_cache,
block_table,
attn_out_tnd,
pa_workspace,
pa_metadata,
layer_cache_base,
Expand Down
67 changes: 62 additions & 5 deletions models/qwen3/14b/kernels/paged_attention_cce/attention/entry.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,17 +25,74 @@ extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { (void)args; }

#else

#include <pto/pto-inst.hpp>

#include "../kernel/fai_body.hpp"

namespace {

static __aicore__ __attribute__((always_inline)) __gm__ int32_t *qwen_fai_barrier_data(GM_ADDR metadata) {
uint64_t raw_barrier = reinterpret_cast<uint64_t>(metadata + qwen_fai_metadata::kBarrierAlignmentOffset);
uint64_t aligned_barrier = (raw_barrier + qwen_fai_metadata::kBarrierAlignmentBytes - 1) &
~(static_cast<uint64_t>(qwen_fai_metadata::kBarrierAlignmentBytes) - 1);
return reinterpret_cast<__gm__ int32_t *>(aligned_barrier);
}

// After the AIV prologue writes q/k/v to GM, synchronize so the AIC attention does
// not read them early. The 48 AIV lanes soft-barrier among themselves over the
// rope-ready slots; each AIC polls those slots until all 48 are non-zero.
// Indexing is args-based (get_block_idx*2 + get_sub_block_id): the hardware
// get_subblockdim() reads 0 in this pypto launch (the dispatch does not set the
// subblockdim register), so pto::SYNCALL's builtin participant model
// (SYNCALL_GET_MIX_PARTICIPANT_*) cannot be used — it under-counts (24 not 72) and
// collides the AIV indices, deadlocking. The vendor attention uses the same
// args-based getters, so this matches its view of the topology. The soft barrier's
// DCCI + dsb(DSB_DDR) gives the cross-core GM visibility the AIC needs.
static __aicore__ __attribute__((always_inline)) void sync_qwen_rope_producers(
__gm__ int64_t *args,
__gm__ int32_t *fai_barrier
) {
__gm__ int32_t *ready = reinterpret_cast<__gm__ int32_t *>(
reinterpret_cast<__gm__ uint8_t *>(fai_barrier) + qwen_fai_metadata::kBarrierBytes
);
pipe_barrier(PIPE_ALL);
#if defined(__DAV_CUBE__)
while (true) {
uint32_t ready_count = 0;
for (uint32_t i = 0; i < qwen_fai_metadata::kRopeReadySlotCount; ++i) {
__gm__ int32_t *slot = ready + i * qwen_fai_metadata::kRopeReadySlotWords;
dcci(reinterpret_cast<__gm__ void *>(slot), SINGLE_CACHE_LINE);
dsb(DSB_DDR);
if (slot[0] != 0) {
++ready_count;
}
}
pipe_barrier(PIPE_ALL);
if (ready_count >= qwen_fai_metadata::kRopeReadySlotCount) {
break;
}
}
#elif defined(__DAV_VEC__)
uint32_t aiv_idx = static_cast<uint32_t>(get_block_idx(args)) * 2 +
static_cast<uint32_t>(get_sub_block_id(args));
pto::SYNCALL_SOFT_AIV_BARRIER(ready, reinterpret_cast<__ubuf__ int32_t *>(0x3000),
static_cast<int32_t>(qwen_fai_metadata::kRopeReadySlotCount),
static_cast<int32_t>(aiv_idx));
#endif
pipe_barrier(PIPE_ALL);
}

} // namespace

extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) {
GM_ADDR metadata = tensor_data<uint8_t>(args, 6);
GM_ADDR metadata = tensor_data<uint8_t>(args, kMetadataArg);
acquire_qwen_fai_metadata(metadata);
run_qwen_rope_qkv(args);
__gm__ int32_t *fai_barrier = qwen_fai_barrier_data(metadata);
sync_qwen_rope_producers(args, fai_barrier);
__gm__ const FAInferTilingData *tiling = reinterpret_cast<__gm__ const FAInferTilingData *>(metadata);
if (tiling->needCoreNum != 0) {
uint64_t raw_barrier = reinterpret_cast<uint64_t>(metadata + qwen_fai_metadata::kBarrierAlignmentOffset);
uint64_t aligned_barrier = (raw_barrier + qwen_fai_metadata::kBarrierAlignmentBytes - 1) &
~(static_cast<uint64_t>(qwen_fai_metadata::kBarrierAlignmentBytes) - 1);
run_qwen_fai<true>(args, reinterpret_cast<__gm__ int32_t *>(aligned_barrier));
run_qwen_fai<true>(args, fai_barrier);
} else {
run_qwen_fai<false>(args);
}
Expand Down
50 changes: 50 additions & 0 deletions models/qwen3/14b/kernels/paged_attention_cce/attention_only/entry.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
/*
* Copyright (c) PyPTO Contributors.
* This program is free software; you can redistribute it and/or modify it under the terms and conditions of
* CANN Open Software License Agreement Version 2.0 (the "License").
* Please refer to the License for details. You may not use this file except in compliance with the License.
* THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
* INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
* See LICENSE in the root of the software repository for the full text of the License.
* -----------------------------------------------------------------------------------------------------------
*/

// Standalone attention entry (no rope prologue, no cross-core sync). Defines
// QWEN_FAI_ATTENTION_ONLY_ABI so fai_body.hpp compiles run_qwen_fai with the
// attention-only arg layout (0..7) and skips the rope_qkv producer. Used by the
// paged_attention_only_cce extern + test_paged_attention_cce.py.

#include <cstdint>

#include "tensor.h"

#ifdef __CPU_SIM
#ifndef __gm__
#define __gm__
#endif
#ifndef __aicore__
#define __aicore__ [aicore]
#endif

extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { (void)args; }

#else

#define QWEN_FAI_ATTENTION_ONLY_ABI
#include "../kernel/fai_body.hpp"

extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) {
GM_ADDR metadata = tensor_data<uint8_t>(args, kMetadataArg);
acquire_qwen_fai_metadata(metadata);
__gm__ const FAInferTilingData *tiling = reinterpret_cast<__gm__ const FAInferTilingData *>(metadata);
if (tiling->needCoreNum != 0) {
uint64_t raw_barrier = reinterpret_cast<uint64_t>(metadata + qwen_fai_metadata::kBarrierAlignmentOffset);
uint64_t aligned_barrier = (raw_barrier + qwen_fai_metadata::kBarrierAlignmentBytes - 1) &
~(static_cast<uint64_t>(qwen_fai_metadata::kBarrierAlignmentBytes) - 1);
run_qwen_fai<true>(args, reinterpret_cast<__gm__ int32_t *>(aligned_barrier));
} else {
run_qwen_fai<false>(args);
}
}

#endif
Loading
Loading