speculative: batched-GEMV kernel for verify path (Phase 24 step 1)

Add launch_gemv_bf16_batched: runs M m=1 GEMVs in a single 3D grid
launch (z = batch row) with numerically identical output to M sequential
launch_gemv_bf16 calls — same K-block partial accumulation, same
fixed-order reduction. Verified on dash5 with 10 prompts × 32 tokens:
matched=true, verify_decode_mismatches=0.

Expose as matmul_batched_gemv(a: [M,K], b: [K,N]) → [M,N] in
xserv-kernels. Replace the old matmul_rows_gemv helper in qwen3
forward_verify_paged_decode_attention; the per-row loop over matmul_2d +
concat_rows is replaced by a single matmul_batched_gemv call that
allocates the partials buffer in one shot and launches 2 kernels instead
of 2*M.

Current speedup_e2e is 0.47× (same ballpark as Phase 23 0.44×);
the batched launch saves ~3 ms overhead but this is small relative to
the total 28 ms spec cost. The path forward (per docs/24 §4) is
higher acceptance rate or cheaper draft, not further kernel optimization.
This commit is contained in:
2026-07-01 16:13:37 +08:00
parent 42e13f33dd
commit e5734b41fa
4 changed files with 155 additions and 20 deletions

View File

@@ -69,6 +69,62 @@ __global__ void gemv_reduce_to_bf16_kernel(
}
}
// Batched variant: M rows, same W. Grid.z = batch row index.
// Numerically identical to calling launch_gemv_bf16 M times in sequence because
// each z-slice executes the same accumulation order on the same data.
// partials buffer must be [M * num_k_blocks * N] floats.
__global__ void gemv_bf16_batched_partial_kernel(
const __nv_bfloat16* __restrict__ x, // [M, K]
const __nv_bfloat16* __restrict__ W, // [K, N]
float* __restrict__ partials, // [M, num_k_blocks, N]
int K, int N
) {
const int block_n = blockIdx.x;
const int block_k = blockIdx.y;
const int row = blockIdx.z;
const int t = threadIdx.x;
const int col = block_n * GEMV_TILE_N + t;
const int k_start = block_k * GEMV_TILE_K;
const int k_end = min(k_start + GEMV_TILE_K, K);
const int k_len = k_end - k_start;
__shared__ float x_shared[GEMV_TILE_K];
const __nv_bfloat16* x_row = x + (long long)row * K;
for (int i = t; i < k_len; i += GEMV_BLOCK) {
x_shared[i] = __bfloat162float(x_row[k_start + i]);
}
__syncthreads();
if (col >= N) return;
float sum = 0.0f;
for (int ki = 0; ki < k_len; ki++) {
sum += x_shared[ki] * __bfloat162float(W[(long long)(k_start + ki) * N + col]);
}
int num_k_blocks = (K + GEMV_TILE_K - 1) / GEMV_TILE_K;
partials[((long long)row * num_k_blocks + block_k) * N + col] = sum;
}
__global__ void gemv_batched_reduce_to_bf16_kernel(
const float* __restrict__ partials, // [M, num_k_blocks, N]
__nv_bfloat16* __restrict__ dst, // [M, N]
int n,
int num_k_blocks
) {
int col = blockIdx.x * blockDim.x + threadIdx.x;
int row = blockIdx.y;
if (col >= n) return;
float sum = 0.0f;
const float* row_partials = partials + (long long)row * num_k_blocks * n;
for (int kb = 0; kb < num_k_blocks; kb++) {
sum += row_partials[(long long)kb * n + col];
}
dst[(long long)row * n + col] = __float2bfloat16(sum);
}
extern "C" {
void launch_gemv_bf16(
@@ -104,4 +160,37 @@ void launch_gemv_bf16(
CUDA_CHECK_LAST_ERROR();
}
void launch_gemv_bf16_batched(
const void* x, // [M, K] BF16
const void* W, // [K, N] BF16
void* y_bf16, // [M, N] BF16
void* y_fp32_buf, // [M * num_k_blocks * N] FP32
int M, int K, int N,
void* stream
) {
cudaStream_t s = (cudaStream_t)stream;
int num_k_blocks = (K + GEMV_TILE_K - 1) / GEMV_TILE_K;
dim3 grid((N + GEMV_TILE_N - 1) / GEMV_TILE_N, num_k_blocks, M);
gemv_bf16_batched_partial_kernel<<<grid, GEMV_BLOCK, 0, s>>>(
(const __nv_bfloat16*)x,
(const __nv_bfloat16*)W,
(float*)y_fp32_buf,
K, N
);
CUDA_CHECK_LAST_ERROR();
int conv_block = 256;
int conv_grid_x = (N + conv_block - 1) / conv_block;
dim3 reduce_grid(conv_grid_x, M);
gemv_batched_reduce_to_bf16_kernel<<<reduce_grid, conv_block, 0, s>>>(
(const float*)y_fp32_buf,
(__nv_bfloat16*)y_bf16,
N,
num_k_blocks
);
CUDA_CHECK_LAST_ERROR();
}
} // extern "C"