From 508df463d9c842c1cf0503191a375da4ca0a427d Mon Sep 17 00:00:00 2001 From: shijiashuai Date: Mon, 14 Sep 2026 22:30:11 +0800 Subject: [PATCH 1/3] feat(cuda): add split-KV decode attention (TLLM-ATTN-SPLITKV PR-B) The decode attention kernel parallelises only over q-heads: at visible=2048 it launches grid=(14,1,1) x 128 threads and achieves 8.33% occupancy with nothing saturated (DRAM 0.44%, L2 0.96%, L1 4.0%). Split the visible KV window across blocks instead: grid=(Hq, num_splits), each block producing a local (m, l, acc), merged by a small combine kernel at grid=(Hq, 1). Graph capture is preserved: split ranges are derived in-kernel from the device visible length, num_splits is a host parameter so the grid is fixed per launch, and the partial buffer is caller-allocated -- no allocation or D2H during capture. Ranges are logical token ranges, so ContiguousRows/PagedRows, the frozen address formula and the zero-row semantics are untouched. The reduction loop is not copied: decode_online_softmax becomes decode_online_softmax_range. The two existing entry points call it with begin=0/end=visible/kFinalize=true and are otherwise unchanged, so all ~18 existing call sites and their tests keep working. Gates (tests/test_attention_splitkv.cpp, 11 cases): - num_splits == 1 is bitwise identical to the single-pass path, contiguous and paged. exp(0) is forced to exactly 1.0f for the max split rather than trusting the fast intrinsic, which makes that exactness structural, not incidental; - num_splits in {2,4,8,16} stays inside the independent oracle's 2e-3 tolerance; - split contiguous and split paged agree bitwise (same reduction, two addressing strategies); - visible=0, visible --- kernels/attention.cu | 286 ++++++++++++++- kernels/attention.cuh | 40 +++ tests/test_attention_splitkv.cpp | 582 +++++++++++++++++++++++++++++++ 3 files changed, 893 insertions(+), 15 deletions(-) create mode 100644 tests/test_attention_splitkv.cpp diff --git a/kernels/attention.cu b/kernels/attention.cu index 7815d36..ee1b241 100644 --- a/kernels/attention.cu +++ b/kernels/attention.cu @@ -126,16 +126,26 @@ struct ContiguousRows { // 共享的 decode online-softmax 循环。累加顺序、__expf 与 final normalize 与抽取前 // 逐字相同。 -template -__device__ __forceinline__ void decode_online_softmax(const half *q_smem, half *output, - float *scores, float *red, float *out_acc, - Rows rows, int visible_len, int head_dim, - float scale, int tid, int nthreads) { +// +// TLLM-ATTN-SPLITKV：新增 [begin, end) 段范围与编译期 kFinalize 开关。 +// kFinalize = true → 归一化后写 output。begin = 0 / end = visible_len 时与抽取前**逐字相同**； +// kFinalize = false → 不归一化，写 split-KV 的局部三元组 (m, l, acc) 到 partial_*。 +// 段范围落在**逻辑 token** 上，因此 ContiguousRows / PagedRows 的取址（含冻结的地址公式与 +// 非法块 id 的零行语义）完全不需要改动。 +template +__device__ __forceinline__ void +decode_online_softmax_range(const half *q_smem, half *output, float *scores, float *red, + float *out_acc, Rows rows, int begin, int end, int head_dim, + float scale, int tid, int nthreads, float *partial_m, float *partial_l, + float *partial_acc) { float running_max = -FLT_MAX; float running_sum = 0.0f; - for (int tile_start = 0; tile_start < visible_len; tile_start += ATTEND_TILE) { - const int tile_size = min(ATTEND_TILE, visible_len - tile_start); + // 空段（begin >= end，例如 visible < num_splits 时的高位 split）不进入循环， + // running_max 保持 -FLT_MAX、running_sum 与 out_acc 保持 0 —— 这正是 combine 需要的 + // 中性三元组。combine 侧 m == M 的特例见下，全空时不会出现 NaN。 + for (int tile_start = begin; tile_start < end; tile_start += ATTEND_TILE) { + const int tile_size = min(ATTEND_TILE, end - tile_start); rows.begin_tile(tile_start, tile_size, tid, nthreads); if (Rows::kNeedsTilePrep) __syncthreads(); @@ -184,10 +194,21 @@ __device__ __forceinline__ void decode_online_softmax(const half *q_smem, half * __syncthreads(); } - // Final normalize. - const float inv_sum = 1.0f / (running_sum + 1e-9f); - for (int d = tid; d < head_dim; d += nthreads) { - output[d] = __float2half(out_acc[d] * inv_sum); + if constexpr (kFinalize) { + // Final normalize. + const float inv_sum = 1.0f / (running_sum + 1e-9f); + for (int d = tid; d < head_dim; d += nthreads) { + output[d] = __float2half(out_acc[d] * inv_sum); + } + } else { + // split-KV：写局部三元组。m 与 l 是块内归约结果（对所有线程一致），换一个线程写。 + if (tid == 0) { + partial_m[0] = running_max; + partial_l[0] = running_sum; + } + for (int d = tid; d < head_dim; d += nthreads) { + partial_acc[d] = out_acc[d]; + } } } @@ -231,8 +252,9 @@ __global__ void attention_decode_kernel(const half *__restrict__ query, __syncthreads(); ContiguousRows rows{k, v, kv_stride}; - decode_online_softmax(q_smem, o, scores, red, out_acc, rows, visible_len, head_dim, scale, tid, - nthreads); + decode_online_softmax_range(q_smem, o, scores, red, out_acc, rows, 0, + visible_len, head_dim, scale, tid, nthreads, + nullptr, nullptr, nullptr); } void attention_decode(const half *query, const half *k_cache, const half *v_cache, half *output, @@ -373,8 +395,9 @@ __global__ void attention_decode_paged_kernel( rows.row_elem = row_elem; rows.tile_start = 0; - decode_online_softmax(q_smem, o, scores, red, out_acc, rows, visible_tokens, head_dim, scale, - tid, nthreads); + decode_online_softmax_range(q_smem, o, scores, red, out_acc, rows, 0, + visible_tokens, head_dim, scale, tid, nthreads, + nullptr, nullptr, nullptr); } void attention_decode_paged(const half *query, const half *k_pool_layer, const half *v_pool_layer, @@ -404,6 +427,239 @@ void attention_decode_paged(const half *query, const half *k_pool_layer, const h head_dim, device_visible_tokens, block_size, max_num_blocks, table_len); } +// ============================================================================ +// Split-KV decode attention（TLLM-ATTN-SPLITKV） +// +// 动机（见 docs/architecture/decode-attention-splitkv-design.md §2）：单 query decode 的 +// 并行度只有 num_q_heads 一个轴，实测 achieved occupancy 8.33%（S=2048），且无任何资源 +// 饱和（DRAM 0.44%、SM 0.72%、L2 0.96%）—— 瓶颈是"没有足够的 warp 隐藏访存延迟"。 +// 因此把**可见 KV 逻辑窗口**切成 num_splits 段，每段一个 block 独立归约，再由一个小的 +// combine kernel 合并。 +// +// 关键性质： +// - 段范围在 device 端由 *device_visible_* 派生，num_splits 是 host 参数 ⇒ grid 在捕获时 +// 固定、可见长度仍可随 replay 变化，**保持 CUDA Graph 可捕获**； +// - num_splits == 1 时 combine 退化为恒等，输出与单遍路径**逐位相同**（见 combine 注释）； +// - 归约循环不复制：两条路径继续共用 decode_online_softmax_range； +// - 段边界落在逻辑 token 上 ⇒ 地址公式与零行语义不需要改。 +// ============================================================================ + +// 把 num_splits 个 (m_i, l_i, acc_i) 按 online-softmax 的合并式折成一个输出： +// M = max_i m_i; L = Σ l_i·w_i; ACC[d] = Σ acc_i[d]·w_i; out = ACC / (L + 1e-9) +// w_i = exp(m_i - M) +// +// m_i == M 时把 w_i **显式取 1.0f**，而不是交给 __expf(0)：这样 num_splits == 1 时 +// L = l_0·1.0f = l_0、ACC = acc_0，与单遍路径的 1/(running_sum + 1e-9) 逐位一致，且不依赖 +// 快速 intrinsics 在 0 点的返回值。数学上 exp(0) 恒为 1，所以这也是更准的形式。 +// +// 全空（visible == 0）时所有 m_i = -FLT_MAX ⇒ M = -FLT_MAX、全部 w_i = 1、L = Σ l_i = 0、 +// ACC = 0 ⇒ out = 0·(1/(0+1e-9)) = 0：与 test_paged_direct.cpp 冻结的「visible = 0 → 全 0」 +// 语义一致，且不产生 NaN。 +__global__ void attention_splitkv_combine_kernel(const float *__restrict__ partial, + half *__restrict__ output, int head_dim, + int num_splits) { + const int q_head = blockIdx.x; + const int tid = threadIdx.x; + const int nthreads = blockDim.x; + const int stride = 2 + head_dim; + const float *base = partial + static_cast(q_head) * num_splits * stride; + + float M = -FLT_MAX; + for (int s = 0; s < num_splits; ++s) { + M = fmaxf(M, base[static_cast(s) * stride]); + } + + float L = 0.0f; + for (int s = 0; s < num_splits; ++s) { + const float m = base[static_cast(s) * stride]; + const float w = (m == M) ? 1.0f : __expf(m - M); + L += base[static_cast(s) * stride + 1] * w; + } + + const float inv_sum = 1.0f / (L + 1e-9f); + half *o = output + static_cast(q_head) * head_dim; + for (int d = tid; d < head_dim; d += nthreads) { + float acc = 0.0f; + for (int s = 0; s < num_splits; ++s) { + const float m = base[static_cast(s) * stride]; + // L 的循环里已经判断过一次；这里为每个 d 重算，combine 的规模（Hq 个 block， + // num_splits 个槽）远小于 attention 本体，不值得引入共享内存广播的复杂度。 + const float w = (m == M) ? 1.0f : __expf(m - M); + acc += base[static_cast(s) * stride + 2 + d] * w; + } + o[d] = __float2half(acc * inv_sum); + } +} + +// 连续 KV 的分段 partial：一个 block 负责 (q_head, split) 的一个逻辑区间。 +__global__ void +attention_decode_splitkv_kernel(const half *__restrict__ query, const half *__restrict__ k_cache, + const half *__restrict__ v_cache, half *__restrict__ output, + float scale, int num_q_heads, int num_kv_heads, + const int *__restrict__ device_visible_len, int head_dim, + float *__restrict__ partial, int num_splits) { + const int visible_len = *device_visible_len; + const int q_head = blockIdx.x; + const int split = blockIdx.y; + const int tid = threadIdx.x; + const int nthreads = blockDim.x; + + const int group_size = num_q_heads / num_kv_heads; + const int kv_head = q_head / group_size; + const int kv_stride = num_kv_heads * head_dim; + + const half *k = k_cache + kv_head * head_dim; + const half *v = v_cache + kv_head * head_dim; + half *o = output + q_head * head_dim; + + extern __shared__ float smem[]; + AttentionSmemLayout layout{head_dim}; + float *scores = smem + layout.scores_offset(); + float *red = smem + layout.red_offset(); + float *out_acc = smem + layout.out_acc_offset(); + half *q_smem = + reinterpret_cast(reinterpret_cast(smem) + layout.q_offset_bytes()); + + for (int d = tid; d < head_dim; d += nthreads) { + q_smem[d] = query[q_head * head_dim + d]; + out_acc[d] = 0.0f; + } + __syncthreads(); + + // 段范围：chunk 由**运行时**可见长度派生，故 grid 固定也能随 replay 变化。 + const int chunk = (visible_len + num_splits - 1) / num_splits; + const int begin = split * chunk; + const int end = (begin + chunk < visible_len) ? (begin + chunk) : visible_len; + + const int stride = 2 + head_dim; + float *pb = + partial + (static_cast(q_head) * num_splits + split) * static_cast(stride); + + ContiguousRows rows{k, v, kv_stride}; + decode_online_softmax_range(q_smem, o, scores, red, out_acc, rows, begin, + end, head_dim, scale, tid, nthreads, pb, + pb + 1, pb + 2); +} + +// 分页 KV 的分段 partial：与上面同构，只换取址策略（池 + 块表）。 +__global__ void attention_decode_paged_splitkv_kernel( + const half *__restrict__ query, const half *__restrict__ k_pool_layer, + const half *__restrict__ v_pool_layer, const int *__restrict__ block_table, + half *__restrict__ output, float scale, int num_q_heads, int num_kv_heads, int head_dim, + const int *__restrict__ device_visible_tokens, int block_size, int max_num_blocks, + int table_len, float *__restrict__ partial, int num_splits) { + const int visible_tokens = *device_visible_tokens; + const int q_head = blockIdx.x; + const int split = blockIdx.y; + const int tid = threadIdx.x; + const int nthreads = blockDim.x; + + const int group_size = num_q_heads / num_kv_heads; + const int kv_head = q_head / group_size; + const int kv_dim = num_kv_heads * head_dim; + + const half *k_base = k_pool_layer + kv_head * head_dim; + const half *v_base = v_pool_layer + kv_head * head_dim; + half *o = output + q_head * head_dim; + + extern __shared__ float smem[]; + AttentionSmemLayout layout{head_dim}; + float *scores = smem + layout.scores_offset(); + float *red = smem + layout.red_offset(); + float *out_acc = smem + layout.out_acc_offset(); + half *q_smem = + reinterpret_cast(reinterpret_cast(smem) + layout.q_offset_bytes()); + + char *tail = reinterpret_cast(smem) + layout.total_bytes(); + int *row_elem = reinterpret_cast(tail); + half *zero_row = reinterpret_cast(tail + ATTEND_TILE * sizeof(int)); + + for (int d = tid; d < head_dim; d += nthreads) { + q_smem[d] = query[q_head * head_dim + d]; + out_acc[d] = 0.0f; + zero_row[d] = __float2half(0.0f); + } + __syncthreads(); + + PagedRows rows; + rows.k_base = k_base; + rows.v_base = v_base; + rows.zero_row = zero_row; + rows.block_table = block_table; + rows.block_size = block_size; + rows.max_num_blocks = max_num_blocks; + rows.table_len = table_len; + rows.kv_dim = kv_dim; + rows.row_elem = row_elem; + rows.tile_start = 0; + + const int chunk = (visible_tokens + num_splits - 1) / num_splits; + const int begin = split * chunk; + const int end = (begin + chunk < visible_tokens) ? (begin + chunk) : visible_tokens; + + const int stride = 2 + head_dim; + float *pb = + partial + (static_cast(q_head) * num_splits + split) * static_cast(stride); + + decode_online_softmax_range(q_smem, o, scores, red, out_acc, rows, begin, end, + head_dim, scale, tid, nthreads, pb, pb + 1, + pb + 2); +} + +void attention_decode_splitkv(const half *query, const half *k_cache, const half *v_cache, + half *output, float scale, int num_q_heads, int num_kv_heads, + const int *device_visible_len, int head_dim, float *partial_workspace, + int num_splits, cudaStream_t stream) { + // 防御性检查与注意：num_splits 会作为 grid.y，受 65535 上界约束。 + if (num_q_heads <= 0 || num_kv_heads <= 0 || head_dim <= 0 || device_visible_len == nullptr) { + return; + } + if (partial_workspace == nullptr || num_splits < 1 || num_splits > 65535) { + return; + } + + AttentionSmemLayout layout{head_dim}; + const size_t shared_size = layout.total_bytes(); + dim3 grid(static_cast(num_q_heads), static_cast(num_splits)); + + attention_decode_splitkv_kernel<<>>( + query, k_cache, v_cache, output, scale, num_q_heads, num_kv_heads, device_visible_len, + head_dim, partial_workspace, num_splits); + attention_splitkv_combine_kernel<<>>(partial_workspace, output, + head_dim, num_splits); +} + +void attention_decode_paged_splitkv(const half *query, const half *k_pool_layer, + const half *v_pool_layer, const int *block_table, half *output, + float scale, int num_q_heads, int num_kv_heads, int head_dim, + const int *device_visible_tokens, int block_size, + int max_num_blocks, int table_len, float *partial_workspace, + int num_splits, cudaStream_t stream) { + if (query == nullptr || k_pool_layer == nullptr || v_pool_layer == nullptr || + block_table == nullptr || output == nullptr || device_visible_tokens == nullptr) { + return; + } + if (num_q_heads <= 0 || num_kv_heads <= 0 || head_dim <= 0 || block_size <= 0 || + max_num_blocks <= 0) { + return; + } + if (partial_workspace == nullptr || num_splits < 1 || num_splits > 65535) { + return; + } + + AttentionSmemLayout layout{head_dim}; + const size_t shared_size = layout.total_bytes() + ATTEND_TILE * sizeof(int) + + static_cast(head_dim) * sizeof(half); + dim3 grid(static_cast(num_q_heads), static_cast(num_splits)); + + attention_decode_paged_splitkv_kernel<<>>( + query, k_pool_layer, v_pool_layer, block_table, output, scale, num_q_heads, num_kv_heads, + head_dim, device_visible_tokens, block_size, max_num_blocks, table_len, partial_workspace, + num_splits); + attention_splitkv_combine_kernel<<>>(partial_workspace, output, + head_dim, num_splits); +} + // Prefill attention: full sequence with causal masking. // Q: [S, Hq, D]; K/V: [S, Hkv, D]. One block per (query_pos, q_head). __global__ void attention_prefill_kernel(const half *__restrict__ query, diff --git a/kernels/attention.cuh b/kernels/attention.cuh index 6af3f15..d512228 100644 --- a/kernels/attention.cuh +++ b/kernels/attention.cuh @@ -76,6 +76,46 @@ void attention_decode_paged(const half *__restrict__ query, const half *__restri const int *device_visible_tokens, int block_size, int max_num_blocks, int table_len, cudaStream_t stream = 0); +// ============================================================================ +// Split-KV decode attention（TLLM-ATTN-SPLITKV） +// +// 把可见 KV **逻辑窗口**切成 num_splits 段，每段一个 block 独立归约（grid = (Hq, num_splits)）， +// 再由一个 combine kernel（grid = (Hq, 1)）按 online-softmax 合并式折成输出。动机是并行度： +// 单 query decode 原本只有 Hq 这一个并行轴，实测 occupancy 8.33%、无任何资源饱和 +// （见 docs/architecture/decode-attention-splitkv-design.md §2）。 +// +// 与 attention_decode / attention_decode_paged 的关系： +// - **语义等价**：两者是同一次归约的两种切分方式，输出差异仅来自 fp32 求和顺序； +// - **num_splits == 1 时逐位相同**（见 kernels/attention.cu 的 combine 注释），这条是 +// 重构正确性的回归锚点； +// - **CUDA Graph 可捕获**：段范围在 device 端由 *device_visible_* 派生，num_splits 是 host +// 参数（grid 捕获时固定），partial 缓冲由调用方预分配、捕获期零分配零 D2H； +// - 两个既有入口**保持不变**（等价于单遍路径），因此既有调用点无需改动。 +// +// partial_workspace 布局（device fp32，连续，无 padding）： +// offset(head, split) = (head * num_splits + split) * (2 + head_dim) +// [+0] = m 该段最大 score（空段为 -FLT_MAX，combine 视作中性） +// [+1] = l 该段 Σ exp +// [+2 .. +1+D] = acc[d] 该段 Σ exp * v +// total floats = num_q_heads * num_splits * (2 + head_dim) +// 调用方负责分配（kernel 零分配、零状态）。每个 block 都会写自己的槽位，含空段。 +// +// 前置条件：指针非空；num_q_heads > 0、num_kv_heads > 0、num_q_heads % num_kv_heads == 0； +// head_dim > 0；1 <= num_splits <= 65535（grid.y 上界）；num_splits == 1 为恒等切分。 +void attention_decode_splitkv(const half *__restrict__ query, const half *__restrict__ k_cache, + const half *__restrict__ v_cache, half *__restrict__ output, + float scale, int num_q_heads, int num_kv_heads, + const int *__restrict__ device_visible_len, int head_dim, + float *__restrict__ partial_workspace, int num_splits, + cudaStream_t stream = 0); + +void attention_decode_paged_splitkv( + const half *__restrict__ query, const half *__restrict__ k_pool_layer, + const half *__restrict__ v_pool_layer, const int *__restrict__ block_table, + half *__restrict__ output, float scale, int num_q_heads, int num_kv_heads, int head_dim, + const int *__restrict__ device_visible_tokens, int block_size, int max_num_blocks, + int table_len, float *__restrict__ partial_workspace, int num_splits, cudaStream_t stream = 0); + // Prefill attention: full sequence with causal masking // Q: [S, Hq, D] // K: [S, Hkv, D] diff --git a/tests/test_attention_splitkv.cpp b/tests/test_attention_splitkv.cpp new file mode 100644 index 0000000..d9a9f78 --- /dev/null +++ b/tests/test_attention_splitkv.cpp @@ -0,0 +1,582 @@ +// TLLM-ATTN-SPLITKV：split-KV decode attention 的 kernel 级门禁。 +// +// 设计包 §8 的矩阵，按强度排序： +// 1. **逐位锚点**：num_splits == 1 时，split 路径与单遍路径输出逐位相同 +// （连续与分页各一次）。这是把"重构风险"与"数值改动"分开的那条门禁； +// 2. **独立 oracle**：num_splits > 1 时落在 TLLM-P0-002 oracle 的容差内 +// （oracle 不调用任何生产 kernel，能发现"两条生产路径一起错"）； +// 3. **同归约、两种取址**：split 连续 vs split 分页逐位相同； +// 4. 边界：visible = 0、visible < num_splits（空段）、非整数倍、非法块 id、短块表； +// 5. **CUDA Graph**：捕获后改变 device 端可见长度并 replay，与非 graph 结果逐位相同 +// ——直接证明"grid 固定 + 段范围走 device int"这一设计前提成立。 +// +// 边界：只测 kernel 级接口；LayerWorkspace 接线与开关属于 PR-C，不在本文件范围。 +// 本文件不产生任何性能数字。 + +#include "attention.cuh" +#include "paged_attention_oracle.h" +#include "paged_kv.cuh" +#include "tiny_llm/cuda_utils.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace tiny_llm; +using namespace tiny_llm::kernels; +using tiny_llm::test::maxAbsDiff; +using tiny_llm::test::oracleAttentionDecode; +using tiny_llm::test::PagedGeometry; +using tiny_llm::test::requiredBlocks; + +namespace { + +constexpr float kOracleTolerance = 2e-3f; +// partial 工作区按最大 split 数分配；测试里只用到 ≤ 16。 +constexpr int kMaxSplits = 32; + +bool hasCudaDevice() { + static bool checked = false; + static bool has_device = false; + if (!checked) { + int n = 0; + cudaError_t e = cudaGetDeviceCount(&n); + has_device = (e == cudaSuccess && n > 0); + checked = true; + } + return has_device; +} + +std::vector randomFp16(size_t n, unsigned seed, float scale = 1.0f) { + std::mt19937 gen(seed); + std::uniform_real_distribution dist(-scale, scale); + std::vector v(n); + for (auto &x : v) + x = __float2half(dist(gen)); + return v; +} + +std::vector toFloat(const std::vector &h) { + std::vector f(h.size()); + for (size_t i = 0; i < h.size(); ++i) + f[i] = __half2float(h[i]); + return f; +} + +bool bitwiseEqual(const std::vector &a, const std::vector &b) { + return a.size() == b.size() && std::memcmp(a.data(), b.data(), a.size() * sizeof(half)) == 0; +} + +PagedGeometry makeGeometry(int hq, int hkv, int hd, int bs, int max_blocks, int layers = 1) { + PagedGeometry g; + g.num_q_heads = hq; + g.num_kv_heads = hkv; + g.head_dim = hd; + g.block_size = bs; + g.max_num_blocks = max_blocks; + g.num_layers = layers; + return g; +} + +std::vector makeTable(int n, int max_blocks, unsigned seed) { + std::vector ids(static_cast(max_blocks)); + for (int i = 0; i < max_blocks; ++i) + ids[static_cast(i)] = i; + std::mt19937 gen(seed); + std::shuffle(ids.begin(), ids.end(), gen); + return std::vector(ids.begin(), ids.begin() + n); +} + +// 一个几何 + pool + 块表 + scratch + partial 工作区的固定装置。 +// 三条路径消费同一份物理 pool：单遍（gather + attention_decode）、split 连续、split 分页。 +class SplitKvFixture { + public: + SplitKvFixture(const PagedGeometry &g, const std::vector &table, int table_len) + : g_(g), table_len_(table_len), d_k_pool_(g.poolElements()), d_v_pool_(g.poolElements()), + d_k_scratch_(static_cast(g.max_num_blocks) * static_cast(g.block_size) * + static_cast(g.kvDim())), + d_v_scratch_(static_cast(g.max_num_blocks) * static_cast(g.block_size) * + static_cast(g.kvDim())), + d_partial_(static_cast(g.num_q_heads) * kMaxSplits * + (2 + static_cast(g.head_dim))), + d_table_(static_cast(g.max_num_blocks)), d_len_(1) { + d_table_.copyFromHost(table.data(), table.size()); + } + + void scatter(int layer, const std::vector &k, const std::vector &v, int tokens, + int position) { + const size_t elems = static_cast(tokens) * g_.kvDim(); + DeviceBuffer d_k(elems), d_v(elems); + d_k.copyFromHost(k.data(), elems); + d_v.copyFromHost(v.data(), elems); + half *kp = d_k_pool_.data() + static_cast(layer) * g_.layerStride(); + half *vp = d_v_pool_.data() + static_cast(layer) * g_.layerStride(); + paged_scatter_blocks(d_k.data(), kp, d_table_.data(), tokens, position, g_.block_size, + g_.kvDim(), g_.max_num_blocks); + paged_scatter_blocks(d_v.data(), vp, d_table_.data(), tokens, position, g_.block_size, + g_.kvDim(), g_.max_num_blocks); + } + + // 单遍连续路径：gather 到 scratch 后调用既有 attention_decode（今天的行为）。 + std::vector runSinglePassContiguous(int layer, const std::vector &q, int visible, + cudaStream_t stream = 0) { + const size_t q_elems = static_cast(g_.num_q_heads) * g_.head_dim; + DeviceBuffer d_q(q_elems), d_out(q_elems); + d_q.copyFromHost(q.data(), q_elems); + d_len_.copyFromHost(&visible, 1, stream); + + const half *kp = d_k_pool_.data() + static_cast(layer) * g_.layerStride(); + const half *vp = d_v_pool_.data() + static_cast(layer) * g_.layerStride(); + paged_gather_blocks(d_k_scratch_.data(), kp, d_table_.data(), visible, g_.block_size, + g_.kvDim(), g_.max_num_blocks, stream); + paged_gather_blocks(d_v_scratch_.data(), vp, d_table_.data(), visible, g_.block_size, + g_.kvDim(), g_.max_num_blocks, stream); + attention_decode(d_q.data(), d_k_scratch_.data(), d_v_scratch_.data(), d_out.data(), + scale(), g_.num_q_heads, g_.num_kv_heads, d_len_.data(), g_.head_dim, + stream); + return readBack(d_out, q_elems); + } + + // 单遍分页路径（PR-2 的 direct kernel，今天的行为）。 + std::vector runSinglePassPaged(int layer, const std::vector &q, int visible, + cudaStream_t stream = 0) { + const size_t q_elems = static_cast(g_.num_q_heads) * g_.head_dim; + DeviceBuffer d_q(q_elems), d_out(q_elems); + d_q.copyFromHost(q.data(), q_elems); + d_len_.copyFromHost(&visible, 1, stream); + + const half *kp = d_k_pool_.data() + static_cast(layer) * g_.layerStride(); + const half *vp = d_v_pool_.data() + static_cast(layer) * g_.layerStride(); + attention_decode_paged(d_q.data(), kp, vp, d_table_.data(), d_out.data(), scale(), + g_.num_q_heads, g_.num_kv_heads, g_.head_dim, d_len_.data(), + g_.block_size, g_.max_num_blocks, table_len_, stream); + return readBack(d_out, q_elems); + } + + // split 连续路径：先 gather（复用单遍的 scratch），再走 split-KV。 + std::vector runSplitContiguous(int layer, const std::vector &q, int visible, + int num_splits, cudaStream_t stream = 0) { + const size_t q_elems = static_cast(g_.num_q_heads) * g_.head_dim; + DeviceBuffer d_q(q_elems), d_out(q_elems); + d_q.copyFromHost(q.data(), q_elems); + d_len_.copyFromHost(&visible, 1, stream); + + const half *kp = d_k_pool_.data() + static_cast(layer) * g_.layerStride(); + const half *vp = d_v_pool_.data() + static_cast(layer) * g_.layerStride(); + paged_gather_blocks(d_k_scratch_.data(), kp, d_table_.data(), visible, g_.block_size, + g_.kvDim(), g_.max_num_blocks, stream); + paged_gather_blocks(d_v_scratch_.data(), vp, d_table_.data(), visible, g_.block_size, + g_.kvDim(), g_.max_num_blocks, stream); + attention_decode_splitkv(d_q.data(), d_k_scratch_.data(), d_v_scratch_.data(), d_out.data(), + scale(), g_.num_q_heads, g_.num_kv_heads, d_len_.data(), + g_.head_dim, d_partial_.data(), num_splits, stream); + return readBack(d_out, q_elems); + } + + // split 分页路径：直接从 pool 寻址。 + std::vector runSplitPaged(int layer, const std::vector &q, int visible, + int num_splits, cudaStream_t stream = 0) { + const size_t q_elems = static_cast(g_.num_q_heads) * g_.head_dim; + DeviceBuffer d_q(q_elems), d_out(q_elems); + d_q.copyFromHost(q.data(), q_elems); + d_len_.copyFromHost(&visible, 1, stream); + + const half *kp = d_k_pool_.data() + static_cast(layer) * g_.layerStride(); + const half *vp = d_v_pool_.data() + static_cast(layer) * g_.layerStride(); + attention_decode_paged_splitkv(d_q.data(), kp, vp, d_table_.data(), d_out.data(), scale(), + g_.num_q_heads, g_.num_kv_heads, g_.head_dim, d_len_.data(), + g_.block_size, g_.max_num_blocks, table_len_, + d_partial_.data(), num_splits, stream); + return readBack(d_out, q_elems); + } + + // 供 CUDA Graph 测试直接拿 device 指针。 + half *kPoolData() { return d_k_pool_.data(); } + half *vPoolData() { return d_v_pool_.data(); } + int *tableData() { return d_table_.data(); } + float *partialData() { return d_partial_.data(); } + int *lenData() { return d_len_.data(); } + int tableLen() const { return table_len_; } + float scale() const { return 1.0f / std::sqrt(static_cast(g_.head_dim)); } + size_t qElems() const { + return static_cast(g_.num_q_heads) * static_cast(g_.head_dim); + } + void setLen(int visible, cudaStream_t stream = 0) { d_len_.copyFromHost(&visible, 1, stream); } + void sync() { ASSERT_EQ(cudaDeviceSynchronize(), cudaSuccess); } + + private: + std::vector readBack(DeviceBuffer &d_out, size_t q_elems) { + std::vector out(q_elems); + d_out.copyToHost(out.data(), q_elems); + EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess); + return out; + } + + PagedGeometry g_; + int table_len_; + DeviceBuffer d_k_pool_; + DeviceBuffer d_v_pool_; + DeviceBuffer d_k_scratch_; + DeviceBuffer d_v_scratch_; + DeviceBuffer d_partial_; + DeviceBuffer d_table_; + DeviceBuffer d_len_; +}; + +// 构造 fixture 并写入逻辑 K/V（按块表 scatter 进 pool）。 +struct Built { + std::unique_ptr fx; + std::vector k_logical, v_logical, q; +}; + +Built buildCase(const PagedGeometry &g, int visible, unsigned seed) { + const int nb = requiredBlocks(visible, g.block_size); + const auto table = makeTable(nb, g.max_num_blocks, seed * 10 + 4); + const size_t k_elems = static_cast(visible) * g.kvDim(); + const size_t q_elems = static_cast(g.num_q_heads) * g.head_dim; + + Built b; + b.fx = std::make_unique(g, table, nb); + b.k_logical = randomFp16(k_elems, seed * 10 + 1); + b.v_logical = randomFp16(k_elems, seed * 10 + 2); + b.q = randomFp16(q_elems, seed * 10 + 3); + b.fx->scatter(0, b.k_logical, b.v_logical, visible, 0); + b.fx->sync(); + return b; +} + +} // namespace + +class SplitKvTest : public ::testing::Test { + protected: + void SetUp() override { + if (!hasCudaDevice()) GTEST_SKIP() << "No CUDA device available"; + cudaSetDevice(0); + } + void TearDown() override { cudaDeviceSynchronize(); } +}; + +// ── 1. 逐位锚点（连续）────────────────────────────────────────────────────── +TEST_F(SplitKvTest, NumSplitsOneIsBitwiseIdenticalToSinglePassContiguous) { + struct Case { + int bs, visible, hq, hkv, hd; + const char *label; + }; + const std::vector cases = { + {16, 1, 4, 4, 32, "S=1"}, + {16, 17, 8, 2, 64, "跨块 GQA"}, + {16, 20, 14, 2, 64, "生产形状"}, + {32, 96, 14, 2, 64, "三块 生产形状"}, + {32, 129, 4, 2, 64, "跨 ATTEND_TILE"}, + {16, 200, 4, 4, 128, "多 tile D=128"}, + {8, 5, 8, 1, 64, "MQA"}, + }; + + for (const auto &c : cases) { + for (unsigned seed : {1u, 99u}) { + SCOPED_TRACE(std::string(c.label) + " seed=" + std::to_string(seed)); + const int nb = requiredBlocks(c.visible, c.bs); + PagedGeometry g = makeGeometry(c.hq, c.hkv, c.hd, c.bs, nb + 3); + ASSERT_TRUE(tiny_llm::test::geometryIsValid(g)); + auto b = buildCase(g, c.visible, seed); + + const auto single = b.fx->runSinglePassContiguous(0, b.q, c.visible); + const auto split1 = b.fx->runSplitContiguous(0, b.q, c.visible, /*num_splits=*/1); + + EXPECT_TRUE(bitwiseEqual(single, split1)) + << "num_splits=1 必须与单遍路径逐位相同（重构锚点），max|diff|=" + << maxAbsDiff(toFloat(single), toFloat(split1)); + } + } +} + +// ── 1b. 逐位锚点（分页）──────────────────────────────────────────────────── +TEST_F(SplitKvTest, NumSplitsOneIsBitwiseIdenticalToSinglePassPaged) { + struct Case { + int bs, visible, hq, hkv, hd; + const char *label; + }; + const std::vector cases = { + {1, 1, 4, 4, 32, "block=1"}, {1, 17, 4, 2, 32, "block=1 跨块"}, + {16, 16, 4, 2, 32, "恰好一块"}, {16, 35, 8, 1, 64, "2*block+tail MQA"}, + {32, 96, 14, 2, 64, "生产形状"}, {32, 129, 4, 2, 64, "跨 ATTEND_TILE"}, + }; + + for (const auto &c : cases) { + for (unsigned seed : {7u, 4321u}) { + SCOPED_TRACE(std::string(c.label) + " seed=" + std::to_string(seed)); + const int nb = requiredBlocks(c.visible, c.bs); + PagedGeometry g = makeGeometry(c.hq, c.hkv, c.hd, c.bs, nb + 3); + ASSERT_TRUE(tiny_llm::test::geometryIsValid(g)); + auto b = buildCase(g, c.visible, seed); + + const auto single = b.fx->runSinglePassPaged(0, b.q, c.visible); + const auto split1 = b.fx->runSplitPaged(0, b.q, c.visible, /*num_splits=*/1); + + EXPECT_TRUE(bitwiseEqual(single, split1)) + << "num_splits=1 必须与单遍 paged 路径逐位相同，max|diff|=" + << maxAbsDiff(toFloat(single), toFloat(split1)); + } + } +} + +// ── 2. num_splits > 1：对独立 oracle 的容差 + 不产生 NaN ──────────────────── +TEST_F(SplitKvTest, SplitMatchesIndependentOracle) { + struct Case { + int bs, visible, hq, hkv, hd; + const char *label; + }; + const std::vector cases = { + {16, 20, 14, 2, 64, "生产形状"}, {32, 96, 14, 2, 64, "三块 生产形状"}, + {16, 33, 8, 2, 64, "block+1 GQA"}, {16, 128, 4, 4, 32, "整 tile 边界"}, + {16, 129, 4, 2, 64, "跨 tile"}, {32, 200, 4, 4, 128, "多 tile D=128"}, + {8, 40, 8, 1, 64, "MQA"}, + }; + + for (const auto &c : cases) { + for (int num_splits : {2, 4, 8, 16}) { + SCOPED_TRACE(std::string(c.label) + " num_splits=" + std::to_string(num_splits)); + const int nb = requiredBlocks(c.visible, c.bs); + PagedGeometry g = makeGeometry(c.hq, c.hkv, c.hd, c.bs, nb + 3); + ASSERT_TRUE(tiny_llm::test::geometryIsValid(g)); + auto b = buildCase(g, c.visible, 1234u); + + const auto out = b.fx->runSplitPaged(0, b.q, c.visible, num_splits); + const auto got = toFloat(out); + for (float x : got) + EXPECT_TRUE(std::isfinite(x)) << "split 输出出现非有限值"; + + const auto ref = oracleAttentionDecode(toFloat(b.q), toFloat(b.k_logical), + toFloat(b.v_logical), c.visible, g); + EXPECT_LT(maxAbsDiff(got, ref), kOracleTolerance) + << "split 与独立 oracle 的偏差超出容差，max|diff|=" << maxAbsDiff(got, ref); + } + } +} + +// ── 3. split 连续 vs split 分页：同一份 pool 上逐位相同 ───────────────────── +TEST_F(SplitKvTest, SplitContiguousAndPagedAgreeBitwise) { + struct Case { + int bs, visible, hq, hkv, hd; + const char *label; + }; + const std::vector cases = { + {16, 20, 14, 2, 64, "生产形状"}, + {1, 17, 4, 2, 32, "block=1 跨块"}, + {16, 129, 4, 2, 64, "跨 tile"}, + {32, 96, 8, 1, 64, "MQA"}, + }; + + for (const auto &c : cases) { + for (int num_splits : {1, 2, 8}) { + SCOPED_TRACE(std::string(c.label) + " num_splits=" + std::to_string(num_splits)); + const int nb = requiredBlocks(c.visible, c.bs); + PagedGeometry g = makeGeometry(c.hq, c.hkv, c.hd, c.bs, nb + 3); + ASSERT_TRUE(tiny_llm::test::geometryIsValid(g)); + auto b = buildCase(g, c.visible, 555u); + + const auto contig = b.fx->runSplitContiguous(0, b.q, c.visible, num_splits); + const auto paged = b.fx->runSplitPaged(0, b.q, c.visible, num_splits); + + EXPECT_TRUE(bitwiseEqual(contig, paged)) << "同一归约、两种取址必须逐位相同，max|diff|=" + << maxAbsDiff(toFloat(contig), toFloat(paged)); + } + } +} + +// ── 4a. visible = 0：全 0 输出，且不与单遍路径分叉 ──────────────────────── +TEST_F(SplitKvTest, EmptyVisibleWindowIsAllZero) { + PagedGeometry g = makeGeometry(4, 2, 64, 16, 4); + const auto q = randomFp16(static_cast(g.num_q_heads) * g.head_dim, 77); + const auto table = makeTable(1, g.max_num_blocks, 78); + SplitKvFixture fx(g, table, /*table_len=*/1); + + for (int num_splits : {1, 8}) { + SCOPED_TRACE("num_splits=" + std::to_string(num_splits)); + const auto out = fx.runSplitPaged(0, q, /*visible=*/0, num_splits); + for (half h : out) { + EXPECT_EQ(__half2float(h), 0.0f) << "visible=0 必须输出全 0（冻结语义）"; + } + // 也不得出现 NaN：上面用 == 比较已经隐含（NaN != 0） + } +} + +// ── 4b. visible < num_splits：大量空段，必须中性且无 NaN ────────────────── +TEST_F(SplitKvTest, VisibleSmallerThanSplitsHasNeutralEmptySplits) { + const int visible = 3; + const int bs = 16; + const int nb = requiredBlocks(visible, bs); + PagedGeometry g = makeGeometry(4, 2, 64, bs, nb + 3); + auto b = buildCase(g, visible, 4242u); + + for (int num_splits : {8, 16, 32}) { + SCOPED_TRACE("num_splits=" + std::to_string(num_splits)); + const auto out = b.fx->runSplitPaged(0, b.q, visible, num_splits); + const auto got = toFloat(out); + for (float x : got) + EXPECT_TRUE(std::isfinite(x)) << "空段导致了 NaN/Inf"; + + const auto ref = oracleAttentionDecode(toFloat(b.q), toFloat(b.k_logical), + toFloat(b.v_logical), visible, g); + EXPECT_LT(maxAbsDiff(got, ref), kOracleTolerance); + } +} + +// ── 4c. visible 不是 num_splits 的整数倍（尾段截断）──────────────────────── +TEST_F(SplitKvTest, VisibleNotMultipleOfSplitsIsCorrect) { + const int visible = 17; + const int bs = 16; + const int nb = requiredBlocks(visible, bs); + PagedGeometry g = makeGeometry(8, 2, 64, bs, nb + 3); + auto b = buildCase(g, visible, 31337u); + + for (int num_splits : {3, 5, 7}) { + SCOPED_TRACE("num_splits=" + std::to_string(num_splits)); + const auto got = toFloat(b.fx->runSplitPaged(0, b.q, visible, num_splits)); + const auto ref = oracleAttentionDecode(toFloat(b.q), toFloat(b.k_logical), + toFloat(b.v_logical), visible, g); + EXPECT_LT(maxAbsDiff(got, ref), kOracleTolerance); + } +} + +// ── 5. 非法块 id：零行语义不变，且 split 与单遍逐位一致 ──────────────────── +TEST_F(SplitKvTest, InvalidBlockIdsFollowZeroRowSemantics) { + const int block_size = 16; + const int visible = 20; + const int nb = requiredBlocks(visible, block_size); + PagedGeometry g = makeGeometry(4, 4, 32, block_size, nb + 1); + + const auto q = randomFp16(static_cast(g.num_q_heads) * g.head_dim, 902); + const auto k = randomFp16(static_cast(visible) * g.kvDim(), 900); + const auto v = randomFp16(static_cast(visible) * g.kvDim(), 901); + + for (int bad : {-1, g.max_num_blocks}) { + SCOPED_TRACE("bad block id = " + std::to_string(bad)); + // table_len = nb，但第二块 id 越界 + std::vector table = {1, bad}; + SplitKvFixture fx(g, table, nb); + fx.scatter(0, k, v, visible, 0); + fx.sync(); + + const auto single = fx.runSinglePassPaged(0, q, visible); + for (int num_splits : {1, 4}) { + const auto split = fx.runSplitPaged(0, q, visible, num_splits); + EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess) << "非法块 id 不得触发 illegal address"; + EXPECT_TRUE(bitwiseEqual(single, split)) + << "非法块 id 下 split 与单遍必须同为零行语义（num_splits=" << num_splits << "）"; + const auto got = toFloat(split); + for (float x : got) + EXPECT_TRUE(std::isfinite(x)); + } + } +} + +// ── 6. 块表长度不足：split 侧按零行处理且不越界读块表 ───────────────────── +TEST_F(SplitKvTest, ShortBlockTableIsGuardedWithoutFault) { + const int block_size = 16; + const int visible = 32; // 需要 2 块 + const int nb = requiredBlocks(visible, block_size); + PagedGeometry g = makeGeometry(4, 2, 64, block_size, 4); + + const auto table = makeTable(nb, g.max_num_blocks, 55); + const auto q = randomFp16(static_cast(g.num_q_heads) * g.head_dim, 56); + const auto k = randomFp16(static_cast(visible) * g.kvDim(), 57); + const auto v = randomFp16(static_cast(visible) * g.kvDim(), 58); + + // table_len = 1 < nb：第二个块必须按零行参与 softmax。 + SplitKvFixture fx(g, table, /*table_len=*/1); + fx.scatter(0, k, v, visible, 0); + fx.sync(); + + const auto got = toFloat(fx.runSplitPaged(0, q, visible, /*num_splits=*/4)); + EXPECT_EQ(cudaDeviceSynchronize(), cudaSuccess) << "块表长度不足不得越界读块表"; + + std::vector k_eff = toFloat(k), v_eff = toFloat(v); + for (int t = block_size; t < visible; ++t) { + for (int c = 0; c < g.kvDim(); ++c) { + k_eff[static_cast(t) * g.kvDim() + static_cast(c)] = 0.0f; + v_eff[static_cast(t) * g.kvDim() + static_cast(c)] = 0.0f; + } + } + const auto ref = oracleAttentionDecode(toFloat(q), k_eff, v_eff, visible, g); + EXPECT_LT(maxAbsDiff(got, ref), kOracleTolerance) + << "table_len 之外的 token 必须按零行参与，而不是读取 pool 内容"; +} + +// ── 7. CUDA Graph：捕获后改变可见长度并 replay，与非 graph 逐位相同 ─────── +TEST_F(SplitKvTest, GraphReplayWithGrowingVisibleMatchesEager) { + const int bs = 16; + const int max_visible = 96; + const int nb = requiredBlocks(max_visible, bs); + PagedGeometry g = makeGeometry(8, 2, 64, bs, nb + 3); + auto b = buildCase(g, max_visible, 8888u); + + const size_t q_elems = b.fx->qElems(); + DeviceBuffer d_q(q_elems), d_out_graph(q_elems); + d_q.copyFromHost(b.q.data(), q_elems); + + cudaStream_t s = nullptr; + ASSERT_EQ(cudaStreamCreate(&s), cudaSuccess); + + // 捕获：注意 num_splits 是 host 常量、grid 固定；可见长度走 device int。 + const int num_splits = 4; + ASSERT_EQ(cudaStreamBeginCapture(s, cudaStreamCaptureModeThreadLocal), cudaSuccess); + attention_decode_paged_splitkv( + d_q.data(), b.fx->kPoolData(), b.fx->vPoolData(), b.fx->tableData(), d_out_graph.data(), + b.fx->scale(), g.num_q_heads, g.num_kv_heads, g.head_dim, b.fx->lenData(), g.block_size, + g.max_num_blocks, b.fx->tableLen(), b.fx->partialData(), num_splits, s); + cudaGraph_t graph = nullptr; + ASSERT_EQ(cudaStreamEndCapture(s, &graph), cudaSuccess); + ASSERT_NE(graph, nullptr); + cudaGraphExec_t exec = nullptr; + ASSERT_EQ(cudaGraphInstantiate(&exec, graph, 0), cudaSuccess); + + // replay：可见长度逐个增长（同一 graph，grid 不变） + for (int visible : {1, 7, 17, 33, 64, max_visible}) { + SCOPED_TRACE("visible=" + std::to_string(visible)); + + b.fx->setLen(visible, s); + ASSERT_EQ(cudaGraphLaunch(exec, s), cudaSuccess); + std::vector graph_out(q_elems); + ASSERT_EQ(cudaMemcpyAsync(graph_out.data(), d_out_graph.data(), q_elems * sizeof(half), + cudaMemcpyDeviceToHost, s), + cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(s), cudaSuccess); + + const auto eager = b.fx->runSplitPaged(0, b.q, visible, num_splits); + EXPECT_TRUE(bitwiseEqual(graph_out, eager)) + << "graph replay 与 eager 不一致，max|diff|=" + << maxAbsDiff(toFloat(graph_out), toFloat(eager)); + } + + ASSERT_EQ(cudaGraphExecDestroy(exec), cudaSuccess); + ASSERT_EQ(cudaGraphDestroy(graph), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(s), cudaSuccess); +} + +// ── 8. non-default stream 与 default stream 结果一致 ────────────────────── +TEST_F(SplitKvTest, NonDefaultStreamMatchesDefaultStream) { + const int bs = 16; + const int visible = 40; + const int nb = requiredBlocks(visible, bs); + PagedGeometry g = makeGeometry(8, 2, 64, bs, nb + 3); + auto b = buildCase(g, visible, 2024u); + + cudaStream_t s = nullptr; + ASSERT_EQ(cudaStreamCreate(&s), cudaSuccess); + + const auto on_default = b.fx->runSplitPaged(0, b.q, visible, /*num_splits=*/4, /*stream=*/0); + const auto on_stream = b.fx->runSplitPaged(0, b.q, visible, /*num_splits=*/4, s); + + EXPECT_TRUE(bitwiseEqual(on_default, on_stream)) + << "非默认 stream 上结果必须一致，max|diff|=" + << maxAbsDiff(toFloat(on_default), toFloat(on_stream)); + + ASSERT_EQ(cudaStreamDestroy(s), cudaSuccess); +} From e38ed0b084a50c82473517f65f81455d56072a9c Mon Sep 17 00:00:00 2001 From: shijiashuai Date: Mon, 14 Sep 2026 22:34:15 +0800 Subject: [PATCH 2/3] test(attention): catch the dropped-rescale mutation with a deterministic late max Mutation testing over the split-KV gates (4 injections): - combine drops the l_i weight -> caught by 6 oracle-based tests, while the num_splits = 1 anchor still passes (w == 1 there). This is the concrete demonstration that the anchor alone is not enough and the independent oracle is load-bearing. - split range off-by-one -> caught by 7 tests, including both bitwise anchors. - drop combine's `m == M -> 1.0f` special case -> not caught, and that is correct: a probe confirms __expf(0.0f) == 1.0f exactly (+0 and -0), so the special case is a no-op on this toolchain. It is kept only so the anchor does not depend on __expf's behaviour at 0. - drop the online rescale in the shared loop -> NOT caught by the original matrix. The last one is a real gap in the tests, not a gate failure: with random data and short sequences the global max almost always lands in the first tile, so old_rescale is exp(0) = 1 and the rescale path is never exercised -- the injected bug passed all 11 gates. LateMaxForcesOnlineRescaleToMatter constructs the opposite deterministically (K all zero except one token in the last tile, which therefore holds the global max), forcing old_rescale != 1; the mutation is now caught by that test alone. Corrects the design record's mutation list, which had claimed the independent oracle would catch a dropped rescale. That was only true once the matrix actually reached a multi-tile rescale, which it did not. Full suite 220 passed / 11 skipped; sanitizer 0 errors on the new tests and the existing paged suites; clang-format 18.1.8 clean. Co-authored-by: CommandCodeBot --- .../decode-attention-splitkv-design.md | 12 ++++- tests/test_attention_splitkv.cpp | 52 ++++++++++++++++++- 2 files changed, 61 insertions(+), 3 deletions(-) diff --git a/docs/architecture/decode-attention-splitkv-design.md b/docs/architecture/decode-attention-splitkv-design.md index 1cef0e3..42dd9af 100644 --- a/docs/architecture/decode-attention-splitkv-design.md +++ b/docs/architecture/decode-attention-splitkv-design.md @@ -268,8 +268,16 @@ out[d] = __float2half(ACC[d] * 1/(L + 1e-9)) // +1e-9 与乘法形式与 2. 段范围 off-by-one（`count_s` 少 1 / `base_s` 偏移）→ 边界用例 + oracle 捕获； 3. 空段写垃圾而非中性值 → `visible < num_splits` 用例出现 NaN/偏差 → 捕获； 4. combine 的 `M` 不取 max（如取 `m_0`）→ 大动态范围 score 的用例捕获； - 5. **删掉在线 rescale**（两条路径同等出错）→ 逐位门禁通过、**独立 oracle 捕获** - （与 PR-2 第三项同一用意，验证分层门禁的必要性）。 + 5. **删掉在线 rescale**（两条路径同等出错）→ 逐位门禁通过 —— 但独立 oracle **在初版 + 矩阵下抓不到**：随机数据 + 短序列时全局 max 几乎总落在第一个 tile，`old_rescale` + 恒为 `exp(0) = 1`，该路径根本没被压到（实测：注入后 11 项门禁全过）。已补 + `LateMaxForcesOnlineRescaleToMatter`：确定性构造后置 max（末 tile 的单个 token 取大 + K、其余 K 为 0），使其必然经过 `old_rescale ≠ 1`，此后该项被捕获。 + **这正是 PR-2 第三项要验证的分层门禁，但前提是测试矩阵必须真的压到多 tile 的 rescale + 路径**——记录在此，以免后人以为"写了 oracle"就自动覆盖。 + 6. 去掉 combine 中 `m == M → 1.0f` 的特例：**在本工具链上不是缺陷**。探针（`__expf(0.0f) + == 1.0f` 为真，+0/-0 均成立）证实该特例是冗余的；保留它只是让 §4.4 的逐位锚点不依赖 + `__expf` 在 0 点的实现细节，而不是它在本机必需。 - **CPU/build 与 GPU correctness 分开报告**：host 校验类用例无 GPU 也运行；kernel 用例在无 device 时 **skip 且计数可见**，不得把 skip 记为 pass。 - **Sanitizer**：`compute-sanitizer --tool memcheck` 跑新增用例，要求 **0 error**；并且对 diff --git a/tests/test_attention_splitkv.cpp b/tests/test_attention_splitkv.cpp index d9a9f78..62dbce8 100644 --- a/tests/test_attention_splitkv.cpp +++ b/tests/test_attention_splitkv.cpp @@ -560,7 +560,57 @@ TEST_F(SplitKvTest, GraphReplayWithGrowingVisibleMatchesEager) { ASSERT_EQ(cudaStreamDestroy(s), cudaSuccess); } -// ── 8. non-default stream 与 default stream 结果一致 ────────────────────── +// ── 9. 多 tile 的在线 rescale 必须被真正压到（变异检验 m4 暴露的缺口）─────── +// +// 背景：随机数据 + 短序列时，全局 max 几乎总落在第一个 tile，old_rescale 恒为 +// exp(0)=1，于是"丢掉 rescale"这类实现错误不会显形（实测：初版矩阵下删掉 +// running_sum / out_acc 的 rescale 后 11 项门禁全过）。这里**确定性构造**后置 max： +// 除最后一个 tile 的一个 token 外 K 全 0、Q 全 1 ⇒ 其余 score 恒为 0，全局 max +// 必然落在后面的 tile，old_rescale 必然 ≠ 1。 +TEST_F(SplitKvTest, LateMaxForcesOnlineRescaleToMatter) { + const int bs = 16, hd = 64, hq = 4, hkv = 4, visible = 300; + const int nb = requiredBlocks(visible, bs); + PagedGeometry g = makeGeometry(hq, hkv, hd, bs, nb + 3); + + std::vector k(static_cast(visible) * g.kvDim(), __float2half(0.0f)); + std::vector v(static_cast(visible) * g.kvDim()); + std::vector q(static_cast(hq) * hd, __float2half(1.0f)); + + // V：每个 token 一个互不相同的常数，使"前段均值"与"最后一个 token"明显不同 + for (int t = 0; t < visible; ++t) + for (int c = 0; c < g.kvDim(); ++c) + v[static_cast(t) * g.kvDim() + c] = + __float2half(static_cast(t % 7) + 1.0f); + + // 后置 max：token 290 落在 [256, 300)，K 全 200 ⇒ score 远大于 0 + const int late = visible - 10; + for (int c = 0; c < g.kvDim(); ++c) + k[static_cast(late) * g.kvDim() + c] = __float2half(200.0f); + + auto b = buildCase(g, visible, 6060u); + b.fx->scatter(0, k, v, visible, 0); + b.fx->sync(); + + const auto ref = oracleAttentionDecode(toFloat(q), toFloat(k), toFloat(v), visible, g); + + // 单遍路径跨 3 个 tile（128/128/44），必然经过 rescale + const auto single = toFloat(b.fx->runSinglePassPaged(0, q, visible)); + EXPECT_LT(maxAbsDiff(single, ref), kOracleTolerance) + << "单遍路径未正确 rescale（旧 max 的权重被错误保留），max|diff|=" + << maxAbsDiff(single, ref); + + // split 路径：num_splits=1 时同样跨 tile；>1 时考验 combine 的跨段权重 + for (int num_splits : {1, 2, 4, 8}) { + SCOPED_TRACE("num_splits=" + std::to_string(num_splits)); + const auto got = toFloat(b.fx->runSplitPaged(0, q, visible, num_splits)); + for (float x : got) + EXPECT_TRUE(std::isfinite(x)); + EXPECT_LT(maxAbsDiff(got, ref), kOracleTolerance) + << "split 路径未正确合并后置 max，max|diff|=" << maxAbsDiff(got, ref); + } +} + +// ── 10. non-default stream 与 default stream 结果一致 ────────────────────── TEST_F(SplitKvTest, NonDefaultStreamMatchesDefaultStream) { const int bs = 16; const int visible = 40; From d2eb39dff4aa80545e5e4a539279a6eda3e5845f Mon Sep 17 00:00:00 2001 From: shijiashuai Date: Mon, 14 Sep 2026 22:44:49 +0800 Subject: [PATCH 3/3] fix(test): use cudaGraphInstantiateWithFlags so the test builds on CUDA 11.8 CI builds against CUDA 11.8 while the dev box has 13.3. The three-argument cudaGraphInstantiate overload is CUDA 12+, so PR-B failed to build there: error: too few arguments to function 'cudaError_t cudaGraphInstantiate(...)' Switch to cudaGraphInstantiateWithFlags, which exists since CUDA 11.4 and is present in both toolchains. The five-argument legacy form was the alternative but is not guaranteed to still exist in CUDA 13. Caught by CI, not by the local build: the kernel changes themselves compiled fine on 11.8, so this was purely a test-side API-version issue. Co-authored-by: CommandCodeBot --- tests/test_attention_splitkv.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_attention_splitkv.cpp b/tests/test_attention_splitkv.cpp index 62dbce8..8ebacd0 100644 --- a/tests/test_attention_splitkv.cpp +++ b/tests/test_attention_splitkv.cpp @@ -535,7 +535,9 @@ TEST_F(SplitKvTest, GraphReplayWithGrowingVisibleMatchesEager) { ASSERT_EQ(cudaStreamEndCapture(s, &graph), cudaSuccess); ASSERT_NE(graph, nullptr); cudaGraphExec_t exec = nullptr; - ASSERT_EQ(cudaGraphInstantiate(&exec, graph, 0), cudaSuccess); + // 用 InstantiateWithFlags（CUDA 11.4 起）：三参数的 cudaGraphInstantiate 重载是 CUDA 12+ + // 才有的，而 CI 跑在 CUDA 11.8 上（本地是 13.3）；五参数旧式在 CUDA 13 已不保证存在。 + ASSERT_EQ(cudaGraphInstantiateWithFlags(&exec, graph, 0), cudaSuccess); // replay：可见长度逐个增长（同一 graph，grid 不变） for (int visible : {1, 7, 17, 33, 64, max_visible}) {