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..4b43b48a 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,74 @@ 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); +} + +// 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(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); +} + +} // 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 4a0d97fa..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 }; @@ -158,4 +183,275 @@ run_qwen_fai(__gm__ int64_t *args, __gm__ int32_t *barrier_state = nullptr) { kernel(params, topology, barrier_state); } +// ══════════════════════════════════════════════════════════════════════════════ +// 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). +// ══════════════════════════════════════════════════════════════════════════════ +#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 +// 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__ + +// 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; +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; +// 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; +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_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..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,11 +28,21 @@ 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; +// 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; // AIV lanes (DEFAULT_BLOCK_DIM * 2) +constexpr uint32_t kRopeReadyBytes = kRopeReadySlotCount * kRopeReadySlotBytes; +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/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/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 96d6ea5a..dc3954d8 100644 --- a/models/qwen3/14b/paged_attention_cce.py +++ b/models/qwen3/14b/paged_attention_cce.py @@ -16,6 +16,8 @@ _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" @@ -64,10 +66,23 @@ 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 +# 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. +# alignment slack before the maximum 48 single-writer barrier slots. The +# AIV 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 ) @@ -85,6 +100,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, @@ -158,7 +204,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, @@ -203,7 +249,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, @@ -216,6 +262,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 +351,10 @@ def qwen_decode_attention_cache_offset_test( "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", "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()