CUDA: prefill mixed scalar fallback has a fixed `scores[512]` and no `n_score` cap (ROCm / decode already check)
- Lenguaje dominante
- C
- Estrellas
- 22.3k
- Forks
- 2.1k
- Merge medio
- 1 d 3 h
- PR fusionados (30 d)
- 4
Descripción
## Summary
`attention_prefill_mixed_kernel` stores every attention score in `__shared__ float scores[512]`. `n_score = raw_count + visible_comp` is not capped, and the last hop of `attention_prefill_mixed_launch` does not call `cuda_attention_score_buffer_fits` or refuse the launch.
The same kernel on ROCm uses `DS4_ROCM_ATTENTION_PREFILL_MIXED_SCORE_CAP` (2048), returns in-kernel if `n_score` is above the cap, and the host refuses the launch with a log line. CUDA decode mixed attention uses `DS4_CUDA_ATTENTION_SCORE_CAP` (8192) plus `cuda_attention_score_buffer_fits` and an online fallback. CUDA prefill mixed last hop has neither.
```c
__shared__ float scores[512];
uint32_t n_score = raw_count + visible_comp;
/* writes scores[r] and scores[raw_count + c] with no cap */
```
`attention_prefill_mixed_launch` tries token-tile, then `attention_static_mixed_heads8_online_kernel` (`head_dim==512`, not `--quality`, `n_tokens>=128`), then cublas (`g_cublas_ready && head_dim==512 && n_tokens>1`), then the 512-slot kernel. `DS4_CUDA_NO_WINDOW_ATTENTION` and `DS4_CUDA_NO_CUBLAS_ATTENTION` skip the two fast paths and land on that last hop.
## Environment
- HEAD `84cc882`
- 2× RTX 2080 Ti (sm_75), Ubuntu, nvcc / compute-sanitizer 12.0, nsys 12.0
- `make cuda-generic`
## Repro
Zero-prefix Flash prefill on a ratio-4 layer with a first chunk of about 1600 tokens reaches `ds4_gpu_attention_prefill_static_mixed_heads_tensor` with `n_comp=400`, `window=128`, `head_dim=512`. Last token `n_score = 128+400 = 528 > 512`. With both fast paths disabled the launch is the 512-slot kernel:
```sh
DS4_CUDA_NO_WINDOW_ATTENTION=1 DS4_CUDA_NO_CUBLAS_ATTENTION=1 \
./ds4 --cuda -m ds4flash.gguf -p ""
```
Same numbers through the executor API the engine uses (`n_tokens=1600, n_comp=400, window=128, ratio=4, head_dim=512`, both env vars set):
- nsys: 100% of GPU time is `attention_prefill_mixed_kernel` (grid `1600 x 1`, block 256).
- `n_score=528` writes 16 floats past `scores[512]` into the same block's `__shared__ float partial[256]`. `cudaDeviceSynchronize` succeeds. compute-sanitizer memcheck treats the whole static shared region as one object and stays silent.
- To make the write leave the shared region (same kernel, larger `n_score`): `n_tokens=1024, n_comp=400, window=0, ratio=1, n_head=1, head_dim=8` → last token `n_score=1424`. compute-sanitizer `--tool memcheck` on 2026-08-14:
```
prefill mixed trigger: n_tokens=1024 n_comp=400 window=0 ratio=1
n_head=1 head_dim=8 last n_score=1424 (cap 512)
attention_prefill_static_mixed_heads rc=1
Invalid __shared__ write of size 4 bytes
at ds4_cuda.cu:7149:attention_prefill_mixed_kernel
Address 0xc10 is out of bounds
ERROR SUMMARY: 1164 errors
```
`:7149` is `scores[raw_count + c] = s`. Bare CUDA then fails `cudaDeviceSynchronize` with `unspecified launch failure`.
## Suggested fix
Before launching the 512-slot kernel, compute `n_score` the same way the kernel does. If it is above 512, refuse the launch or take `attention_static_mixed_heads8_online_kernel` when `head_dim==512`. Add `if (n_score > cap) return` inside the kernel, matching ROCm. Do the same for `ds4_gpu_attention_prefill_masked_mixed_heads_tensor`; it shares this launch.
Guía de contribución
Evaluación
Este issue todavía no se ha evaluado.