From 16c8d90ccaf118d2f430fee8a214f4b221035842 Mon Sep 17 00:00:00 2001 From: Youhezhen Date: Thu, 16 Jul 2026 18:17:06 -0700 Subject: [PATCH 1/3] feat(qwen3): Stage 1 isolated rope_qkv CCEC kernel (compute byte-exact on a2a3) First half of fusing rope_qkv into paged_attention_cce: a standalone CCEC extern `paged_attention_rope_only_cce` that runs ONLY the rope_qkv producer (per-head QK-norm + NeoX RoPE + paged KV-cache writes), with a torch golden test (test_rope_qkv_cce.py) validating query + the written k_cache/v_cache rows byte-exact against decode_fwd.py's rope_qkv. decode_fwd.py and the attention extern are untouched (serving path unchanged); this only adds the isolated kernel + test. Stage 2 (fuse into attention + cross-core sync) is pending. The C++ rope compute (`run_qwen_rope_qkv` in fai_body.hpp, vec-only, AIC no-op) mirrors decode_fwd.py: x=proj*inv_rms; inv_head=rsqrt(sum(x^2)/128+eps); x=(x*weight)*inv_head; NeoX half-split rotate; cast BF16; write paged cache. Four C220 (a2a3) hazards were pinned by dumping intermediates to GM via the on-device golden test (iterated locally through task-submit) and fixed: 1. Dynamic-index GetValue on a UB LocalTensor misreads (~128x). The S-side sum loop `for i square.GetValue(i)` returned ss~7.3 vs the true tree-sum 0.057 (inv_head 4.19 vs 84.9). Constant-index UB GetValue (weight/cos GetValue(0)) and dynamic-index GM GetValue (inv_rms.GetValue(batch)) are both fine. Fixed by summing the 8 tree partials with literal constant-index GetValues. 2. C220 V-pipe does NOT serialize dependent ops on the same buffer (intermittent RAW; surfaces as a clean ratio when it fails). Added PipeBarrier() between every dependent V op: Muls(x,inv_rms)->Mul(square,x,x) (was reading unscaled x, ss inflated 1/inv_rms^2); the reduce Muls/Adds/Rsqrt chain; the rope lo/hi Mul->Sub/Add; and the V-path Muls(x,inv_rms)->Cast. 3. count<8 (sub-block) AscendC::Add HANGS on C220 (507018). The tree uses 64/32/16/8; the final 8->1 sum is done in S (constant-index GetValue), not via more V Adds. 4. MTE3->MTE2 DMA tear: MTE2 (GM->UB load) and MTE3 (UB->GM store) share the external-memory DMA path, so a load right after a store with only a qwen_mte3_to_v (MTE3->V) fence tears the in-flight store (intermittent whole-head GM garbage, proportional to write count; key_cache was spared because its write->read gap ran through V's Muls+Cast and drained naturally). Added qwen_mte3_to_mte2() (HardEvent::MTE3_MTE2) after every UB->GM store. Result: test_rope_qkv_cce.py PASSES query+key_cache+value_cache on a2a3. --- .../paged_attention_cce/kernel/fai_body.hpp | 261 ++++++++++++++++++ .../paged_attention_cce/rope_only/entry.cpp | 38 +++ models/qwen3/14b/paged_attention_cce.py | 79 ++++++ models/qwen3/14b/test_rope_qkv_cce.py | 163 +++++++++++ 4 files changed, 541 insertions(+) create mode 100644 models/qwen3/14b/kernels/paged_attention_cce/rope_only/entry.cpp create mode 100644 models/qwen3/14b/test_rope_qkv_cce.py diff --git a/models/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hpp b/models/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hpp index 4a0d97fa..f771271e 100644 --- a/models/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hpp +++ b/models/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hpp @@ -158,4 +158,265 @@ run_qwen_fai(__gm__ int64_t *args, __gm__ int32_t *barrier_state = nullptr) { kernel(params, topology, barrier_state); } +// ══════════════════════════════════════════════════════════════════════════════ +// rope_qkv producer (Stage 1: validated in isolation via the rope_only entry). +// Gated so the attention path (which does not define QWEN_FAI_ROPE_ONLY_ABI) is +// untouched. Mirrors decode_fwd.py:471-590 (the standalone rope_qkv SPMD task). +// ══════════════════════════════════════════════════════════════════════════════ +#ifdef QWEN_FAI_ROPE_ONLY_ABI +#include "kernel_operator.h" + +// rope runs on AIV lanes only; on the AIC variant of the mixed extern every symbol +// below would be unused, so the whole compute is compiled under __DAV_C220_VEC__. +// run_qwen_rope_qkv stays declared for both core types (cube gets the no-op stub). +#ifdef __DAV_C220_VEC__ + +// Arg indices for the rope_only extern ABI (paged_attention_rope_only_cce). +constexpr int32_t kRopeQProjArg = 0; +constexpr int32_t kRopeKProjArg = 1; +constexpr int32_t kRopeVProjArg = 2; +constexpr int32_t kRopeInvRmsArg = 3; +constexpr int32_t kRopeQNormArg = 4; +constexpr int32_t kRopeKNormArg = 5; +constexpr int32_t kRopeSeqLensArg = 6; +constexpr int32_t kRopeSlotMappingArg = 7; +constexpr int32_t kRopeCosArg = 8; +constexpr int32_t kRopeSinArg = 9; +constexpr int32_t kRopeQueryArg = 10; // InOut BF16: rotated Q (q_tnd_flat) +constexpr int32_t kRopeKeyCacheArg = 11; // InOut BF16: paged K cache +constexpr int32_t kRopeValueCacheArg = 12; // InOut BF16: paged V cache +constexpr int32_t kRopeCacheRowOffsetArg = 13; + +// Static batch-16 architecture constants (Qwen3-14B decode contract). +constexpr uint32_t kQwenBatch = 16; +constexpr uint32_t kQwenNumHeads = 40; +constexpr uint32_t kQwenNumKvHeads = 8; +constexpr uint32_t kQwenHeadsPerKv = 5; // Q_PER_KV +constexpr float kQwenRopeEps = 1.0e-6F; +constexpr float kQwenRopeHeadDimInv = 1.0F / static_cast(kQwenFaiHeadDim); // 1/128 + +static __aicore__ __attribute__((always_inline)) void qwen_mte2_to_v() { + event_t e = static_cast(GetTPipePtr()->FetchEventID(AscendC::HardEvent::MTE2_V)); + AscendC::SetFlag(e); + AscendC::WaitFlag(e); +} +static __aicore__ __attribute__((always_inline)) void qwen_v_to_s() { + event_t e = static_cast(GetTPipePtr()->FetchEventID(AscendC::HardEvent::V_S)); + AscendC::SetFlag(e); + AscendC::WaitFlag(e); +} +// Return fence after an S-side GetValue: pto-isa's SumLastBlockElements does the +// FULL V->S ... S->V handshake. Without this, the scalar read by GetValue is not +// committed before the next V op (Muls) consumes it, so across repeated Q-head +// calls the scalar goes stale -> x=0 -> query=0. +static __aicore__ __attribute__((always_inline)) void qwen_s_to_v() { + event_t e = static_cast(GetTPipePtr()->FetchEventID(AscendC::HardEvent::S_V)); + AscendC::SetFlag(e); + AscendC::WaitFlag(e); +} +static __aicore__ __attribute__((always_inline)) void qwen_v_to_mte3() { + event_t e = static_cast(GetTPipePtr()->FetchEventID(AscendC::HardEvent::V_MTE3)); + AscendC::SetFlag(e); + AscendC::WaitFlag(e); +} +static __aicore__ __attribute__((always_inline)) void qwen_mte3_to_v() { + event_t e = static_cast(GetTPipePtr()->FetchEventID(AscendC::HardEvent::MTE3_V)); + AscendC::SetFlag(e); + AscendC::WaitFlag(e); +} +// Drain a UB->GM store (MTE3) before the next GM->UB load (MTE2). On C220 MTE2/MTE3 +// share the external-memory DMA path; without this fence a rapid load issued right +// after a store can tear the in-flight store (manifests as whole-head GM garbage, +// intermittent, proportional to write count). qwen_mte3_to_v alone only fences V. +static __aicore__ __attribute__((always_inline)) void qwen_mte3_to_mte2() { + event_t e = static_cast(GetTPipePtr()->FetchEventID(AscendC::HardEvent::MTE3_MTE2)); + AscendC::SetFlag(e); + AscendC::WaitFlag(e); +} + +// Normalize one 128-element head row and apply NeoX half-split RoPE, matching +// decode_fwd.py rope_qkv exactly: +// x = proj_row * inv_rms (deferred input-RMSNorm factor; gamma upstream) +// ih = rsqrt(sum(x^2)/128 + eps) (per-head QK-norm) +// x = (x * weight) * ih (weight broadcast; kernel's mul order) +// rot_lo = lo*cos_lo - hi*sin_lo ; rot_hi = hi*cos_hi + lo*sin_hi (cos/sin cols dup) +// All intermediates FP32; the rotated row is left FP32 in `out`. +static __aicore__ __attribute__((always_inline)) void qwen_normalize_rope_row( + const AscendC::LocalTensor &x, + const AscendC::LocalTensor &weight, + const AscendC::LocalTensor &cos, + const AscendC::LocalTensor &sin, + const AscendC::LocalTensor &square, + AscendC::LocalTensor &reduce, + const AscendC::LocalTensor &reduce_work, + AscendC::LocalTensor &out, + const AscendC::LocalTensor &tmp, + float inv_rms) { + // C220 V-pipe does NOT serialize dependent ops on the same buffer: a later op can + // read a buffer before an earlier op's write commits (RAW hazard). EVERY place a V + // op reads another V op's output needs a PipeBarrier() between them (mirror + // the vendor CombineScale, which barriers between every dependent V op). + AscendC::Muls(x, x, inv_rms, kQwenFaiHeadDim); + AscendC::PipeBarrier(); // square must read the inv_rms-scaled x + AscendC::Mul(square, x, x, kQwenFaiHeadDim); + AscendC::PipeBarrier(); + // Tree-reduce the 128 squares to 8 partial sums via in-place Adds (64/32/16/8). + AscendC::Add(square, square, square[kQwenFaiHeadDim / 2], kQwenFaiHeadDim / 2); // 128 -> 64 + AscendC::PipeBarrier(); + AscendC::Add(square, square, square[kQwenFaiHeadDim / 4], kQwenFaiHeadDim / 4); // 64 -> 32 + AscendC::PipeBarrier(); + AscendC::Add(square, square, square[kQwenFaiHeadDim / 8], kQwenFaiHeadDim / 8); // 32 -> 16 + AscendC::PipeBarrier(); + AscendC::Add(square, square, square[kQwenFaiHeadDim / 16], kQwenFaiHeadDim / 16); // 16 -> 8 + AscendC::PipeBarrier(); + // Sum the 8 partial sums in S with LITERAL CONSTANT-index GetValue. A DYNAMIC-index + // GetValue on a UB LocalTensor misreads on C220 (~128x); constant-index UB GetValue + // is fine. Reducing 8->1 in V needs count<8 Adds, which HANG on C220, so sum in S. + qwen_v_to_s(); + float ss = square.GetValue(0) + square.GetValue(1) + square.GetValue(2) + square.GetValue(3) + + square.GetValue(4) + square.GetValue(5) + square.GetValue(6) + square.GetValue(7); + reduce.SetValue(0, ss); + qwen_s_to_v(); + AscendC::Muls(reduce, reduce, kQwenRopeHeadDimInv, 8); + AscendC::PipeBarrier(); + AscendC::Adds(reduce, reduce, kQwenRopeEps, 8); + AscendC::PipeBarrier(); + AscendC::Rsqrt(reduce, reduce, 8); + AscendC::PipeBarrier(); + qwen_v_to_s(); + float inv_head = reduce.GetValue(0); + qwen_s_to_v(); // commit the scalar before V consumes it (full V->S->S->V handshake) + AscendC::Mul(x, x, weight, kQwenFaiHeadDim); + AscendC::PipeBarrier(); + AscendC::Muls(x, x, inv_head, kQwenFaiHeadDim); + AscendC::PipeBarrier(); + AscendC::Mul(out, x, cos, kQwenFaiHeadDim / 2); // lo * cos_lo + AscendC::PipeBarrier(); + AscendC::Mul(tmp, x[kQwenFaiHeadDim / 2], sin, kQwenFaiHeadDim / 2); + AscendC::PipeBarrier(); + AscendC::Sub(out, out, tmp, kQwenFaiHeadDim / 2); // rot_lo + AscendC::PipeBarrier(); + AscendC::Mul(out[kQwenFaiHeadDim / 2], x[kQwenFaiHeadDim / 2], cos[kQwenFaiHeadDim / 2], + kQwenFaiHeadDim / 2); // hi * cos_hi + AscendC::PipeBarrier(); + AscendC::Mul(tmp[kQwenFaiHeadDim / 2], x, sin[kQwenFaiHeadDim / 2], kQwenFaiHeadDim / 2); + AscendC::PipeBarrier(); + AscendC::Add(out[kQwenFaiHeadDim / 2], out[kQwenFaiHeadDim / 2], tmp[kQwenFaiHeadDim / 2], + kQwenFaiHeadDim / 2); // rot_hi + AscendC::PipeBarrier(); +} + +#endif // __DAV_C220_VEC__ + +// Per (kv_head, batch): QK-norm + RoPE on Q (5 heads) and K (1 head) -> BF16 query +// buffer and paged K cache; V scaled by inv_rms -> paged V cache (no norm, no rope). +// On the cube (AIC) variant this is a no-op — RoPE is vector-only. +static __aicore__ void run_qwen_rope_qkv(__gm__ int64_t *args) { +#ifdef __DAV_C220_VEC__ + constexpr uint32_t kItems = kQwenBatch * kQwenNumKvHeads; // 128 + + AscendC::GlobalTensor q_proj; + AscendC::GlobalTensor k_proj; + AscendC::GlobalTensor v_proj; + AscendC::GlobalTensor inv_rms; + AscendC::GlobalTensor q_norm; + AscendC::GlobalTensor k_norm; + AscendC::GlobalTensor seq_lens; + AscendC::GlobalTensor slot_mapping; + AscendC::GlobalTensor rope_cos; + AscendC::GlobalTensor rope_sin; + AscendC::GlobalTensor query; + AscendC::GlobalTensor key_cache; + AscendC::GlobalTensor value_cache; + q_proj.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(tensor_data(args, kRopeQProjArg))); + k_proj.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(tensor_data(args, kRopeKProjArg))); + v_proj.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(tensor_data(args, kRopeVProjArg))); + inv_rms.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(tensor_data(args, kRopeInvRmsArg))); + q_norm.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(tensor_data(args, kRopeQNormArg))); + k_norm.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(tensor_data(args, kRopeKNormArg))); + seq_lens.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(tensor_data(args, kRopeSeqLensArg))); + slot_mapping.SetGlobalBuffer(reinterpret_cast<__gm__ int32_t *>(tensor_data(args, kRopeSlotMappingArg))); + rope_cos.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(tensor_data(args, kRopeCosArg))); + rope_sin.SetGlobalBuffer(reinterpret_cast<__gm__ float *>(tensor_data(args, kRopeSinArg))); + query.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(tensor_data(args, kRopeQueryArg))); + key_cache.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(tensor_data(args, kRopeKeyCacheArg))); + value_cache.SetGlobalBuffer(reinterpret_cast<__gm__ bfloat16_t *>(tensor_data(args, kRopeValueCacheArg))); + + AscendC::TPipe pipe; + AscendC::TBuf buf; + pipe.InitBuffer(buf, 8192); + AscendC::LocalTensor mem = buf.Get(); + AscendC::LocalTensor x = mem[0]; + AscendC::LocalTensor square = mem[128]; + AscendC::LocalTensor weight = mem[256]; + AscendC::LocalTensor cos = mem[384]; + AscendC::LocalTensor sin = mem[512]; + AscendC::LocalTensor out = mem[640]; + AscendC::LocalTensor tmp = mem[768]; + AscendC::LocalTensor reduce = mem[896]; + AscendC::LocalTensor reduce_work = mem[1024]; + AscendC::LocalTensor out_bf16 = buf.Get()[3072]; + + uint64_t cache_row_offset = static_cast(args[kRopeCacheRowOffsetArg]); + uint32_t physical_lane = + static_cast(get_block_idx(args)) * 2 + static_cast(get_sub_block_id(args)); + uint32_t physical_lanes = static_cast(get_block_num(args)) * 2; + + for (uint32_t item = physical_lane; item < kItems; item += physical_lanes) { + uint32_t kv_head = item / kQwenBatch; + uint32_t batch = item - kv_head * kQwenBatch; + uint32_t pos = static_cast(seq_lens.GetValue(batch)) - 1; + uint32_t slot = static_cast(slot_mapping.GetValue(batch)); + float inv_rms_value = inv_rms.GetValue(batch); + + AscendC::DataCopy(cos, rope_cos[static_cast(pos) * kQwenFaiHeadDim], kQwenFaiHeadDim); + AscendC::DataCopy(sin, rope_sin[static_cast(pos) * kQwenFaiHeadDim], kQwenFaiHeadDim); + + AscendC::DataCopy(weight, q_norm, kQwenFaiHeadDim); + qwen_mte2_to_v(); + for (uint32_t q = 0; q < kQwenHeadsPerKv; ++q) { + uint32_t head = kv_head * kQwenHeadsPerKv + q; + AscendC::DataCopy(x, q_proj[batch * kQwenNumHeads * kQwenFaiHeadDim + head * kQwenFaiHeadDim], + kQwenFaiHeadDim); + qwen_mte2_to_v(); + qwen_normalize_rope_row(x, weight, cos, sin, square, reduce, reduce_work, out, tmp, inv_rms_value); + AscendC::Cast(out_bf16, out, AscendC::RoundMode::CAST_ROUND, kQwenFaiHeadDim); + qwen_v_to_mte3(); + AscendC::DataCopy(query[batch * kQwenNumHeads * kQwenFaiHeadDim + head * kQwenFaiHeadDim], out_bf16, + kQwenFaiHeadDim); + qwen_mte3_to_v(); + qwen_mte3_to_mte2(); + } + + AscendC::DataCopy(weight, k_norm, kQwenFaiHeadDim); + AscendC::DataCopy(x, k_proj[batch * kQwenNumKvHeads * kQwenFaiHeadDim + kv_head * kQwenFaiHeadDim], + kQwenFaiHeadDim); + qwen_mte2_to_v(); + qwen_normalize_rope_row(x, weight, cos, sin, square, reduce, reduce_work, out, tmp, inv_rms_value); + AscendC::Cast(out_bf16, out, AscendC::RoundMode::CAST_ROUND, kQwenFaiHeadDim); + qwen_v_to_mte3(); + uint64_t cache_row = cache_row_offset + static_cast(slot) * kQwenNumKvHeads + kv_head; + AscendC::DataCopy(key_cache[cache_row * kQwenFaiHeadDim], out_bf16, kQwenFaiHeadDim); + qwen_mte3_to_v(); + qwen_mte3_to_mte2(); + + AscendC::DataCopy(x, v_proj[batch * kQwenNumKvHeads * kQwenFaiHeadDim + kv_head * kQwenFaiHeadDim], + kQwenFaiHeadDim); + qwen_mte2_to_v(); + AscendC::Muls(x, x, inv_rms_value, kQwenFaiHeadDim); + AscendC::PipeBarrier(); // Cast must read the inv_rms-scaled x + AscendC::Cast(out_bf16, x, AscendC::RoundMode::CAST_ROUND, kQwenFaiHeadDim); + qwen_v_to_mte3(); + AscendC::DataCopy(value_cache[cache_row * kQwenFaiHeadDim], out_bf16, kQwenFaiHeadDim); + qwen_mte3_to_v(); + qwen_mte3_to_mte2(); + } + pipe.Destroy(); +#else + (void)args; +#endif +} + +#endif // QWEN_FAI_ROPE_ONLY_ABI + #endif diff --git a/models/qwen3/14b/kernels/paged_attention_cce/rope_only/entry.cpp b/models/qwen3/14b/kernels/paged_attention_cce/rope_only/entry.cpp new file mode 100644 index 00000000..8f13e1e7 --- /dev/null +++ b/models/qwen3/14b/kernels/paged_attention_cce/rope_only/entry.cpp @@ -0,0 +1,38 @@ +/* + * 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. + * ----------------------------------------------------------------------------------------------------------- + */ + +#include + +#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 + +// Stage 1 isolation entry: runs ONLY the rope_qkv producer (no attention, no +// cross-core sync) so the C++ rope compute can be validated byte-exact against +// decode_fwd.py's rope_qkv before being fused into the attention kernel. +#define QWEN_FAI_ROPE_ONLY_ABI +#include "../kernel/fai_body.hpp" + +extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { + run_qwen_rope_qkv(args); +} + +#endif diff --git a/models/qwen3/14b/paged_attention_cce.py b/models/qwen3/14b/paged_attention_cce.py index 96d6ea5a..9d05b961 100644 --- a/models/qwen3/14b/paged_attention_cce.py +++ b/models/qwen3/14b/paged_attention_cce.py @@ -16,6 +16,7 @@ _KERNEL_DIR = Path(__file__).parent / "kernels" / "paged_attention_cce" _ATTENTION_ENTRY = _KERNEL_DIR / "attention" / "entry.cpp" +_ROPE_ONLY_ENTRY = _KERNEL_DIR / "rope_only" / "entry.cpp" _TILING_ENTRY = _KERNEL_DIR / "tiling" / "entry.cpp" @@ -216,6 +217,82 @@ def qwen_decode_attention_cache_offset_test( return out +# Stage 1 isolation kernel: runs ONLY the rope_qkv producer (Q/K norm + RoPE + +# KV-cache writes), no attention. Used by test_rope_qkv_cce.py to validate the +# C++ rope compute byte-exact against decode_fwd.py's rope_qkv before fusion. +# Mixed + dual_aiv_dispatch so its AIV lane model (block_idx*2 + sub_block_id) +# matches the fused kernel; AICs are a no-op (run_qwen_rope_qkv is vec-only). +@pl.jit.extern( + core_type="mixed", + aic_source=_ROPE_ONLY_ENTRY, + aiv_source=_ROPE_ONLY_ENTRY, + include_dirs=_CANN_INCLUDE_DIRS, + dual_aiv_dispatch=True, +) +def paged_attention_rope_only_cce( + q_proj: pl.Tensor, + k_proj: pl.Tensor, + v_proj: pl.Tensor, + inv_rms_states: pl.Tensor, + q_norm_weight: pl.Tensor, + k_norm_weight: pl.Tensor, + seq_lens: pl.Tensor, + slot_mapping: pl.Tensor, + rope_cos: pl.Tensor, + rope_sin: pl.Tensor, + query: pl.Out[pl.Tensor], + key_cache: pl.InOut[pl.Tensor], + value_cache: pl.InOut[pl.Tensor], + cache_row_offset: pl.Scalar[pl.INDEX], +) -> pl.Tensor: ... + + +@pl.jit +def qwen_decode_rope_only_cce( + q_proj: pl.Tensor, + k_proj: pl.Tensor, + v_proj: pl.Tensor, + inv_rms_states: pl.Tensor, + q_norm_weight: pl.Tensor, + k_norm_weight: pl.Tensor, + seq_lens: pl.Tensor, + slot_mapping: pl.Tensor, + rope_cos: pl.Tensor, + rope_sin: pl.Tensor, + query: pl.Out[pl.Tensor], + key_cache: pl.InOut[pl.Tensor], + value_cache: pl.InOut[pl.Tensor], +) -> pl.Tensor: + """Stage 1 driver: run the fused-prologue rope_qkv in isolation (no attention). + + ``cache_row_offset`` is fixed at 0 (single-layer, slot 0 base) so the driver + has no scalar parameter and the harness needs only tensor specs. + """ + # sync_start=True is mandatory for a mixed + dual_aiv_dispatch extern: it makes + # all 24 AIC blocks (and their 2 AIV sub-lanes) launch atomically so the AIC/AIV + # pairing is coherent. Without it the dispatch 507018s on device. + with pl.spmd( + DEFAULT_BLOCK_DIM, name_hint="rope_only", allow_early_resolve=True, sync_start=True + ) as _rope_tid: + query = paged_attention_rope_only_cce( + q_proj, + k_proj, + v_proj, + inv_rms_states, + q_norm_weight, + k_norm_weight, + seq_lens, + slot_mapping, + rope_cos, + rope_sin, + query, + key_cache, + value_cache, + 0, + ) + return query + + __all__ = [ "BATCH", "BLOCK_SIZE", @@ -229,7 +306,9 @@ def qwen_decode_attention_cache_offset_test( "WORKSPACE_BYTES", "build_paged_attention_metadata", "paged_attention_cce", + "paged_attention_rope_only_cce", "paged_attention_tiling_cce", "qwen_decode_attention_cache_offset_test", "qwen_decode_attention_cce", + "qwen_decode_rope_only_cce", ] diff --git a/models/qwen3/14b/test_rope_qkv_cce.py b/models/qwen3/14b/test_rope_qkv_cce.py new file mode 100644 index 00000000..b36aec0c --- /dev/null +++ b/models/qwen3/14b/test_rope_qkv_cce.py @@ -0,0 +1,163 @@ +# 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. +# ----------------------------------------------------------------------------------------------------------- +# ci: no-sim +"""Stage 1 on-device correctness driver for the fused-prologue rope_qkv kernel. + +Runs ONLY the C++ rope compute (Q/K-norm + RoPE + KV-cache writes) — no attention — +and validates its GM outputs (query, and the written k_cache / v_cache rows) against a +torch reference that mirrors decode_fwd.py's rope_qkv byte-for-byte. Proving this in +isolation separates a compute defect from the cross-core sync added in Stage 2. +""" + +import argparse + +import torch + +from golden import TensorSpec, run_jit +from paged_attention_cce import ( + BATCH, + BLOCK_SIZE, + HEAD_DIM, + KV_HIDDEN, + NUM_HEADS, + NUM_KV_HEADS, + SUPPORTED_PLATFORMS, + qwen_decode_rope_only_cce, +) + +HIDDEN = NUM_HEADS * HEAD_DIM # 5120 +HALF = HEAD_DIM // 2 # 64 +EPS = 1.0e-6 +HEAD_DIM_INV = 1.0 / HEAD_DIM +MAX_SEQ = 512 +NUM_BLOCKS = 1 # one page (BLOCK_SIZE=128 tokens) covers the 16 decode slots + + +def _rope_tables() -> tuple[torch.Tensor, torch.Tensor]: + """NeoX half-split RoPE tables (cols [0:64] and [64:128] duplicated).""" + posv = torch.arange(MAX_SEQ).float().unsqueeze(1) + inv_freq = 1.0 / (1.0e4 ** (torch.arange(0, HALF).float() / HALF)) + ang = posv * inv_freq.unsqueeze(0) + return torch.cat([ang.cos(), ang.cos()], dim=1).float(), torch.cat([ang.sin(), ang.sin()], dim=1).float() + + +def build_specs(seed: int = 1234) -> list[TensorSpec]: + torch.manual_seed(seed) + g = torch.Generator().manual_seed(seed) + + def rn(shape, std=1.0): + return torch.empty(shape).normal_(0.0, std, generator=g) + + rope_cos, rope_sin = _rope_tables() + seq_lens = torch.randint(1, MAX_SEQ + 1, (BATCH,), generator=g, dtype=torch.int32) + # Distinct physical token slots per batch (no aliasing); cache_row = slot*NUM_KV_HEADS + ki. + slot_mapping = torch.arange(BATCH, dtype=torch.int32) + inv_rms = torch.rand(BATCH, generator=g) * 0.9 + 0.3 # deferred input-RMSNorm factor, > 0 + + return [ + TensorSpec("q_proj", [BATCH, HIDDEN], torch.float32, init_value=rn([BATCH, HIDDEN], 0.02)), + TensorSpec("k_proj", [BATCH, KV_HIDDEN], torch.float32, init_value=rn([BATCH, KV_HIDDEN], 0.02)), + TensorSpec("v_proj", [BATCH, KV_HIDDEN], torch.float32, init_value=rn([BATCH, KV_HIDDEN], 0.02)), + TensorSpec("inv_rms_states", [BATCH], torch.float32, init_value=inv_rms), + TensorSpec("q_norm_weight", [1, HEAD_DIM], torch.float32, init_value=rn([1, HEAD_DIM], 0.1) + 1.0), + TensorSpec("k_norm_weight", [1, HEAD_DIM], torch.float32, init_value=rn([1, HEAD_DIM], 0.1) + 1.0), + TensorSpec("seq_lens", [BATCH], torch.int32, init_value=seq_lens), + TensorSpec("slot_mapping", [BATCH], torch.int32, init_value=slot_mapping), + TensorSpec("rope_cos", [MAX_SEQ, HEAD_DIM], torch.float32, init_value=rope_cos), + TensorSpec("rope_sin", [MAX_SEQ, HEAD_DIM], torch.float32, init_value=rope_sin), + TensorSpec("query", [BATCH, NUM_HEADS, HEAD_DIM], torch.bfloat16, is_output=True), + TensorSpec( + "key_cache", + [NUM_BLOCKS, BLOCK_SIZE, KV_HIDDEN], + torch.bfloat16, + is_output=True, + init_value=rn([NUM_BLOCKS, BLOCK_SIZE, KV_HIDDEN], 0.01).to(torch.bfloat16), + ), + TensorSpec( + "value_cache", + [NUM_BLOCKS, BLOCK_SIZE, KV_HIDDEN], + torch.bfloat16, + is_output=True, + init_value=(rn([NUM_BLOCKS, BLOCK_SIZE, KV_HIDDEN], 0.02) + 0.3).to(torch.bfloat16), + ), + ] + + +def golden_rope(values: dict[str, torch.Tensor]) -> None: + """Torch reference mirroring decode_fwd.py rope_qkv (and golden_decode_layer's rope part). + + Kernel mul order is preserved: normed = (proj*inv_rms * norm_w) * per_head_inv_rms. + """ + inv_rms = values["inv_rms_states"].float().reshape(BATCH, 1) + qn = values["q_norm_weight"].float().reshape(HEAD_DIM) + kn = values["k_norm_weight"].float().reshape(HEAD_DIM) + seq_lens = values["seq_lens"] + slot_mapping = values["slot_mapping"] + rope_cos = values["rope_cos"].float() + rope_sin = values["rope_sin"].float() + + def norm_rope(proj: torch.Tensor, weight: torch.Tensor, n_heads: int) -> torch.Tensor: + raw = (proj.float() * inv_rms).reshape(BATCH, n_heads, HEAD_DIM) + inv_head = torch.rsqrt(raw.pow(2).sum(-1, keepdim=True) * HEAD_DIM_INV + EPS) + normed = (raw * weight) * inv_head # (proj*inv_rms)*norm_w, then *per-head rms + lo, hi = normed[..., :HALF], normed[..., HALF:] + cos = rope_cos[seq_lens.long() - 1] # [BATCH, HEAD_DIM] + sin = rope_sin[seq_lens.long() - 1] + clo = cos[:, :HALF].unsqueeze(1) # [BATCH, 1, HALF] broadcast over heads + chi = cos[:, HALF:].unsqueeze(1) + slo = sin[:, :HALF].unsqueeze(1) + shi = sin[:, HALF:].unsqueeze(1) + rot_lo = lo * clo - hi * slo + rot_hi = hi * chi + lo * shi + return torch.cat([rot_lo, rot_hi], dim=-1) + + values["query"] = norm_rope(values["q_proj"], qn, NUM_HEADS).to(torch.bfloat16) + k_rot = norm_rope(values["k_proj"], kn, NUM_KV_HEADS) # [BATCH, NUM_KV_HEADS, HEAD_DIM] + v_heads = (values["v_proj"].float() * inv_rms).reshape(BATCH, NUM_KV_HEADS, HEAD_DIM) + + key_cache = values["key_cache"].clone() + value_cache = values["value_cache"].clone() + for b in range(BATCH): + slot = int(slot_mapping[b].item()) + block = slot // BLOCK_SIZE + off = slot % BLOCK_SIZE + for ki in range(NUM_KV_HEADS): + cols = slice(ki * HEAD_DIM, (ki + 1) * HEAD_DIM) + key_cache[block, off, cols] = k_rot[b, ki].to(torch.bfloat16) + value_cache[block, off, cols] = v_heads[b, ki].to(torch.bfloat16) + values["key_cache"] = key_cache + values["value_cache"] = value_cache + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("-p", "--platform", default="a2a3", choices=SUPPORTED_PLATFORMS) + parser.add_argument("-d", "--device", type=int, default=0) + parser.add_argument("--compile-only", action="store_true") + args = parser.parse_args() + + specs = build_specs() + result = run_jit( + fn=qwen_decode_rope_only_cce, + specs=specs, + golden_fn=golden_rope, + runtime_cfg={"platform": args.platform, "device_id": args.device}, + compile_only=args.compile_only or args.platform.endswith("sim"), + rtol=5e-3, + atol=2e-2, + save_data=False, + ) + if not result.passed: + if result.error: + print(result.error) + raise SystemExit(1) + + +if __name__ == "__main__": + main() From 59d2d55e467ed688edce881dec237480b15583e5 Mon Sep 17 00:00:00 2001 From: Youhezhen Date: Thu, 16 Jul 2026 19:06:20 -0700 Subject: [PATCH 2/3] =?UTF-8?q?feat(qwen3):=20Stage=202=20=E2=80=94=20fuse?= =?UTF-8?q?=20rope=5Fqkv=20prologue=20into=20paged=5Fattention=5Fcce?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fuses the QKV prologue (QK-norm + RoPE + paged KV-cache writes) into the mixed-core attention extern so the decode path no longer runs a separate rope_qkv SPMD task. Stage 1's isolated rope compute (proven byte-exact on a2a3, commit fa4832a) is now the AIV-side prologue of the attention launch. fai_body.hpp - run_qwen_fai arg indices switched by QWEN_FAI_ATTENTION_ONLY_ABI: standalone attention keeps 0..7; the fused extern prepends the rope inputs (q_proj..rope_sin 0..9, block_table 10) so the attention tensors shift to out=11, query/k/v=12..14, workspace/metadata=15..16, cache_row_offset=17. - rope_qkv producer guarded by !QWEN_FAI_ATTENTION_ONLY_ABI (compiled for both the fused entry and the rope_only isolation entry); its output arg indices switch by QWEN_FAI_ROPE_ONLY_ABI. The Stage 1 compute fixes (UB constant GetValue, V-pipe PipeBarrier between every dependent op incl MTE3->MTE2) carry over. attention/entry.cpp (fused) - acquire_metadata -> run_qwen_rope_qkv -> sync_qwen_rope_producers -> run_qwen_fai. - sync_qwen_rope_producers: AIV lanes barrier over the rope-ready metadata slots (pto::SYNCALL_SOFT_AIV_BARRIER), then one sub-lane per AIC pair signals its AIC via CrossCoreSetFlag<0x2, PIPE_MTE3>; the AIC CrossCoreWaitFlag before attention. PIPE_MTE3 so the signal fires only after rope's UB->GM stores flush (the cross-pipe GM visibility a plain SYNCALL lacks). MODE=0x2 + pto primitives verified valid on a2a3. attention_only/entry.cpp (NEW) - defines QWEN_FAI_ATTENTION_ONLY_ABI -> run_qwen_fai only (no rope/sync); preserves the standalone attention ABI for test_paged_attention_cce.py. paged_attention_cce.py - paged_attention_cce is now the fused 18-arg extern (attention/entry.cpp). - paged_attention_only_cce (NEW) = the old 8-arg attention ABI (attention_only entry); qwen_decode_attention_cce + cache_offset_test driver route to it. - METADATA_BYTES bumped for the rope-ready soft-barrier region (27840 -> 29376, matches metadata_layout.h kMetadataBytes). metadata_layout.h + tiling/entry.cpp - rope-ready slot region constants (48 lanes x 32B) after the aligned FAI barrier; the tiling kernel zeroes it (clear_rope_ready) each launch. decode_fwd.py - deleted the standalone rope_qkv SPMD task; fa_fused now depends on [q_proj,k_proj,v_proj,rms,pa_tiling,attn_out_seed,mlp_out_seed] and calls the fused paged_attention_cce passing the raw projections + norm/rope/slot inputs, with attn_out_tnd as the Out return and q_tnd/k_cache/v_cache as InOut (written by the AIV prologue). Local validation: test_rope_qkv_cce (rope_only path) still PASSES on a2a3; both extern paths codegen-clean. The fused serving path (cross-core sync) is validated via CI (test_qwen3_accuracy). --- models/qwen3/14b/decode_fwd.py | 142 +++--------------- .../paged_attention_cce/attention/entry.cpp | 55 ++++++- .../attention_only/entry.cpp | 50 ++++++ .../paged_attention_cce/kernel/fai_body.hpp | 63 ++++++-- .../kernel/metadata_layout.h | 14 +- .../paged_attention_cce/tiling/entry.cpp | 19 ++- models/qwen3/14b/paged_attention_cce.py | 51 ++++++- 7 files changed, 243 insertions(+), 151 deletions(-) create mode 100644 models/qwen3/14b/kernels/paged_attention_cce/attention_only/entry.cpp diff --git a/models/qwen3/14b/decode_fwd.py b/models/qwen3/14b/decode_fwd.py index 0c08bdd7..dcb5fe0f 100644 --- a/models/qwen3/14b/decode_fwd.py +++ b/models/qwen3/14b/decode_fwd.py @@ -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 @@ -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, diff --git a/models/qwen3/14b/kernels/paged_attention_cce/attention/entry.cpp b/models/qwen3/14b/kernels/paged_attention_cce/attention/entry.cpp index cadbef80..d9c21c7b 100644 --- a/models/qwen3/14b/kernels/paged_attention_cce/attention/entry.cpp +++ b/models/qwen3/14b/kernels/paged_attention_cce/attention/entry.cpp @@ -25,17 +25,62 @@ extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { (void)args; } #else +#include + #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(metadata + qwen_fai_metadata::kBarrierAlignmentOffset); + uint64_t aligned_barrier = (raw_barrier + qwen_fai_metadata::kBarrierAlignmentBytes - 1) & + ~(static_cast(qwen_fai_metadata::kBarrierAlignmentBytes) - 1); + return reinterpret_cast<__gm__ int32_t *>(aligned_barrier); +} + +// AIV lanes finish the rope GM writes, barrier among themselves over the rope-ready +// metadata slots, then one sub-lane per AIC pair signals its AIC. The AIC waits for +// that flag before starting attention (which reads the rope-produced q/k/v cache). +// CrossCoreSetFlag uses PIPE_MTE3 so the signal fires only after rope's UB->GM +// stores have flushed — the cross-pipe GM visibility that a plain SYNCALL lacks. +static __aicore__ __attribute__((always_inline)) void sync_qwen_rope_producers( + __gm__ int64_t *args, + __gm__ int32_t *fai_barrier +) { + pipe_barrier(PIPE_ALL); +#if defined(__DAV_CUBE__) + AscendC::CrossCoreWaitFlag(qwen_fai_metadata::kRopeReadyCrossCoreFlag); +#elif defined(__DAV_VEC__) + __gm__ int32_t *ready = reinterpret_cast<__gm__ int32_t *>( + reinterpret_cast<__gm__ uint8_t *>(fai_barrier) + qwen_fai_metadata::kBarrierBytes + ); + uint32_t physical_lane = static_cast(get_block_idx(args)) * 2 + + static_cast(get_sub_block_id(args)); + __ubuf__ int32_t *ub_workspace = reinterpret_cast<__ubuf__ int32_t *>(0x3000); + pto::SYNCALL_SOFT_AIV_BARRIER( + ready, + ub_workspace, + qwen_fai_metadata::kRopeReadySlotCount, + physical_lane + ); + if (get_sub_block_id(args) == 0) { + AscendC::CrossCoreSetFlag<0x2, PIPE_MTE3>(qwen_fai_metadata::kRopeReadyCrossCoreFlag); + } +#endif + pipe_barrier(PIPE_ALL); +} + +} // namespace + extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { - GM_ADDR metadata = tensor_data(args, 6); + GM_ADDR metadata = tensor_data(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(metadata + qwen_fai_metadata::kBarrierAlignmentOffset); - uint64_t aligned_barrier = (raw_barrier + qwen_fai_metadata::kBarrierAlignmentBytes - 1) & - ~(static_cast(qwen_fai_metadata::kBarrierAlignmentBytes) - 1); - run_qwen_fai(args, reinterpret_cast<__gm__ int32_t *>(aligned_barrier)); + run_qwen_fai(args, fai_barrier); } else { run_qwen_fai(args); } diff --git a/models/qwen3/14b/kernels/paged_attention_cce/attention_only/entry.cpp b/models/qwen3/14b/kernels/paged_attention_cce/attention_only/entry.cpp new file mode 100644 index 00000000..9f07da18 --- /dev/null +++ b/models/qwen3/14b/kernels/paged_attention_cce/attention_only/entry.cpp @@ -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 + +#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(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(metadata + qwen_fai_metadata::kBarrierAlignmentOffset); + uint64_t aligned_barrier = (raw_barrier + qwen_fai_metadata::kBarrierAlignmentBytes - 1) & + ~(static_cast(qwen_fai_metadata::kBarrierAlignmentBytes) - 1); + run_qwen_fai(args, reinterpret_cast<__gm__ int32_t *>(aligned_barrier)); + } else { + run_qwen_fai(args); + } +} + +#endif diff --git a/models/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hpp b/models/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hpp index f771271e..26b296f0 100644 --- a/models/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hpp +++ b/models/qwen3/14b/kernels/paged_attention_cce/kernel/fai_body.hpp @@ -35,6 +35,31 @@ constexpr uint64_t kQwenFaiHeadDim = 128; +// Arg indices into the kernel args[] array, switched by QWEN_FAI_ATTENTION_ONLY_ABI. +// The standalone attention extern passes (query, k/v_cache, block_table, out, +// workspace, metadata, cache_row_offset) at 0..7. The fused extern prepends the rope +// inputs (q_proj..rope_sin at 0..9, block_table 10), so the attention tensors shift +// to 11..17 (out first — it is the extern return / first writable). +#ifdef QWEN_FAI_ATTENTION_ONLY_ABI +constexpr int32_t kQueryArg = 0; +constexpr int32_t kKeyCacheArg = 1; +constexpr int32_t kValueCacheArg = 2; +constexpr int32_t kBlockTableArg = 3; +constexpr int32_t kOutArg = 4; +constexpr int32_t kWorkspaceArg = 5; +constexpr int32_t kMetadataArg = 6; +constexpr int32_t kCacheRowOffsetArg = 7; +#else +constexpr int32_t kQueryArg = 12; +constexpr int32_t kKeyCacheArg = 13; +constexpr int32_t kValueCacheArg = 14; +constexpr int32_t kBlockTableArg = 10; +constexpr int32_t kOutArg = 11; +constexpr int32_t kWorkspaceArg = 15; +constexpr int32_t kMetadataArg = 16; +constexpr int32_t kCacheRowOffsetArg = 17; +#endif + static __aicore__ __attribute__((always_inline)) void acquire_qwen_fai_metadata(GM_ADDR metadata) { uint64_t first_line = reinterpret_cast(metadata) & ~(static_cast(qwen_fai_metadata::kDcciLineBytes) - 1); @@ -125,24 +150,24 @@ run_qwen_fai(__gm__ int64_t *args, __gm__ int32_t *barrier_state = nullptr) { FaiKernel::MaskType::NO_MASK, FaiKernel::inputLayout::TND>; using Kernel = std::conditional_t; - GM_ADDR metadata = tensor_data(args, 6); - uint64_t cache_row_offset = static_cast(args[7]); + GM_ADDR metadata = tensor_data(args, kMetadataArg); + uint64_t cache_row_offset = static_cast(args[kCacheRowOffsetArg]); uint64_t cache_byte_offset = cache_row_offset * kQwenFaiHeadDim * sizeof(uint16_t); - GM_ADDR key = tensor_data(args, 1) + cache_byte_offset; - GM_ADDR value = tensor_data(args, 2) + cache_byte_offset; + GM_ADDR key = tensor_data(args, kKeyCacheArg) + cache_byte_offset; + GM_ADDR value = tensor_data(args, kValueCacheArg) + cache_byte_offset; FAIKernelParams params{ - tensor_data(args, 0), + tensor_data(args, kQueryArg), key, value, nullptr, nullptr, - tensor_data(args, 3), + tensor_data(args, kBlockTableArg), metadata + qwen_fai_metadata::kCumulativeQOffset, metadata + qwen_fai_metadata::kKvLengthsOffset, - tensor_data(args, 4), + tensor_data(args, kOutArg), nullptr, - tensor_data(args, 5), + tensor_data(args, kWorkspaceArg), metadata + qwen_fai_metadata::kTilingOffset, nullptr }; @@ -159,11 +184,12 @@ run_qwen_fai(__gm__ int64_t *args, __gm__ int32_t *barrier_state = nullptr) { } // ══════════════════════════════════════════════════════════════════════════════ -// rope_qkv producer (Stage 1: validated in isolation via the rope_only entry). -// Gated so the attention path (which does not define QWEN_FAI_ROPE_ONLY_ABI) is -// untouched. Mirrors decode_fwd.py:471-590 (the standalone rope_qkv SPMD task). +// rope_qkv producer. Compiled for both the fused attention entry (AIV runs rope, +// then cross-core-syncs the AIC attention) and the rope_only isolation entry; the +// standalone attention_only entry defines QWEN_FAI_ATTENTION_ONLY_ABI to skip it. +// On the AIC (cube) variant run_qwen_rope_qkv is a no-op (rope is vector-only). // ══════════════════════════════════════════════════════════════════════════════ -#ifdef QWEN_FAI_ROPE_ONLY_ABI +#ifndef QWEN_FAI_ATTENTION_ONLY_ABI #include "kernel_operator.h" // rope runs on AIV lanes only; on the AIC variant of the mixed extern every symbol @@ -171,7 +197,7 @@ run_qwen_fai(__gm__ int64_t *args, __gm__ int32_t *barrier_state = nullptr) { // run_qwen_rope_qkv stays declared for both core types (cube gets the no-op stub). #ifdef __DAV_C220_VEC__ -// Arg indices for the rope_only extern ABI (paged_attention_rope_only_cce). +// Rope input args (shared by rope_only and the fused extern): q_proj..rope_sin at 0..9. constexpr int32_t kRopeQProjArg = 0; constexpr int32_t kRopeKProjArg = 1; constexpr int32_t kRopeVProjArg = 2; @@ -182,10 +208,19 @@ constexpr int32_t kRopeSeqLensArg = 6; constexpr int32_t kRopeSlotMappingArg = 7; constexpr int32_t kRopeCosArg = 8; constexpr int32_t kRopeSinArg = 9; +// Rope output args differ by ABI: rope_only packs them right after sin (10..13); +// the fused extern has block_table at 10 and the attention tensors at 12..14,17. +#ifdef QWEN_FAI_ROPE_ONLY_ABI constexpr int32_t kRopeQueryArg = 10; // InOut BF16: rotated Q (q_tnd_flat) constexpr int32_t kRopeKeyCacheArg = 11; // InOut BF16: paged K cache constexpr int32_t kRopeValueCacheArg = 12; // InOut BF16: paged V cache constexpr int32_t kRopeCacheRowOffsetArg = 13; +#else +constexpr int32_t kRopeQueryArg = 12; +constexpr int32_t kRopeKeyCacheArg = 13; +constexpr int32_t kRopeValueCacheArg = 14; +constexpr int32_t kRopeCacheRowOffsetArg = 17; +#endif // Static batch-16 architecture constants (Qwen3-14B decode contract). constexpr uint32_t kQwenBatch = 16; @@ -417,6 +452,6 @@ static __aicore__ void run_qwen_rope_qkv(__gm__ int64_t *args) { #endif } -#endif // QWEN_FAI_ROPE_ONLY_ABI +#endif // !QWEN_FAI_ATTENTION_ONLY_ABI (rope_qkv producer) #endif diff --git a/models/qwen3/14b/kernels/paged_attention_cce/kernel/metadata_layout.h b/models/qwen3/14b/kernels/paged_attention_cce/kernel/metadata_layout.h index b21f1991..6e0f0bf3 100644 --- a/models/qwen3/14b/kernels/paged_attention_cce/kernel/metadata_layout.h +++ b/models/qwen3/14b/kernels/paged_attention_cce/kernel/metadata_layout.h @@ -28,11 +28,19 @@ constexpr uint32_t kBarrierSlotBytes = 512; constexpr uint32_t kBarrierSlotWords = kBarrierSlotBytes / sizeof(int32_t); constexpr uint32_t kBarrierSlotCount = 48; constexpr uint32_t kBarrierBytes = kBarrierSlotCount * kBarrierSlotBytes; -constexpr uint32_t kMetadataBytes = 27840; +// Producer-only RoPE-ready soft-barrier region, laid out right after the aligned +// FAI barrier. The fused kernel's AIV lanes barrier among themselves here, then one +// sub-lane per AIC pair signals the AIC attention via CrossCoreSetFlag. +constexpr uint32_t kRopeReadySlotBytes = 32; +constexpr uint32_t kRopeReadySlotWords = kRopeReadySlotBytes / sizeof(int32_t); +constexpr uint32_t kRopeReadySlotCount = 48; +constexpr uint32_t kRopeReadyBytes = kRopeReadySlotCount * kRopeReadySlotBytes; +constexpr uint16_t kRopeReadyCrossCoreFlag = 0; +constexpr uint32_t kMetadataBytes = 29376; static_assert( - kBarrierAlignmentOffset + kBarrierAlignmentBytes - 1 + kBarrierBytes <= kMetadataBytes, - "metadata buffer does not cover the aligned barrier region" + kBarrierAlignmentOffset + kBarrierAlignmentBytes - 1 + kBarrierBytes + kRopeReadyBytes <= kMetadataBytes, + "metadata buffer does not cover the aligned FAI and RoPE-ready regions" ); } // namespace qwen_fai_metadata diff --git a/models/qwen3/14b/kernels/paged_attention_cce/tiling/entry.cpp b/models/qwen3/14b/kernels/paged_attention_cce/tiling/entry.cpp index 3c48e349..1e0ac326 100644 --- a/models/qwen3/14b/kernels/paged_attention_cce/tiling/entry.cpp +++ b/models/qwen3/14b/kernels/paged_attention_cce/tiling/entry.cpp @@ -58,6 +58,21 @@ static __aicore__ void clear_barrier(__gm__ int32_t *barrier) { dsb(DSB_DDR); } +static __aicore__ __attribute__((always_inline)) __gm__ int32_t *rope_ready_data(__gm__ int32_t *barrier) { + return reinterpret_cast<__gm__ int32_t *>( + reinterpret_cast<__gm__ uint8_t *>(barrier) + qwen_fai_metadata::kBarrierBytes + ); +} + +static __aicore__ void clear_rope_ready(__gm__ int32_t *ready) { + for (uint32_t slot = 0; slot < qwen_fai_metadata::kRopeReadySlotCount; ++slot) { + __gm__ int32_t *slot_data = ready + slot * qwen_fai_metadata::kRopeReadySlotWords; + slot_data[0] = 0; + dcci(slot_data, SINGLE_CACHE_LINE, CACHELINE_OUT); + } + dsb(DSB_DDR); +} + static __aicore__ void flush_metadata_prefix(__gm__ uint8_t *metadata) { uint64_t first_line = reinterpret_cast(metadata) & ~(static_cast(qwen_fai_metadata::kDcciLineBytes) - 1); @@ -82,7 +97,9 @@ extern "C" __aicore__ void kernel_entry(__gm__ int64_t *args) { uint32_t max_blocks_per_batch = static_cast(args[2]); uint32_t num_blocks = static_cast(args[3]); - clear_barrier(barrier_data(metadata)); + __gm__ int32_t *barrier = barrier_data(metadata); + clear_barrier(barrier); + clear_rope_ready(rope_ready_data(barrier)); if (batch == 0 || batch > qwen_fai_tiler::kMaxBatch) { return; } diff --git a/models/qwen3/14b/paged_attention_cce.py b/models/qwen3/14b/paged_attention_cce.py index 9d05b961..2bfef69f 100644 --- a/models/qwen3/14b/paged_attention_cce.py +++ b/models/qwen3/14b/paged_attention_cce.py @@ -16,6 +16,7 @@ _KERNEL_DIR = Path(__file__).parent / "kernels" / "paged_attention_cce" _ATTENTION_ENTRY = _KERNEL_DIR / "attention" / "entry.cpp" +_ATTENTION_ONLY_ENTRY = _KERNEL_DIR / "attention_only" / "entry.cpp" _ROPE_ONLY_ENTRY = _KERNEL_DIR / "rope_only" / "entry.cpp" _TILING_ENTRY = _KERNEL_DIR / "tiling" / "entry.cpp" @@ -65,10 +66,20 @@ def _cann_include_dirs() -> tuple[Path, ...]: METADATA_PREFIX_BYTES = KV_LENGTHS_OFFSET + BATCH * 8 BARRIER_SLOT_BYTES = 512 BARRIER_PHYSICAL_LANES = DEFAULT_BLOCK_DIM * 2 +ROPE_READY_SLOT_BYTES = 32 +ROPE_READY_PHYSICAL_LANES = DEFAULT_BLOCK_DIM * 2 # The CCE wrapper aligns the barrier start at runtime, so reserve one slot of -# alignment slack before the maximum 48 single-writer barrier slots. +# alignment slack before the maximum 48 single-writer barrier slots. The +# producer-only RoPE-ready soft-barrier region follows the aligned FAI barrier. METADATA_BYTES = ( - (METADATA_PREFIX_BYTES + BARRIER_SLOT_BYTES - 1 + BARRIER_PHYSICAL_LANES * BARRIER_SLOT_BYTES + 31) + ( + METADATA_PREFIX_BYTES + + BARRIER_SLOT_BYTES + - 1 + + BARRIER_PHYSICAL_LANES * BARRIER_SLOT_BYTES + + ROPE_READY_PHYSICAL_LANES * ROPE_READY_SLOT_BYTES + + 31 + ) // 32 * 32 ) @@ -86,6 +97,37 @@ def _cann_include_dirs() -> tuple[Path, ...]: dual_aiv_dispatch=True, ) def paged_attention_cce( + q_proj: pl.Tensor, + k_proj: pl.Tensor, + v_proj: pl.Tensor, + inv_rms_states: pl.Tensor, + q_norm_weight: pl.Tensor, + k_norm_weight: pl.Tensor, + seq_lens: pl.Tensor, + slot_mapping: pl.Tensor, + rope_cos: pl.Tensor, + rope_sin: pl.Tensor, + block_table: pl.Tensor, + # The extern return is the first writable tensor, so the real attention output + # must precede the query/cache state written by the fused rope prologue. + out: pl.Out[pl.Tensor], + query: pl.InOut[pl.Tensor], + key_cache: pl.InOut[pl.Tensor], + value_cache: pl.InOut[pl.Tensor], + workspace: pl.InOut[pl.Tensor], + metadata: pl.InOut[pl.Tensor], + cache_row_offset: pl.Scalar[pl.INDEX], +) -> pl.Tensor: ... + + +@pl.jit.extern( + core_type="mixed", + aic_source=_ATTENTION_ONLY_ENTRY, + aiv_source=_ATTENTION_ONLY_ENTRY, + include_dirs=_CANN_INCLUDE_DIRS, + dual_aiv_dispatch=True, +) +def paged_attention_only_cce( query: pl.Tensor, key_cache: pl.Tensor, value_cache: pl.Tensor, @@ -159,7 +201,7 @@ def qwen_decode_attention_cce( sync_start=True, deps=[tiling_tid], ) as _attention_tid: - out = paged_attention_cce( + out = paged_attention_only_cce( query, key_cache, value_cache, @@ -204,7 +246,7 @@ def qwen_decode_attention_cache_offset_test( sync_start=True, deps=[tiling_tid], ) as _attention_tid: - out = paged_attention_cce( + out = paged_attention_only_cce( query, key_cache, value_cache, @@ -306,6 +348,7 @@ def qwen_decode_rope_only_cce( "WORKSPACE_BYTES", "build_paged_attention_metadata", "paged_attention_cce", + "paged_attention_only_cce", "paged_attention_rope_only_cce", "paged_attention_tiling_cce", "qwen_decode_attention_cache_offset_test", From f2c48ae9970ad1cc1c6ebeeb7a3e7857269dfdb3 Mon Sep 17 00:00:00 2001 From: Youhezhen Date: Fri, 17 Jul 2026 03:22:22 -0700 Subject: [PATCH 3/3] fix(qwen3): args-based indexing for fused rope_qkv cross-core sync The Stage2 sync_qwen_rope_producers (pto::SYNCALL) deadlocked on a2a3: pto-isa's no-arg topology builtins (get_subblockdim/get_subblockid) read 0 in this pypto mixed+dual_aiv launch -- the dispatch sets no subblockdim register -- so SYNCALL_GET_MIX_PARTICIPANT_COUNT under-counted (24 not 72) and collided the AIV participant indices (507018 running-stalled, stuck fa_fused). The vendor FAI kernel avoids this by using only the args-based getters (get_block_idx(args)/get_sub_block_id(args)/get_block_num(args)). Replace the MIX syncall with an args-based handoff that matches the vendor's topology view: - AIV lanes soft-barrier via pto::SYNCALL_SOFT_AIV_BARRIER with idx get_block_idx(args)*2 + get_sub_block_id(args) (48 lanes); - each AIC polls the 48 rope-ready GM slots (DCCI+dsb) until all non-zero. Reverts the rope-ready region to 48 AIV slots (METADATA_BYTES 29376). Fused decode_fwd (rope_qkv inside the mixed-core attention extern) now PASSES on a2a3: exit=0, swimlane complete, Total Test Time ~3631us. --- .../paged_attention_cce/attention/entry.cpp | 52 ++++++++++++------- .../kernel/metadata_layout.h | 12 +++-- models/qwen3/14b/paged_attention_cce.py | 5 +- 3 files changed, 43 insertions(+), 26 deletions(-) diff --git a/models/qwen3/14b/kernels/paged_attention_cce/attention/entry.cpp b/models/qwen3/14b/kernels/paged_attention_cce/attention/entry.cpp index d9c21c7b..4b43b48a 100644 --- a/models/qwen3/14b/kernels/paged_attention_cce/attention/entry.cpp +++ b/models/qwen3/14b/kernels/paged_attention_cce/attention/entry.cpp @@ -38,34 +38,46 @@ static __aicore__ __attribute__((always_inline)) __gm__ int32_t *qwen_fai_barrie return reinterpret_cast<__gm__ int32_t *>(aligned_barrier); } -// AIV lanes finish the rope GM writes, barrier among themselves over the rope-ready -// metadata slots, then one sub-lane per AIC pair signals its AIC. The AIC waits for -// that flag before starting attention (which reads the rope-produced q/k/v cache). -// CrossCoreSetFlag uses PIPE_MTE3 so the signal fires only after rope's UB->GM -// stores have flushed — the cross-pipe GM visibility that a plain SYNCALL lacks. +// 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 ) { - pipe_barrier(PIPE_ALL); -#if defined(__DAV_CUBE__) - AscendC::CrossCoreWaitFlag(qwen_fai_metadata::kRopeReadyCrossCoreFlag); -#elif defined(__DAV_VEC__) __gm__ int32_t *ready = reinterpret_cast<__gm__ int32_t *>( reinterpret_cast<__gm__ uint8_t *>(fai_barrier) + qwen_fai_metadata::kBarrierBytes ); - uint32_t physical_lane = static_cast(get_block_idx(args)) * 2 + - static_cast(get_sub_block_id(args)); - __ubuf__ int32_t *ub_workspace = reinterpret_cast<__ubuf__ int32_t *>(0x3000); - pto::SYNCALL_SOFT_AIV_BARRIER( - ready, - ub_workspace, - qwen_fai_metadata::kRopeReadySlotCount, - physical_lane - ); - if (get_sub_block_id(args) == 0) { - AscendC::CrossCoreSetFlag<0x2, PIPE_MTE3>(qwen_fai_metadata::kRopeReadyCrossCoreFlag); + 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(get_block_idx(args)) * 2 + + static_cast(get_sub_block_id(args)); + pto::SYNCALL_SOFT_AIV_BARRIER(ready, reinterpret_cast<__ubuf__ int32_t *>(0x3000), + static_cast(qwen_fai_metadata::kRopeReadySlotCount), + static_cast(aiv_idx)); #endif pipe_barrier(PIPE_ALL); } diff --git a/models/qwen3/14b/kernels/paged_attention_cce/kernel/metadata_layout.h b/models/qwen3/14b/kernels/paged_attention_cce/kernel/metadata_layout.h index 6e0f0bf3..de507a9f 100644 --- a/models/qwen3/14b/kernels/paged_attention_cce/kernel/metadata_layout.h +++ b/models/qwen3/14b/kernels/paged_attention_cce/kernel/metadata_layout.h @@ -28,14 +28,16 @@ constexpr uint32_t kBarrierSlotBytes = 512; constexpr uint32_t kBarrierSlotWords = kBarrierSlotBytes / sizeof(int32_t); constexpr uint32_t kBarrierSlotCount = 48; constexpr uint32_t kBarrierBytes = kBarrierSlotCount * kBarrierSlotBytes; -// Producer-only RoPE-ready soft-barrier region, laid out right after the aligned -// FAI barrier. The fused kernel's AIV lanes barrier among themselves here, then one -// sub-lane per AIC pair signals the AIC attention via CrossCoreSetFlag. +// AIV rope-ready soft-barrier region, laid out right after the aligned FAI barrier. +// After the AIV prologue writes q/k/v, the 48 AIV lanes soft-barrier over these +// slots (args-based idx; the hardware get_subblockdim() is 0 in this launch, so +// pto::SYNCALL's builtin indexing cannot be used). Each AIC polls the slots until +// all 48 are non-zero. The soft barrier's DCCI+dsb gives the cross-core GM +// visibility the AIC attention needs. constexpr uint32_t kRopeReadySlotBytes = 32; constexpr uint32_t kRopeReadySlotWords = kRopeReadySlotBytes / sizeof(int32_t); -constexpr uint32_t kRopeReadySlotCount = 48; +constexpr uint32_t kRopeReadySlotCount = 48; // AIV lanes (DEFAULT_BLOCK_DIM * 2) constexpr uint32_t kRopeReadyBytes = kRopeReadySlotCount * kRopeReadySlotBytes; -constexpr uint16_t kRopeReadyCrossCoreFlag = 0; constexpr uint32_t kMetadataBytes = 29376; static_assert( diff --git a/models/qwen3/14b/paged_attention_cce.py b/models/qwen3/14b/paged_attention_cce.py index 2bfef69f..dc3954d8 100644 --- a/models/qwen3/14b/paged_attention_cce.py +++ b/models/qwen3/14b/paged_attention_cce.py @@ -67,10 +67,13 @@ def _cann_include_dirs() -> tuple[Path, ...]: BARRIER_SLOT_BYTES = 512 BARRIER_PHYSICAL_LANES = DEFAULT_BLOCK_DIM * 2 ROPE_READY_SLOT_BYTES = 32 +# AIV lanes that soft-barrier after the fused rope prologue (DEFAULT_BLOCK_DIM * 2 = 48). +# Args-based indexing (get_block_idx*2 + get_sub_block_id); the hardware get_subblockdim() +# is 0 in this launch, so pto::SYNCALL's builtin participant model cannot be used. ROPE_READY_PHYSICAL_LANES = DEFAULT_BLOCK_DIM * 2 # The CCE wrapper aligns the barrier start at runtime, so reserve one slot of # alignment slack before the maximum 48 single-writer barrier slots. The -# producer-only RoPE-ready soft-barrier region follows the aligned FAI barrier. +# AIV RoPE-ready soft-barrier region follows the aligned FAI barrier. METADATA_BYTES = ( ( METADATA_PREFIX_BYTES