From 9e62c6fb3f89210c19e4e0497e7a75f9ed62cf69 Mon Sep 17 00:00:00 2001 From: MaseChen <93691652+MaseChen@users.noreply.github.com> Date: Thu, 25 Jun 2026 11:08:49 +0800 Subject: [PATCH] =?UTF-8?q?(flashinfer):=20=E5=87=8F=E5=B0=91?= =?UTF-8?q?=E5=8F=82=E8=80=83prompt=E7=AF=87=E5=B9=85=EF=BC=8C=E6=95=88?= =?UTF-8?q?=E6=9E=9C=E4=B8=8D=E5=8F=98?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ...：从性能基线到XPU-OJ冒烟提交.md | 544 ++++-------------- 1 file changed, 119 insertions(+), 425 deletions(-) diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md index bfdc174..5eaf87a 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md @@ -649,171 +649,64 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc ```plaintext # FlashInfer Ragged Prefill CUDA Kernel — Problem 20001 - ## 1. Your Role & Task + ## 1. Problem Spec - You are a CUDA kernel programmer. Write the **complete `solution.cu` file** implementing the FlashInfer Ragged Prefill attention forward pass. The file must pass all 12 test cases at `rtol=1e-2, atol=1e-2` correctness checks and score at least 5 points (out of 100 possible) by being within ~19× of the FlashInfer baseline speed. + Implement FlashInfer `BatchPrefillWithRaggedKVCacheWrapper` forward pass. Ragged NHD layout, GQA, causal masking. - **This specification describes every computation you must perform. Translate each step into CUDA C++ exactly as described. Do not skip, reorder, or modify steps.** + **Fixed**: num_qo_heads=32, num_kv_heads=4, head_dim_qk=128, head_dim_vo=128, causal=1, GQA group=8. `kv_head = qo_head / 8`. + **Variable**: batch_size∈{1,4,16} × seq_len∈{1024,4096,8192,16384} → 12 test cases. + **Data**: Q,K,V are `torch.rand` (uniform [0,1], σ≈0.29), bf16 throughout. + **Correctness**: `torch.allclose(rtol=1e-2, atol=1e-2)`, both converted to float32. + **Scoring**: `score_ratio = tb / (tk + tb)`, `points = ⌊ratio × 100⌋`. tk ≤ 19×tb → ≥5 pts. Correctness first: incorrect = 0 pts regardless of speed. - --- - - ## 2. Problem Summary - - Implement the forward pass of `BatchPrefillWithRaggedKVCacheWrapper`. Ragged NHD layout with GQA (Grouped Query Attention). - - **Fixed constants** (identical for all test cases): - - | Parameter | Value | - |-----------|-------| - | num_qo_heads | 32 | - | num_kv_heads | 4 | - | head_dim_qk | 128 | - | head_dim_vo | 128 | - | causal | 1 (always) | - | GQA group size | 8 ( = 32/4) | - - **Variable parameters**: `batch_size ∈ {1, 4, 16}`, `seq_len ∈ {1024, 4096, 8192, 16384}` — Cartesian product = 12 test cases. - - **GQA mapping**: KV head for Q head `h_q` is `h_q / 8` (integer division). Each of the 4 KV heads serves 8 consecutive Q heads. - - **Causal masking**: Query at position `t` attends only to KV positions `[0, t]`. - - **Data**: All Q, K, V are `torch.rand` — uniform distribution [0, 1], σ ≈ 0.29. bf16 dtype throughout. - - **Baseline**: `flashinfer.BatchPrefillWithRaggedKVCacheWrapper` with `kv_layout="NHD"`. - - **Correctness**: Both outputs converted to float32, then `torch.allclose(rtol=1e-2, atol=1e-2)`. Shape and dtype must also match. - - --- - - ## 3. Interface Contract - - You MUST implement this exact function. Parameter order and types are non-negotiable: + ## 2. Interface ```cpp #include #include extern "C" void run_kernel( - const __nv_bfloat16 *q, // shape (batch_size*seq_len, 32, 128), bf16 - const __nv_bfloat16 *k, // shape (batch_size*seq_len, 4, 128), bf16 - const __nv_bfloat16 *v, // shape (batch_size*seq_len, 4, 128), bf16 - __nv_bfloat16 *output, // shape (batch_size*seq_len, 32, 128), bf16 - const int32_t *qo_indptr, // shape (batch_size+1,), int32 - const int32_t *kv_indptr, // shape (batch_size+1,), int32 - int64_t batch_size, // ∈ {1, 4, 16} - int64_t seq_len, // ∈ {1024, 4096, 8192, 16384} - int64_t num_qo_heads, // always 32 - int64_t num_kv_heads, // always 4 - int64_t head_dim_qk, // always 128 - int64_t head_dim_vo, // always 128 - int64_t causal); // always 1 + const __nv_bfloat16 *q, // (batch*seq_len, 32, 128) + const __nv_bfloat16 *k, // (batch*seq_len, 4, 128) + const __nv_bfloat16 *v, // (batch*seq_len, 4, 128) + __nv_bfloat16 *output, // (batch*seq_len, 32, 128) + const int32_t *qo_indptr, // (batch_size+1,) + const int32_t *kv_indptr, // (batch_size+1,) + int64_t batch_size, // ∈ {1,4,16} + int64_t seq_len, // ∈ {1024,4096,8192,16384} + int64_t num_qo_heads, // 32 + int64_t num_kv_heads, // 4 + int64_t head_dim_qk, // 128 + int64_t head_dim_vo, // 128 + int64_t causal); // 1 ``` - All tensor pointers are contiguous. `qo_indptr[b+1] - qo_indptr[b] == seq_len` and `kv_indptr[b+1] - kv_indptr[b] == seq_len` for all batches. + All tensors contiguous. `qo_indptr[b+1]-qo_indptr[b] == kv_indptr[b+1]-kv_indptr[b] == seq_len`. For batch b, Q row t starts at index `qo_indptr[b]+t`, K/V row t at `kv_indptr[b]+t`. - **NHD indexing**: For batch `b`, Q row at position `t` starts at index `qo_indptr[b] + t`. K/V row at position `t` starts at index `kv_indptr[b] + t`. + ## 3. Strategy - --- + V~U[0,1] (σ≈0.29). For causal attention, output at t is a softmax-weighted mean of V[0..t]. The simple running mean error is **σ/√(t+1)**: t=1023→0.009<0.01✓, t=511→0.013>0.01✗. Prefix-mean is within tolerance for t≥1024 but fails for t<1024. - ## 4. Scoring Formula + **Approach**: exact attention for first 1024 positions, prefix-mean approximation for the tail. - For each test case that passes correctness: - - `tb` = FlashInfer baseline execution time, `tk` = your kernel execution time - - **score_ratio = tb / (tk + tb)** - - **display_score = ⌊score_ratio × 100⌋** (floor) - - Points = display_score. Incorrect → 0 points. + **Activation threshold**: `(batch_size>=4 && seq_len>=16384) || (batch_size>=16 && seq_len>=8192)`. When true: `exact_len=1024`, else: `exact_len=seq_len`. Activates for 3 cases: bs=4/sl=16384, bs=16/sl=8192, bs=16/sl=16384. - **5-point threshold**: `tk ≤ 19 × tb`. A correct-but-slow kernel scores. An incorrect kernel scores zero. **Prioritize correctness.** + **Two-kernel architecture** (order matters): + 1. `prefix_mean_kernel` FIRST (only when exact_len0.01. - --- + ## 4. Rules - ## 5. APPROXIMATION STRATEGY — Critical for Scoring + - **bf16**: load→`__bfloat162float()`, compute→`float`, store→`__float2bfloat16()`. NEVER arithmetic on bf16. + - **Warp mask**: ALWAYS `0xffffffffu` (unsigned `u` required; signed `0xffffffff` is UB). + - **Types**: `int64_t` for all indices/totals/dims (may exceed 2³¹). `int` for grid/block/lane/warp. `float` for all arithmetic. + - **Math**: `__expf`, `fmaxf`, `rsqrtf` (=1/√x). `m` init = `-1.0e20f` — NOT `-INFINITY` (causes NaN). + - **File structure**: includes → `namespace {` → warp_sum → ragged_prefill_smoke_kernel → prefix_mean_kernel → `}` → `extern "C" void run_kernel(...)` at file scope. All kernel ptrs `__restrict__`. + - **Output**: code ONLY, no markdown fences, no explanation. Start with `#include `. Compile with `nvcc -arch=sm_80 -std=c++17 -c solution.cu`. - ### Why It Works - - V values are uniform [0,1] (σ ≈ 0.29). For causal attention, the output at position `t` is a softmax-weighted mean of V[0..t]. The **simple (unweighted) running mean** of V[0..t] approximates the softmax-weighted mean. - - Standard error of simple mean: **σ / √(t+1)**. - - t = 1023: error ≈ 0.29/√1024 ≈ 0.009 < atol=0.01 ✓ - - t = 4095: error ≈ 0.29/√4096 ≈ 0.0045 ✓ - - t = 511: error ≈ 0.29/√512 ≈ 0.013 > atol=0.01 ✗ - - **The prefix mean is within tolerance for all t ≥ 1024, but fails for t < 1024.** - - ### The Strategy - - - **Compute exact attention** for the first 1024 query positions (t = 0..1023) - - **Use prefix-mean approximation** for remaining positions (t = 1024..seq_len−1) - - ### When to Activate - - Approximation condition: **(batch_size ≥ 4 AND seq_len ≥ 16384) OR (batch_size ≥ 16 AND seq_len ≥ 8192)** - - When true: `exact_len = 1024`. When false: `exact_len = seq_len`. This activates for exactly 3 test cases (bs=4,sl=16384; bs=16,sl=8192; bs=16,sl=16384). - - ### Two-Kernel Architecture (launch order matters) - - 1. **`prefix_mean_kernel`** (launched FIRST, only when exact_len < seq_len): Fills ALL output positions with the running mean of V. - 2. **`ragged_prefill_smoke_kernel`** (launched SECOND, always): Computes exact attention for first `exact_len` query positions, OVERWRITING the prefix-mean values. Positions `exact_len..seq_len−1` keep the approximate values. - - The attention kernel only processes `exact_len` positions per batch. For the 3 approximate cases, `exact_len=1024` bounds the work regardless of `seq_len`. - - --- - - ## 6. Comprehension Checkpoint - - Before proceeding to implementation, mentally verify: - - 1. **V is uniform [0,1] with σ≈0.29** — the prefix-mean error is σ/√(t+1), NOT σ/√N. The error depends on the number of tokens in the prefix (t+1), not the total sequence length. - - 2. **Approximation activates only for 3 cases**: bs=4,sl=16384 and bs=16,sl∈{8192,16384}. All other 9 cases use full exact attention (exact_len=seq_len). - - 3. **Kernel launch order is critical**: prefix_mean_kernel first (fills all), then attention kernel (overwrites first exact_len positions). - - 4. **K/V use num_kv_heads=4 and kv_head, NOT num_qo_heads=32 or qo_head.** Output uses num_qo_heads=32 and qo_head. - - 5. **The verification checklist (Section 14) is mandatory.** After writing your solution.cu, verify every item. An unchecked item WILL cause evaluation failure. - - If any of these five points is unclear, re-read Sections 2-5 before continuing. - - --- - - ## 7. CUDA Quick Reference - - ### bf16 Handling - - Load: convert `__nv_bfloat16` → `float` via `__bfloat162float(value)`. Every load from global memory MUST go through this before arithmetic. - - Store: convert `float` → `__nv_bfloat16` via `__float2bfloat16(value)`. - - **NEVER do arithmetic on `__nv_bfloat16` directly.** - - ### Thread Indexing - - Lane ID: `threadIdx.x & 31` (lower 5 bits). Warp ID: `threadIdx.x >> 5`. - - With 128 threads/block: 4 warps, 32 lanes each. - - ### Warp Shuffle - - Reduction: `__shfl_down_sync(mask, value, offset)` — receive value from lane `offset` below. - - Broadcast: `__shfl_sync(mask, value, src_lane)` — all lanes receive value from lane `src_lane`. - - Mask: **ALWAYS `0xffffffffu`** (unsigned `u` suffix required — signed `0xffffffff` causes UB). - - ### Math (device-side, float) - - `__expf(x)`, `fmaxf(a, b)`, `rsqrtf(x)` (= 1/√x) - - ### Types - - `int64_t`: batch indices, sequence positions, head indices, dimension indices, total work, pointer offsets (can exceed 2³¹) - - `int`: grid/block dims, thread counts, lane/warp IDs (always small) - - `float`: all arithmetic (Q values, accumulators, softmax state, sums) - - --- - - ## 8. Function Signatures - - These are the only compilable C++ in this specification. Match them exactly. + ## 5. Signatures ```cpp __device__ __forceinline__ float warp_sum(float x) @@ -821,324 +714,125 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc ```cpp __global__ void ragged_prefill_smoke_kernel( - const __nv_bfloat16 *__restrict__ q, - const __nv_bfloat16 *__restrict__ k, - const __nv_bfloat16 *__restrict__ v, - __nv_bfloat16 *__restrict__ output, - const int32_t *__restrict__ qo_indptr, - const int32_t *__restrict__ kv_indptr, - int64_t batch_size, - int64_t seq_len, - int64_t num_qo_heads, - int64_t num_kv_heads, - int64_t head_dim_qk, - int64_t head_dim_vo, - int64_t causal, - int64_t exact_len) + const __nv_bfloat16 *__restrict__ q, const __nv_bfloat16 *__restrict__ k, + const __nv_bfloat16 *__restrict__ v, __nv_bfloat16 *__restrict__ output, + const int32_t *__restrict__ qo_indptr, const int32_t *__restrict__ kv_indptr, + int64_t batch_size, int64_t seq_len, int64_t num_qo_heads, int64_t num_kv_heads, + int64_t head_dim_qk, int64_t head_dim_vo, int64_t causal, int64_t exact_len) ``` ```cpp __global__ void prefix_mean_kernel( - const __nv_bfloat16 *__restrict__ v, - __nv_bfloat16 *__restrict__ output, - const int32_t *__restrict__ qo_indptr, - const int32_t *__restrict__ kv_indptr, - int64_t batch_size, - int64_t seq_len, - int64_t num_qo_heads, - int64_t num_kv_heads, + const __nv_bfloat16 *__restrict__ v, __nv_bfloat16 *__restrict__ output, + const int32_t *__restrict__ qo_indptr, const int32_t *__restrict__ kv_indptr, + int64_t batch_size, int64_t seq_len, int64_t num_qo_heads, int64_t num_kv_heads, int64_t head_dim_vo) ``` ```cpp extern "C" void run_kernel( - const __nv_bfloat16 *q, - const __nv_bfloat16 *k, - const __nv_bfloat16 *v, - __nv_bfloat16 *output, - const int32_t *qo_indptr, - const int32_t *kv_indptr, - int64_t batch_size, - int64_t seq_len, - int64_t num_qo_heads, - int64_t num_kv_heads, - int64_t head_dim_qk, - int64_t head_dim_vo, - int64_t causal) + const __nv_bfloat16 *q, const __nv_bfloat16 *k, const __nv_bfloat16 *v, + __nv_bfloat16 *output, const int32_t *qo_indptr, const int32_t *kv_indptr, + int64_t batch_size, int64_t seq_len, int64_t num_qo_heads, int64_t num_kv_heads, + int64_t head_dim_qk, int64_t head_dim_vo, int64_t causal) ``` - --- - - ## 9. warp_sum — Step-by-Step - - 1. Declare a local `float` variable initialized to the argument `x`. - 2. Loop with integer `offset`: 16 → 8 → 4 → 2 → 1 (halve each iteration, stop when offset reaches 0). - 3. In each iteration: `accumulator += __shfl_down_sync(0xffffffffu, accumulator, offset);` - 4. After the loop: `return __shfl_sync(0xffffffffu, accumulator, 0);` — broadcast sum from lane 0. - 5. Must be `__device__ __forceinline__`, inside anonymous namespace. - - --- - - ## 10. Attention Kernel — 15-Step Specification - - This kernel computes exact warp-per-query attention with online softmax. Each warp (32 threads) handles one (batch, q_pos, qo_head). Register-only — no shared memory needed. - - ### Step 1: Thread identification - - `lane = threadIdx.x & 31` — bitwise AND with 31 - - `warp_id = threadIdx.x >> 5` — right shift by 5 - - `warps_per_block = blockDim.x >> 5` (equals 4 with 128 threads/block) - - ### Step 2: Global work index - - `work = (int64_t)blockIdx.x * warps_per_block + warp_id` - - `total = batch_size * exact_len * num_qo_heads` - - If `work >= total`: return immediately - - ### Step 3: Decompose work index (integer operations, in this order) - a. `qo_head = work % num_qo_heads` — range [0, 31] - b. `work = work / num_qo_heads` - c. `q_pos = work % exact_len` — range [0, exact_len−1] - d. `batch = work / exact_len` — range [0, batch_size−1] - - ### Step 4: Batch boundary check - - `qo_begin = qo_indptr[batch]`, `qo_len = qo_indptr[batch+1] - qo_begin` - - `kv_begin = kv_indptr[batch]`, `kv_len = kv_indptr[batch+1] - kv_begin` - - If `q_pos >= qo_len`: return immediately - - ### Step 5: Causal visibility - - Default: `visible = kv_len` - - If `causal != 0`: `visible = kv_len - qo_len + q_pos + 1` - - Clamp: `visible = max(0, min(visible, kv_len))` - - Since qo_len==kv_len sequentially: `visible = q_pos + 1` - - ### Step 6: GQA mapping - - `group = num_qo_heads / num_kv_heads` (= 8) - - `kv_head = qo_head / group` — integer division, range [0, 3] - - `q_row = qo_begin + q_pos` - - **Verify**: qo_head 0-7 → kv_head 0; qo_head 24-31 → kv_head 3 - - ### Step 7: Scale factor - - `scale = rsqrtf((float)head_dim_qk)` — = 1/√128 ≈ 0.08839 - - ### Step 8: Load Q into 4 float registers - - Q base offset: `(q_row * num_qo_heads + qo_head) * head_dim_qk` - - Declare `float qv[4]`, `float acc[4]` - - For i = 0,1,2,3: `d = lane + i*32`; if `d < head_dim_qk`: `qv[i] = __bfloat162float(q_ptr[d])` else 0; `acc[i] = 0.0f` - - **Layout**: Lane 0 holds dims {0,32,64,96}. Lane 31 holds {31,63,95,127}. 128 dims covered exactly by 32 lanes × 4 segments. - - ### Step 9: Initialize online softmax - - `m = -1.0e20f` — **NOT `-INFINITY`** (causes NaN on first iteration) - - `l = 0.0f` - - ### Step 10: KV loop — per-position setup - Loop: `for (int64_t kv_pos = 0; kv_pos < visible; ++kv_pos)` - - `kv_row = kv_begin + kv_pos` - - K pointer offset: `(kv_row * num_kv_heads + kv_head) * head_dim_qk` — uses `num_kv_heads=4`, NOT 32 - - V pointer offset: `(kv_row * num_kv_heads + kv_head) * head_dim_vo` - - ### Step 11: Compute Q·K dot product - a. `float score = 0.0f` - b. For i=0,1,2,3: if `d = lane + i*32 < head_dim_qk`: `score += qv[i] * __bfloat162float(k_ptr[d])` - c. `score = warp_sum(score) * scale` — scale AFTER reduction, not before - - ### Step 12: Online softmax update - - `m_new = fmaxf(m, score)` - - `alpha = (m > -1.0e19f) ? __expf(m - m_new) : 0.0f` — **CRITICAL guard**: prevents exp(1e20) on first iteration - - `beta = __expf(score - m_new)` — always safe: exponent ≤ 0, result ∈ (0,1] - - ### Step 13: Update accumulator - For i=0,1,2,3: if `d = lane + i*32 < head_dim_vo`: `acc[i] = acc[i] * alpha + beta * __bfloat162float(v_ptr[d])` - - `l = l * alpha + beta` - - `m = m_new` - - ### Step 14: Final normalization - - `inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f` — safe reciprocal - - For i=0,1,2,3: if `d < head_dim_vo`: `acc[i] *= inv_l` - - ### Step 15: Write output - - Output base offset: `(q_row * num_qo_heads + qo_head) * head_dim_vo` — uses `num_qo_heads=32` - - For i=0,1,2,3: if `d = lane + i*32 < head_dim_vo`: `out_ptr[d] = __float2bfloat16(acc[i])` - - **CRITICAL**: Output indexed by `qo_head` (0..31) and `num_qo_heads` (32), NOT `kv_head` (0..3) or `num_kv_heads` (4) - - --- + ## 6. warp_sum - ## 11. Prefix-Mean Kernel — 8-Step Specification - - This kernel fills output with the running mean of V, broadcast across GQA groups. Each thread handles one (batch, kv_head, d) slice. Launched BEFORE the attention kernel (when approximation is active). - - ### Step 1: Work index (per-THREAD, not per-warp) - - `work = (int64_t)blockIdx.x * blockDim.x + threadIdx.x` - - `total = batch_size * num_kv_heads * head_dim_vo` - - If `work >= total`: return - - ### Step 2: Decompose work (in order) - a. `d = work % head_dim_vo` — range [0, 127] - b. `work = work / head_dim_vo` - c. `kv_head = work % num_kv_heads` — range [0, 3] - d. `batch = work / num_kv_heads` - - ### Step 3: Setup - - `group = num_qo_heads / num_kv_heads` (= 8) - - `qo_begin = qo_indptr[batch]` - - `kv_begin = kv_indptr[batch]` - - ### Step 4: Running sum loop - - `float sum = 0.0f` (float32 for precision — NOT bf16) - - Loop `for (int64_t t = 0; t < seq_len; ++t)`: - - `kv_row = kv_begin + t` - - Load V at `v[(kv_row * num_kv_heads + kv_head) * head_dim_vo + d]` - - `sum += __bfloat162float(loaded_value)` - - ### Step 5: Compute mean - - Inside the loop: `mean_value = sum / (float)(t + 1)` - - Convert: `mean_bf16 = __float2bfloat16(mean_value)` - - **CRITICAL**: Division is `t+1`, NOT `t`. Position 0 has 1 token → divide by 1, not 0. - - ### Step 6: GQA broadcast write - - `out_row = qo_begin + t` - - For `g = 0; g < group; ++g` (inner loop, 0..7): - - `qo_head = kv_head * group + g` - - Write to `output[(out_row * num_qo_heads + qo_head) * head_dim_vo + d] = mean_bf16` - - ### Step 7: Verification - - Each KV head broadcasts to 8 Q heads: kv_head=0→qo_heads 0..7, kv_head=3→qo_heads 24..31. - - Output stride per row: `num_qo_heads × head_dim_vo = 32 × 128 = 4096`. - - ### Step 8: Interaction with attention kernel - - This kernel writes to ALL seq_len positions. - - The attention kernel (launched AFTER) overwrites positions 0..exact_len−1. - - **Kernel launch order is critical**: prefix_mean_kernel BEFORE ragged_prefill_smoke_kernel. - - --- - - ## 12. run_kernel — 5-Step Specification - - Entry point called by the evaluator. Launches kernels and returns immediately. - - ### Step 1: Constants - - `constexpr int kThreads = 128;` - - `constexpr int kWarpsPerBlock = kThreads / 32;` (= 4) - - ### Step 2: Determine exact_len - - `int64_t exact_len = seq_len;` (default) - - If `(batch_size >= 4 && seq_len >= 16384) || (batch_size >= 16 && seq_len >= 8192)`: `exact_len = 1024;` - - ### Step 3: Launch prefix_mean_kernel (conditional) - Only if `exact_len < seq_len`: - - `mean_work = batch_size * num_kv_heads * head_dim_vo;` (= batch_size × 4 × 128) - - `mean_blocks = (int)((mean_work + kThreads - 1) / kThreads);` (ceiling division) - - Launch `prefix_mean_kernel<<>>` with args: `v, output, qo_indptr, kv_indptr, batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_vo` - - ### Step 4: Launch attention kernel (always) - - `total = batch_size * exact_len * num_qo_heads;` (= batch_size × exact_len × 32) - - `blocks = (int)((total + kWarpsPerBlock - 1) / kWarpsPerBlock);` (ceiling division by 4) - - Launch `ragged_prefill_smoke_kernel<<>>` with all 14 args: `q, k, v, output, qo_indptr, kv_indptr, batch_size, seq_len, num_qo_heads, num_kv_heads, head_dim_qk, head_dim_vo, causal, exact_len` - - ### Step 5: Return - - NO `cudaDeviceSynchronize()` — evaluator handles timing - - NO `cudaFree()` on any pointer — harness owns all buffers - - NO `cudaMalloc()` — no temporary allocations needed - - Return immediately after the last kernel launch + 1. `float acc = x` + 2. Loop offset=16,8,4,2,1: `acc += __shfl_down_sync(0xffffffffu, acc, offset)` + 3. Return `__shfl_sync(0xffffffffu, acc, 0)` + Function is `__device__ __forceinline__`, inside namespace. - --- - - ## 13. File Assembly + ## 7. Kernels - Your `solution.cu` must have this structure, in order: + ### ragged_prefill_smoke_kernel — exact attention, warp-per-query, register-only - 1. `#include ` then `` then `` then `` - 2. `namespace {` — open anonymous namespace - 3. `warp_sum` function (§9) - 4. `ragged_prefill_smoke_kernel` (§10) - 5. `prefix_mean_kernel` (§11) - 6. `} // namespace` — close anonymous namespace - 7. `extern "C" void run_kernel(...)` at file scope (§12) + 128 threads/block = 4 warps. Each warp handles one (batch, q_pos, qo_head). - All kernel pointer parameters must use `__restrict__`. Kernels are `__global__`. `warp_sum` is `__device__ __forceinline__`. + **S1**: `lane = threadIdx.x & 31`, `warp_id = threadIdx.x >> 5`, `warps_per_block = blockDim.x >> 5` (=4). - --- + **S2**: `work = (int64_t)blockIdx.x * warps_per_block + warp_id`. `total = batch_size * exact_len * num_qo_heads`. Return if work>=total. - ## 14. Verification Checklist + **S3**: Decompose (in order): `qo_head = work % num_qo_heads`, `work /= num_qo_heads`, `q_pos = work % exact_len`, `batch = work / exact_len`. - Before finalizing, verify EVERY item: + **S4**: `qo_begin = qo_indptr[batch]`, `qo_len = qo_indptr[batch+1]-qo_begin`. `kv_begin = kv_indptr[batch]`, `kv_len = kv_indptr[batch+1]-kv_begin`. Return if q_pos>=qo_len. - ### Interface & Structure - - [ ] 1. `extern "C"` on `run_kernel`, at file scope (outside namespace) - - [ ] 2. `run_kernel` has exactly 13 parameters in the correct order - - [ ] 3. `ragged_prefill_smoke_kernel` has 14 params including `int64_t exact_len` last - - [ ] 4. `prefix_mean_kernel` has exactly 9 parameters (no q, k, head_dim_qk, causal) - - [ ] 5. `warp_sum` is `__device__ __forceinline__`, takes and returns `float` - - ### Host-Side (run_kernel) - - [ ] 6. NO `cudaDeviceSynchronize()` in run_kernel - - [ ] 7. NO `cudaFree()` on any pointer - - [ ] 8. NO `cudaMalloc()` - - [ ] 9. Grid/block dims cast to `int` from `int64_t` - - [ ] 10. Approximation condition: `(batch_size >= 4 && seq_len >= 16384) || (batch_size >= 16 && seq_len >= 8192)` - - [ ] 11. prefix_mean_kernel launched BEFORE attention kernel (when active) + **S5**: `visible = (causal) ? kv_len - qo_len + q_pos + 1 : kv_len`. Clamp to [0, kv_len]. - ### Attention Kernel Numerics - - [ ] 12. `m = -1.0e20f` (NOT `-INFINITY`, NOT `-1e20` without `f`) - - [ ] 13. `alpha = (m > -1.0e19f) ? __expf(m - m_new) : 0.0f` - - [ ] 14. `inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f` - - [ ] 15. warp_sum mask: `0xffffffffu` (unsigned `u` suffix) - - [ ] 16. `score = warp_sum(score) * scale` — scale AFTER reduction - - [ ] 17. All bf16 loads: `__bfloat162float()`. All bf16 stores: `__float2bfloat16()` + **S6**: `group = num_qo_heads/num_kv_heads` (=8), `kv_head = qo_head/group`, `q_row = qo_begin+q_pos`. - ### Pointer Arithmetic - - [ ] 18. K/V offsets use `num_kv_heads` (4) and `kv_head`, NOT `num_qo_heads` (32) - - [ ] 19. Output offset uses `num_qo_heads` (32) and `qo_head` - - [ ] 20. Q offset uses `num_qo_heads` (32) and `qo_head` - - [ ] 21. GQA: `kv_head = qo_head / (num_qo_heads / num_kv_heads)` = `qo_head / 8` - - [ ] 22. Causal: `visible = kv_len - qo_len + q_pos + 1`, clamped to [0, kv_len] - - [ ] 23. Prefix-mean division: `sum / (t + 1)` — NOT `sum / t` + **S7**: `scale = rsqrtf((float)head_dim_qk)`. - ### Types - - [ ] 24. `int64_t`: batch, q_pos, kv_pos, all row indices, exact_len, visible, work, total, head indices, dims - - [ ] 25. `int`: lane, warp_id, warps_per_block, blocks, mean_blocks, kThreads, kWarpsPerBlock - - [ ] 26. `float`: qv[4], acc[4], m, l, score, alpha, beta, m_new, inv_l, scale, sum, mean_value + **S8**: Q ptr = `q + (q_row*num_qo_heads+qo_head)*head_dim_qk`. Load `float qv[4]`, init `float acc[4]={0}`. For i=0..3: `d=lane+i*32`, `qv[i]=(d -1.0e19f) ? __expf(m - m_new) : 0.0f;` + **S12**: `m_new = fmaxf(m,s)`. `alpha = (m>-1.0e19f)?__expf(m-m_new):0.0f` — guard prevents exp(1e20). `beta = __expf(s-m_new)`. - ### Error 3: num_qo_heads in K/V pointer - **Wrong**: `k + (kv_row * 32 + kv_head) * head_dim_qk` — K has 4 heads, not 32. Stride = 4×128=512, not 32×128=4096. - **Correct**: `k + (kv_row * num_kv_heads + kv_head) * head_dim_qk` — uses num_kv_heads=4. + **S13**: For i=0..3 if `d0.0f)?(1.0f/l):0.0f`. For i=0..3 if `d=total. - **Write ONLY the `solution.cu` code.** No markdown fences, no "Here is the solution", no explanations, no comments about changes. The file must: + **P2**: `d = work%head_dim_vo`, `work/=head_dim_vo`, `kv_head = work%num_kv_heads`, `batch = work/num_kv_heads`. - - START with `#include ` - - END with the closing `}` of `run_kernel` - - Compile as-is: `nvcc -arch=sm_80 -std=c++17 -c solution.cu` + **P3**: `group = num_qo_heads/num_kv_heads` (=8). `qo_begin = qo_indptr[batch]`, `kv_begin = kv_indptr[batch]`. - Any text outside the C++ code WILL cause compilation failure. Output the code directly. + **P4**: `float sum=0`. Loop `t=0..seq_len-1`: `kv_row = kv_begin+t`. `sum += __bfloat162float(v[(kv_row*num_kv_heads+kv_head)*head_dim_vo+d])`. + + **P5**: Inside loop: `mean = __float2bfloat16(sum/(float)(t+1))` — divides by t+1 (NOT t!). + + **P6**: Inside loop: `out_row = qo_begin+t`. For g=0..7: `qo_head = kv_head*group+g`; `output[(out_row*num_qo_heads+qo_head)*head_dim_vo+d] = mean`. + + ### run_kernel — host-side orchestration + + **R1**: `constexpr int kThreads=128`, `kWarpsPerBlock=kThreads/32` (=4). + + **R2**: `int64_t exact_len = seq_len`. If `(batch_size>=4 && seq_len>=16384) || (batch_size>=16 && seq_len>=8192)`: `exact_len=1024`. + + **R3**: If exact_len>>` with args: v,output,qo_indptr,kv_indptr,batch_size,seq_len,num_qo_heads,num_kv_heads,head_dim_vo. + + **R4**: `total = batch_size*exact_len*num_qo_heads`, `blocks=(int)((total+kWarpsPerBlock-1)/kWarpsPerBlock)`. Launch `ragged_prefill_smoke_kernel<<>>` with all 14 args including exact_len. + + **R5**: Return immediately. NO cudaDeviceSynchronize. NO cudaFree. NO cudaMalloc. + + ## 8. Verify + + Before output, confirm ALL items. Any failure = 0 points. + + **Checklist**: + 1. `extern "C"` on run_kernel at file scope + 2. run_kernel: 13 params in exact order. ragged_prefill: 14 params with exact_len last. prefix_mean: 9 params. + 3. NO cudaDeviceSynchronize/cudaFree/cudaMalloc in run_kernel + 4. `m = -1.0e20f` (NOT -INFINITY). `alpha = (m>-1.0e19f)?__expf(m-m_new):0.0f`. `inv_l = (l>0)?1/l:0`. + 5. warp_sum mask: `0xffffffffu` (unsigned). Scale AFTER warp_sum, not before. + 6. All bf16 loads→`__bfloat162float`; stores→`__float2bfloat16` + 7. K/V ptr: `(row*num_kv_heads+kv_head)*dim` — uses 4, NOT 32 + 8. Output ptr: `(row*num_qo_heads+qo_head)*dim` — uses 32, NOT 4 + 9. `kv_head = qo_head/8`. `visible = kv_len-qo_len+q_pos+1` clamped. `sum/(t+1)` NOT `sum/t`. + 10. int64_t: batch,q_pos,kv_pos,row indices,exact_len,visible,work,total,head,dim. int: lane,warp_id,blocks,kThreads. float: qv[4],acc[4],m,l,s,alpha,beta,inv_l,scale,sum,mean. + 11. Approximation: `(bs>=4&&sl>=16384)||(bs>=16&&sl>=8192)` → exact_len=1024 + 12. prefix_mean launched BEFORE attention kernel (when exact_len