antirez / antirez/ds4

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

Offen
#803 1 Kommentar 0 Reaktionen 0 zugewiesene Personen Auf GitHub ansehen
Vorherrschende Sprache
C
Sterne
22.3k
Forks
2.1k
Ø Merge
1 T. 3 Std.
Gemergte PRs (30 T.)
4

Beschreibung

## 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.

Beitragsleitfaden

Beitragsleitfaden öffnen

Bewertung

Dieses Issue wurde noch nicht bewertet.

Neue Issues direkt in Ihr Postfach

Eine kurze Übersicht über anfängerfreundliche GitHub-Issues.