gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

cuda: deterministic BF16 gemv + paged attention reductions

Browse Source 
BF16 greedy decode was sensitive to inter-block scheduling when logits
were close, which broke speculative-decoding verify-vs-decode parity.

- gemv.cu: write per-K-block partials, then reduce in fixed block order
  in a second kernel instead of atomicAdd across K-blocks. Scratch
  buffer size is now n * ceil(k / GEMV_TILE_K); gemv_scratch_elems()
  exposes this to callers, and decode_graph.rs sizes fp32_hidden/q/kv/
  intermediate/vocab from it.
- paged_attention.cu: replace atomicAdd merge of warp outputs with
  per-warp shared partials reduced in warp-id order for both the base
  and sinks kernels.
This commit is contained in:
Gahow Wang
2026-07-01 14:14:55 +08:00
parent 0314b4f3ac
commit 5b350ee5f0
4 changed files with 59 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
csrc/attention/paged_attention.cu
View File

@@ -118,7 +118,7 @@ __global__ void paged_decode_attention_bf16_kernel(
// ---- Block-level online softmax reduction ----
__shared__ float smem_max[32];
__shared__ float smem_sum[32];
__shared__ float smem_O[PAGED_HEAD_DIM_MAX];
__shared__ float smem_O_warp[32][PAGED_HEAD_DIM_MAX];
int lane = tid & 31;
int warp_id = tid >> 5;
@@ -164,8 +164,12 @@ __global__ void paged_decode_attention_bf16_kernel(
__syncthreads();
global_sum = smem_sum[0];
// Step 4: reduce O across block, dim by dim
for (int d = tid; d < head_dim; d += PAGED_THREADS) smem_O[d] = 0.0f;
// Step 4: reduce O across block, dim by dim. Store one partial per warp
// and sum in warp-id order; atomicAdd made greedy decode nondeterministic
// when logits were close.
for (int i = tid; i < 32 * PAGED_HEAD_DIM_MAX; i += PAGED_THREADS) {
reinterpret_cast<float*>(smem_O_warp)[i] = 0.0f;
}
__syncthreads();
for (int d = 0; d < head_dim; d++) {
@@ -173,13 +177,15 @@ __global__ void paged_decode_attention_bf16_kernel(
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_down_sync(0xffffffff, val, offset);
if (lane == 0) atomicAdd(&smem_O[d], val);
if (lane == 0) smem_O_warp[warp_id][d] = val;
}
__syncthreads();
float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
for (int d = tid; d < head_dim; d += PAGED_THREADS) {
O_ptr[d] = __float2bfloat16(smem_O[d] * inv_sum);
float out = 0.0f;
for (int i = 0; i < num_warps; i++) out += smem_O_warp[i][d];
O_ptr[d] = __float2bfloat16(out * inv_sum);
}
}
@@ -289,7 +295,7 @@ __global__ void paged_decode_attention_sinks_bf16_kernel(
// ---- Block-level online softmax reduction (same as base kernel) ----
__shared__ float smem_max[32];
__shared__ float smem_sum[32];
__shared__ float smem_O[PAGED_HEAD_DIM_MAX];
__shared__ float smem_O_warp[32][PAGED_HEAD_DIM_MAX];
int lane = tid & 31;
int warp_id = tid >> 5;
@@ -332,7 +338,9 @@ __global__ void paged_decode_attention_sinks_bf16_kernel(
__syncthreads();
global_sum = smem_sum[0];
for (int d = tid; d < head_dim; d += PAGED_THREADS) smem_O[d] = 0.0f;
for (int i = tid; i < 32 * PAGED_HEAD_DIM_MAX; i += PAGED_THREADS) {
reinterpret_cast<float*>(smem_O_warp)[i] = 0.0f;
}
__syncthreads();
for (int d = 0; d < head_dim; d++) {
@@ -340,13 +348,15 @@ __global__ void paged_decode_attention_sinks_bf16_kernel(
#pragma unroll
for (int offset = 16; offset > 0; offset >>= 1)
val += __shfl_down_sync(0xffffffff, val, offset);
if (lane == 0) atomicAdd(&smem_O[d], val);
if (lane == 0) smem_O_warp[warp_id][d] = val;
}
__syncthreads();
float inv_sum = (global_sum > 0.0f) ? (1.0f / global_sum) : 0.0f;
for (int d = tid; d < head_dim; d += PAGED_THREADS) {
O_ptr[d] = __float2bfloat16(smem_O[d] * inv_sum);
float out = 0.0f;
for (int i = 0; i < num_warps; i++) out += smem_O_warp[i][d];
O_ptr[d] = __float2bfloat16(out * inv_sum);
}
}