ViperEkura
/
AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

perf: post-multiply attention scale in float instead of pre-scaling Q in bf16 

- Replace bf16 pre-scale Q loading with direct 32-bit aligned bf16x2 reads
- Apply scale in float32 after Q@K^T, before online softmax
- Reduces causal max error from 2^-6 to 2^-8 with zero perf cost
This commit is contained in:
ViperEkura 2026-07-11 00:12:12 +08:00
parent d923ebe38d
commit 41cd40363a
2 changed files with 21 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
csrc/kernels/attn_decode_split_kv_mma.cuh
View File

@ -49,7 +49,6 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) {
__shared__ __align__(16) bf16 sV[BC * HEAD_DIM]; __shared__ __align__(16) bf16 sV[BC * HEAD_DIM];
__shared__ __align__(16) bf16 sQ[BR * HEAD_DIM]; __shared__ __align__(16) bf16 sQ[BR * HEAD_DIM];
bf16 scale_bf16 = __float2bfloat16(p.scale);
for (int i = lane; i < BR * HEAD_DIM; i += 32) { for (int i = lane; i < BR * HEAD_DIM; i += 32) {
int r = i / HEAD_DIM, d = i % HEAD_DIM; int r = i / HEAD_DIM, d = i % HEAD_DIM;
bf16 val = __float2bfloat16(0.0f); bf16 val = __float2bfloat16(0.0f);
@ -57,7 +56,7 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) {
int qh = q_head0 + r; int qh = q_head0 + r;
val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d]; val = p.q[(batch * p.q_head + qh) * HEAD_DIM + d];
} }
sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = __hmul(val, scale_bf16); sQ[r * LD + swiz_col(d, r, SWIZ_MASK)] = val;
} }
__syncwarp(); __syncwarp();
@ -116,6 +115,11 @@ __global__ void attn_decode_split_kv_mma_kernel(AttentionParams p) {
float Sacc[NC8][4]; float Sacc[NC8][4];
mma_compute_scores<KD, NC8>(Qa, sK, LD, SWIZ_MASK, lane, Sacc); mma_compute_scores<KD, NC8>(Qa, sK, LD, SWIZ_MASK, lane, Sacc);
#pragma unroll
for (int n8 = 0; n8 < NC8; n8++)
Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len; int maxc = p.is_causal ? min(p.kv_len, p.causal_offset + 1) : p.kv_len;
mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc, mma_softmax_tile<NC8, DN8>(kv0, maxc, maxc,
mask_base, p.mask, has_mask, mask_base, p.mask, has_mask,

26
csrc/kernels/attn_prefill_split_q_mma.cuh
View File

@ -81,16 +81,14 @@ void attn_prefill_split_q_mma_kernel(AttentionParams p) {
#pragma unroll #pragma unroll
for (int kt = 0; kt < KD; kt++) { for (int kt = 0; kt < KD; kt++) {
int c = kt * 16 + tid4 * 2; int c = kt * 16 + tid4 * 2;
const bf16* pa = &p.q[q_base + qra * HEAD_DIM + c]; const unsigned* pau = reinterpret_cast<const unsigned*>(
const bf16* pb = &p.q[q_base + qrb * HEAD_DIM + c]; &p.q[q_base + qra * HEAD_DIM + c]);
Qa[kt][0] = va ? pk2(__bfloat162float(pa[0]) * p.scale, const unsigned* pbu = reinterpret_cast<const unsigned*>(
__bfloat162float(pa[1]) * p.scale) : 0u; &p.q[q_base + qrb * HEAD_DIM + c]);
Qa[kt][1] = vb ? pk2(__bfloat162float(pb[0]) * p.scale, Qa[kt][0] = va ? pau[0] : 0u;
__bfloat162float(pb[1]) * p.scale) : 0u; Qa[kt][1] = vb ? pbu[0] : 0u;
Qa[kt][2] = va ? pk2(__bfloat162float(pa[8]) * p.scale, Qa[kt][2] = va ? pau[4] : 0u;
__bfloat162float(pa[9]) * p.scale) : 0u; Qa[kt][3] = vb ? pbu[4] : 0u;
Qa[kt][3] = vb ? pk2(__bfloat162float(pb[8]) * p.scale,
__bfloat162float(pb[9]) * p.scale) : 0u;
} }
float Oacc[DN8][4]; float Oacc[DN8][4];
@ -163,10 +161,16 @@ void attn_prefill_split_q_mma_kernel(AttentionParams p) {
// Warp-level causal skip // Warp-level causal skip
if (!use_skip || kv0 <= max_kv) { if (!use_skip || kv0 <= max_kv) {
// S = Q @ K^T + online softmax + O += P @ V (shared MMA functions) // S = Q @ K^T + scale + online softmax + O += P @ V
float Sacc[NC8][4]; float Sacc[NC8][4];
mma_compute_scores<KD, NC8>(Qa, bK, LD, SWIZ_MASK, lane, Sacc); mma_compute_scores<KD, NC8>(Qa, bK, LD, SWIZ_MASK, lane, Sacc);
// post-multiply scale in float (no bf16 precision loss from pre-scaling Q)
#pragma unroll
for (int n8 = 0; n8 < NC8; n8++)
Sacc[n8][0] *= p.scale, Sacc[n8][1] *= p.scale,
Sacc[n8][2] *= p.scale, Sacc[n8][3] *= p.scale;
int maxc0 = p.is_causal ? min(p.kv_len, qr0 + p.causal_offset + 1) int maxc0 = p.is_causal ? min(p.kv_len, qr0 + p.causal_offset + 1)
: p.kv_len; : p.kv_len;
int maxc1 = p.is_causal ? min(p.kv_len, qr1 + p.causal_offset + 1) int maxc1 = p.is_causal ? min(p.kv_len, qr1 + p.causal_offset + 1)