forked from metax-maca/op_optimization
<docs>(flashinfer): 减少参考prompt篇幅，效果不变
This commit is contained in:
MaseChen
parent
c0dbfb6c52
commit
9e62c6fb3f
|
|
@ -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 <stdint.h>
|
||||
#include <cuda_bf16.h>
|
||||
|
||||
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_len<seq_len): fills ALL output with running mean of V
|
||||
2. `ragged_prefill_smoke_kernel` SECOND (always): exact attention for first exact_len positions, overwriting prefix-mean
|
||||
|
||||
| Test Cases (all 12) | batch_size | seq_len | Total Q×KV pairs |
|
||||
|---------------------|------------|---------|-------------------|
|
||||
| 1-4 | 1 | 1024, 4096, 8192, 16384 | 524K – 134M |
|
||||
| 5-8 | 4 | 1024, 4096, 8192, 16384 | 2.1M – 537M |
|
||||
| 9-12 | 16 | 1024, 4096, 8192, 16384 | 8.4M – 2,147M |
|
||||
Key: K/V strides use num_kv_heads=4 and kv_head (NOT 32/qo_head). Output strides use num_qo_heads=32 and qo_head (NOT 4/kv_head). If V were randn (σ≈1.0), the approximation would fail: 1.0/√1024≈0.031>0.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 <stdint.h>`. 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<<<mean_blocks, kThreads>>>` 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<<<blocks, kThreads>>>` 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 <stdint.h>` then `<cuda_bf16.h>` then `<cuda_runtime.h>` then `<math.h>`
|
||||
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<head_dim_qk)?__bfloat162float(q_ptr[d]):0`.
|
||||
|
||||
---
|
||||
**S9**: `m = -1.0e20f` (NOT -INFINITY!), `l = 0.0f`.
|
||||
|
||||
## 15. Common CUDA Errors
|
||||
**S10**: Loop `kv_pos=0..visible-1`: `kv_row = kv_begin+kv_pos`. K ptr = `k + (kv_row*num_kv_heads+kv_head)*head_dim_qk` — uses num_kv_heads=4. V ptr = `v + (kv_row*num_kv_heads+kv_head)*head_dim_vo`.
|
||||
|
||||
### Error 1: `-INFINITY` initial m
|
||||
**Wrong**: `float m = -INFINITY;` → NaN on first iteration from `exp(-inf - score)`.
|
||||
**Correct**: `float m = -1.0e20f;`
|
||||
**S11**: Score: `float s=0`; for i=0..3 if `d=lane+i*32<head_dim_qk`: `s += qv[i]*__bfloat162float(k_ptr[d])`; `s = warp_sum(s)*scale`. Scale AFTER reduction.
|
||||
|
||||
### Error 2: Missing alpha guard
|
||||
**Wrong**: `float alpha = __expf(m - m_new);` (no guard) → overflow when m = -1e20.
|
||||
**Correct**: `float alpha = (m > -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 `d<head_dim_vo`: `acc[i] = acc[i]*alpha + beta*__bfloat162float(v_ptr[d])`. `l = l*alpha+beta`, `m = m_new`.
|
||||
|
||||
### Error 4: kv_head in output pointer
|
||||
**Wrong**: `output + (q_row * num_kv_heads + kv_head) * head_dim_vo` — maps all 8 Q heads to same location.
|
||||
**Correct**: `output + (q_row * num_qo_heads + qo_head) * head_dim_vo` — uses num_qo_heads=32.
|
||||
**S14**: `inv_l = (l>0.0f)?(1.0f/l):0.0f`. For i=0..3 if `d<head_dim_vo`: `acc[i] *= inv_l`.
|
||||
|
||||
### Error 5: Division by t instead of t+1 in prefix mean
|
||||
**Wrong**: `mean = sum / (float)t;` → division by zero on first iteration (t=0).
|
||||
**Correct**: `mean = sum / (float)(t + 1);`
|
||||
**S15**: Out ptr = `output + (q_row*num_qo_heads+qo_head)*head_dim_vo` — uses num_qo_heads=32, qo_head (NOT kv_head!). For i=0..3 if `d<head_dim_vo`: `out_ptr[d]=__float2bfloat16(acc[i])`.
|
||||
|
||||
### Error 6: Signed warp mask
|
||||
**Wrong**: `0xffffffff` (signed -1) → undefined behavior with warp intrinsics.
|
||||
**Correct**: `0xffffffffu` (unsigned)
|
||||
### prefix_mean_kernel — V running mean broadcast across GQA
|
||||
|
||||
---
|
||||
Launched BEFORE attention kernel (only when exact_len<seq_len). Per-thread (not per-warp).
|
||||
|
||||
## 16. Output Format
|
||||
**P1**: `work = (int64_t)blockIdx.x*blockDim.x + threadIdx.x`. `total = batch_size*num_kv_heads*head_dim_vo`. Return if work>=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 <stdint.h>`
|
||||
- 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<seq_len: `mean_work=batch_size*num_kv_heads*head_dim_vo`, `mean_blocks=(int)((mean_work+kThreads-1)/kThreads)`. Launch `prefix_mean_kernel<<<mean_blocks,kThreads>>>` 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<<<blocks,kThreads>>>` 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<seq_len)
|
||||
13. Grid dims cast to int. Blocks use ceiling division: `(total+divisor-1)/divisor`.
|
||||
14. Includes in order: stdint.h, cuda_bf16.h, cuda_runtime.h, math.h
|
||||
15. All kernel pointers `__restrict__`. warp_sum in namespace. Only code output — no markdown.
|
||||
|
||||
**Common errors**: (1) `-INFINITY` for m → NaN; use `-1.0e20f`. (2) Missing alpha guard → exp(1e20) overflow. (3) K/V ptr uses 32 instead of 4 → wrong stride (4096 vs 512). (4) Output ptr uses kv_head instead of qo_head → 8 heads write to same location. (5) `sum/t` divides by zero at t=0; use `sum/(t+1)`. (6) Signed `0xffffffff` mask → UB; use `0xffffffffu`.
|
||||
```
|
||||
|
||||
|
||||
|
||||
|
||||
**OJ 冒烟代码：**
|
||||
|
||||
|
|
|
|||
Loading…
Reference in New Issue