kernels: fix sink-decode build (b4db953 didn't compile)
b4db953 committed before a green build: the existing
launch_decode_attention_bf16 call wasn't updated for the new kernel
signature (too-few-args at flash_attention.cu:433) and the Rust FFI
binding + decode_attention_sink wrapper never landed (failed string
matches). This adds the missing pieces; xserv now builds green on dash5
(release, 17s). qwen3 decode path unchanged (passes nullptr/0).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -429,7 +429,30 @@ void launch_decode_attention_bf16(
|
||||
(__nv_bfloat16*)O,
|
||||
num_q_heads, num_kv_heads,
|
||||
kv_len, head_dim,
|
||||
scale
|
||||
scale,
|
||||
nullptr, 0 // no sinks, no sliding window (qwen3 path unchanged)
|
||||
);
|
||||
CUDA_CHECK_LAST_ERROR();
|
||||
}
|
||||
|
||||
// gpt-oss decode: per-head attention sinks + optional sliding window.
|
||||
void launch_decode_attention_sink_bf16(
|
||||
const void* Q, const void* K, const void* V, void* O,
|
||||
int batch, int num_q_heads, int num_kv_heads,
|
||||
int kv_len, int head_dim,
|
||||
float scale, const void* sinks, int window, void* stream
|
||||
) {
|
||||
int grid = batch * num_q_heads;
|
||||
int block = DECODE_THREADS;
|
||||
decode_attention_bf16_kernel<<<grid, block, 0, (cudaStream_t)stream>>>(
|
||||
(const __nv_bfloat16*)Q,
|
||||
(const __nv_bfloat16*)K,
|
||||
(const __nv_bfloat16*)V,
|
||||
(__nv_bfloat16*)O,
|
||||
num_q_heads, num_kv_heads,
|
||||
kv_len, head_dim,
|
||||
scale,
|
||||
(const float*)sinks, window
|
||||
);
|
||||
CUDA_CHECK_LAST_ERROR();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user