antirez / antirez/ds4

CUDA: prefill mixed scalar fallback has a fixed `scores[512]` and no `n_score` cap (ROCm / decode already check)

Abierto
#803 1 comentario 0 reacciones 0 asignados Ver en GitHub
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

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.