diff --git a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md index 4af322b..bda6723 100644 --- a/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md +++ b/基于AI Agent开发范式的国产GPU大模型推理算子库优化/FlashInfer Benchmark实战：从性能基线到XPU-OJ冒烟提交.md @@ -663,376 +663,378 @@ FlashInfer 方向包含 **4 个可选算子题目**，每个对应独立的 benc **参考 prompt** 及 **20001 FlashInfer Ragged Prefill 参考冒烟代码**： - ```plaintext - # FlashInfer Ragged Prefill CUDA Kernel — Problem 20001 + - 参考 prompt + ```plaintext + # FlashInfer Ragged Prefill CUDA Kernel — Problem 20001 - ## 1. Problem Spec + ## 1. Problem Spec - Implement FlashInfer `BatchPrefillWithRaggedKVCacheWrapper` forward pass. Ragged NHD layout, GQA, causal masking. + Implement FlashInfer `BatchPrefillWithRaggedKVCacheWrapper` forward pass. Ragged NHD layout, GQA, causal masking. - **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. + **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. Interface + ## 2. Interface + ```cpp + #include + #include + + extern "C" void run_kernel( + 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 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`. + + ## 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. + + **Approach**: exact attention for first 1024 positions, prefix-mean approximation for the tail. + + **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. + + **Two-kernel architecture** (order matters): + 1. `prefix_mean_kernel` FIRST (only when exact_len0.01. + + ## 4. Rules + + - **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`. + + ## 5. Signatures + + ```cpp + __device__ __forceinline__ float warp_sum(float x) + ``` + + ```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) + ``` + + ```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, + 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) + ``` + + ## 6. warp_sum + + 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. + + ## 7. Kernels + + ### ragged_prefill_smoke_kernel — exact attention, warp-per-query, register-only + + 128 threads/block = 4 warps. Each warp handles one (batch, q_pos, qo_head). + + **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. + + **S3**: Decompose (in order): `qo_head = work % num_qo_heads`, `work /= num_qo_heads`, `q_pos = work % exact_len`, `batch = work / exact_len`. + + **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. + + **S5**: `visible = (causal) ? kv_len - qo_len + q_pos + 1 : kv_len`. Clamp to [0, kv_len]. + + **S6**: `group = num_qo_heads/num_kv_heads` (=8), `kv_head = qo_head/group`, `q_row = qo_begin+q_pos`. + + **S7**: `scale = rsqrtf((float)head_dim_qk)`. + + **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` — guard prevents exp(1e20). `beta = __expf(s-m_new)`. + + **S13**: For i=0..3 if `d0.0f)?(1.0f/l):0.0f`. For i=0..3 if `d=total. + + **P2**: `d = work%head_dim_vo`, `work/=head_dim_vo`, `kv_head = work%num_kv_heads`, `batch = work/num_kv_heads`. + + **P3**: `group = num_qo_heads/num_kv_heads` (=8). `qo_begin = qo_indptr[batch]`, `kv_begin = kv_indptr[batch]`. + + **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 + #include + #include - extern "C" void run_kernel( - 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 - ``` + #include - 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`. + namespace { - ## 3. Strategy + __device__ __forceinline__ float warp_sum(float x) { + for (int offset = 16; offset > 0; offset >>= 1) { + x += __shfl_down_sync(0xffffffffu, x, offset); + } + return __shfl_sync(0xffffffffu, x, 0); + } - 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. - - **Approach**: exact attention for first 1024 positions, prefix-mean approximation for the tail. - - **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. - - **Two-kernel architecture** (order matters): - 1. `prefix_mean_kernel` FIRST (only when exact_len0.01. - - ## 4. Rules - - - **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`. - - ## 5. Signatures - - ```cpp - __device__ __forceinline__ float warp_sum(float x) - ``` - - ```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) { + const int lane = threadIdx.x & 31; + const int warp_id = threadIdx.x >> 5; + const int warps_per_block = blockDim.x >> 5; - ```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, - int64_t head_dim_vo) - ``` + int64_t work = static_cast(blockIdx.x) * warps_per_block + warp_id; + const int64_t total = batch_size * exact_len * num_qo_heads; + if (work >= total) return; - ```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 int64_t qo_head = work % num_qo_heads; + work /= num_qo_heads; + const int64_t q_pos = work % exact_len; + const int64_t batch = work / exact_len; - ## 6. warp_sum + const int64_t qo_begin = qo_indptr[batch]; + const int64_t qo_len = qo_indptr[batch + 1] - qo_begin; + if (q_pos >= qo_len) return; - 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. + const int64_t kv_begin = kv_indptr[batch]; + const int64_t kv_len = kv_indptr[batch + 1] - kv_begin; + int64_t visible = kv_len; + if (causal) { + visible = kv_len - qo_len + q_pos + 1; + if (visible < 0) visible = 0; + if (visible > kv_len) visible = kv_len; + } - ## 7. Kernels + const int64_t group = num_qo_heads / num_kv_heads; + const int64_t kv_head = qo_head / group; + const int64_t q_row = qo_begin + q_pos; + const float scale = rsqrtf(static_cast(head_dim_qk)); - ### ragged_prefill_smoke_kernel — exact attention, warp-per-query, register-only - - 128 threads/block = 4 warps. Each warp handles one (batch, q_pos, qo_head). - - **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. - - **S3**: Decompose (in order): `qo_head = work % num_qo_heads`, `work /= num_qo_heads`, `q_pos = work % exact_len`, `batch = work / exact_len`. - - **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. - - **S5**: `visible = (causal) ? kv_len - qo_len + q_pos + 1 : kv_len`. Clamp to [0, kv_len]. - - **S6**: `group = num_qo_heads/num_kv_heads` (=8), `kv_head = qo_head/group`, `q_row = qo_begin+q_pos`. - - **S7**: `scale = rsqrtf((float)head_dim_qk)`. - - **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` — guard prevents exp(1e20). `beta = __expf(s-m_new)`. - - **S13**: For i=0..3 if `d0.0f)?(1.0f/l):0.0f`. For i=0..3 if `d=total. - - **P2**: `d = work%head_dim_vo`, `work/=head_dim_vo`, `kv_head = work%num_kv_heads`, `batch = work/num_kv_heads`. - - **P3**: `group = num_qo_heads/num_kv_heads` (=8). `qo_begin = qo_indptr[batch]`, `kv_begin = kv_indptr[batch]`. - - **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 - - #include - #include - - #include - - namespace { - - __device__ __forceinline__ float warp_sum(float x) { - for (int offset = 16; offset > 0; offset >>= 1) { - x += __shfl_down_sync(0xffffffffu, x, offset); - } - return __shfl_sync(0xffffffffu, x, 0); - } - - __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 int lane = threadIdx.x & 31; - const int warp_id = threadIdx.x >> 5; - const int warps_per_block = blockDim.x >> 5; - - int64_t work = static_cast(blockIdx.x) * warps_per_block + warp_id; - const int64_t total = batch_size * exact_len * num_qo_heads; - if (work >= total) return; - - const int64_t qo_head = work % num_qo_heads; - work /= num_qo_heads; - const int64_t q_pos = work % exact_len; - const int64_t batch = work / exact_len; - - const int64_t qo_begin = qo_indptr[batch]; - const int64_t qo_len = qo_indptr[batch + 1] - qo_begin; - if (q_pos >= qo_len) return; - - const int64_t kv_begin = kv_indptr[batch]; - const int64_t kv_len = kv_indptr[batch + 1] - kv_begin; - int64_t visible = kv_len; - if (causal) { - visible = kv_len - qo_len + q_pos + 1; - if (visible < 0) visible = 0; - if (visible > kv_len) visible = kv_len; - } - - const int64_t group = num_qo_heads / num_kv_heads; - const int64_t kv_head = qo_head / group; - const int64_t q_row = qo_begin + q_pos; - const float scale = rsqrtf(static_cast(head_dim_qk)); - - const __nv_bfloat16* q_ptr = q + (q_row * num_qo_heads + qo_head) * head_dim_qk; - float qv[4]; - float acc[4]; - for (int i = 0; i < 4; ++i) { - const int d = lane + i * 32; - qv[i] = (d < head_dim_qk) ? __bfloat162float(q_ptr[d]) : 0.0f; - acc[i] = 0.0f; - } - - float m = -1.0e20f; - float l = 0.0f; - for (int64_t kv_pos = 0; kv_pos < visible; ++kv_pos) { - const int64_t kv_row = kv_begin + kv_pos; - const __nv_bfloat16* k_ptr = k + (kv_row * num_kv_heads + kv_head) * head_dim_qk; - const __nv_bfloat16* v_ptr = v + (kv_row * num_kv_heads + kv_head) * head_dim_vo; - - float score = 0.0f; + const __nv_bfloat16* q_ptr = q + (q_row * num_qo_heads + qo_head) * head_dim_qk; + float qv[4]; + float acc[4]; for (int i = 0; i < 4; ++i) { const int d = lane + i * 32; - if (d < head_dim_qk) { - score += qv[i] * __bfloat162float(k_ptr[d]); - } + qv[i] = (d < head_dim_qk) ? __bfloat162float(q_ptr[d]) : 0.0f; + acc[i] = 0.0f; } - score = warp_sum(score) * scale; - const float m_new = fmaxf(m, score); - const float alpha = (m > -1.0e19f) ? __expf(m - m_new) : 0.0f; - const float beta = __expf(score - m_new); + float m = -1.0e20f; + float l = 0.0f; + for (int64_t kv_pos = 0; kv_pos < visible; ++kv_pos) { + const int64_t kv_row = kv_begin + kv_pos; + const __nv_bfloat16* k_ptr = k + (kv_row * num_kv_heads + kv_head) * head_dim_qk; + const __nv_bfloat16* v_ptr = v + (kv_row * num_kv_heads + kv_head) * head_dim_vo; + float score = 0.0f; + for (int i = 0; i < 4; ++i) { + const int d = lane + i * 32; + if (d < head_dim_qk) { + score += qv[i] * __bfloat162float(k_ptr[d]); + } + } + score = warp_sum(score) * scale; + + const float m_new = fmaxf(m, score); + const float alpha = (m > -1.0e19f) ? __expf(m - m_new) : 0.0f; + const float beta = __expf(score - m_new); + + for (int i = 0; i < 4; ++i) { + const int d = lane + i * 32; + if (d < head_dim_vo) { + acc[i] = acc[i] * alpha + beta * __bfloat162float(v_ptr[d]); + } + } + l = l * alpha + beta; + m = m_new; + } + + __nv_bfloat16* out_ptr = output + (q_row * num_qo_heads + qo_head) * head_dim_vo; + const float inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f; for (int i = 0; i < 4; ++i) { const int d = lane + i * 32; if (d < head_dim_vo) { - acc[i] = acc[i] * alpha + beta * __bfloat162float(v_ptr[d]); + out_ptr[d] = __float2bfloat16(acc[i] * inv_l); } } - l = l * alpha + beta; - m = m_new; } - __nv_bfloat16* out_ptr = output + (q_row * num_qo_heads + qo_head) * head_dim_vo; - const float inv_l = (l > 0.0f) ? (1.0f / l) : 0.0f; - for (int i = 0; i < 4; ++i) { - const int d = lane + i * 32; - if (d < head_dim_vo) { - out_ptr[d] = __float2bfloat16(acc[i] * inv_l); + __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, + int64_t head_dim_vo) { + int64_t work = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t total = batch_size * num_kv_heads * head_dim_vo; + if (work >= total) return; + + const int64_t d = work % head_dim_vo; + work /= head_dim_vo; + const int64_t kv_head = work % num_kv_heads; + const int64_t batch = work / num_kv_heads; + const int64_t group = num_qo_heads / num_kv_heads; + const int64_t qo_begin = qo_indptr[batch]; + const int64_t kv_begin = kv_indptr[batch]; + + float sum = 0.0f; + for (int64_t t = 0; t < seq_len; ++t) { + const int64_t kv_row = kv_begin + t; + sum += __bfloat162float(v[(kv_row * num_kv_heads + kv_head) * head_dim_vo + d]); + const __nv_bfloat16 mean = __float2bfloat16(sum / static_cast(t + 1)); + const int64_t out_row = qo_begin + t; + for (int64_t g = 0; g < group; ++g) { + const int64_t qo_head = kv_head * group + g; + output[(out_row * num_qo_heads + qo_head) * head_dim_vo + d] = mean; + } } } - } - __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, - int64_t head_dim_vo) { - int64_t work = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - const int64_t total = batch_size * num_kv_heads * head_dim_vo; - if (work >= total) return; + } // namespace - const int64_t d = work % head_dim_vo; - work /= head_dim_vo; - const int64_t kv_head = work % num_kv_heads; - const int64_t batch = work / num_kv_heads; - const int64_t group = num_qo_heads / num_kv_heads; - const int64_t qo_begin = qo_indptr[batch]; - const int64_t kv_begin = kv_indptr[batch]; + 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) { + constexpr int kThreads = 128; + constexpr int kWarpsPerBlock = kThreads / 32; - float sum = 0.0f; - for (int64_t t = 0; t < seq_len; ++t) { - const int64_t kv_row = kv_begin + t; - sum += __bfloat162float(v[(kv_row * num_kv_heads + kv_head) * head_dim_vo + d]); - const __nv_bfloat16 mean = __float2bfloat16(sum / static_cast(t + 1)); - const int64_t out_row = qo_begin + t; - for (int64_t g = 0; g < group; ++g) { - const int64_t qo_head = kv_head * group + g; - output[(out_row * num_qo_heads + qo_head) * head_dim_vo + d] = mean; + int64_t exact_len = seq_len; + if ((batch_size >= 4 && seq_len >= 16384) || (batch_size >= 16 && seq_len >= 8192)) { + exact_len = 1024; + const int64_t mean_work = batch_size * num_kv_heads * head_dim_vo; + const int mean_blocks = static_cast((mean_work + kThreads - 1) / kThreads); + prefix_mean_kernel<<>>( + v, output, qo_indptr, kv_indptr, batch_size, seq_len, + num_qo_heads, num_kv_heads, head_dim_vo); } + + const int64_t total = batch_size * exact_len * num_qo_heads; + const int blocks = static_cast((total + kWarpsPerBlock - 1) / kWarpsPerBlock); + ragged_prefill_smoke_kernel<<>>( + 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); } - } - - } // namespace - - 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) { - constexpr int kThreads = 128; - constexpr int kWarpsPerBlock = kThreads / 32; - - int64_t exact_len = seq_len; - if ((batch_size >= 4 && seq_len >= 16384) || (batch_size >= 16 && seq_len >= 8192)) { - exact_len = 1024; - const int64_t mean_work = batch_size * num_kv_heads * head_dim_vo; - const int mean_blocks = static_cast((mean_work + kThreads - 1) / kThreads); - prefix_mean_kernel<<>>( - v, output, qo_indptr, kv_indptr, batch_size, seq_len, - num_qo_heads, num_kv_heads, head_dim_vo); - } - - const int64_t total = batch_size * exact_len * num_qo_heads; - const int blocks = static_cast((total + kWarpsPerBlock - 1) / kWarpsPerBlock); - ragged_prefill_smoke_kernel<<>>( - 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); - } - ``` + ``` 更多 OpenCode 使用教程和 Agent 使用技巧可见 - [*OpenCode 官方文档：简介*](https://opencode.ai/docs/zh-cn/)