cuda: reduce long-context visual-attention scratch

[JordiPosthumus]

Summary

Reuse the CUDA visual-attention score/output workspace across groups of up to
eight heads when the full score matrix exceeds 256 MiB. Smaller calls keep
their original cuBLAS batch size and unpack kernel. No model, sampling, attention
mask, KV format or server configuration changes.

Why

Visual prefill currently materializes heads × queries × (raw + compressed keys)
FP32 scores at once. A model can fit comfortably in memory and still request
tens of GiB of temporary storage when it encounters an image late in a conversation.

On GB10 with DeepSeek-V4-Flash-Vision-Exp IQ2XXS/w2-Q2_K and Q8 dense projections,
a 140,127-token request with an image after 136K tokens requested 18.13 GiB
for this workspace. A test-only allocation guard rejected it before it could
exhaust the host. The patched production-shaped build completed the identical
request with a 2.33 GiB maximum allocation and about 9 GiB host headroom.
The guard and watchdog exist only in the external diagnostic harness.

Each attention head is independent. This patch keeps each head's complete key
dimension, uses the existing score GEMM, softmax and value GEMM, and scatters
the finished group to its normal output positions before reusing scratch.
It does not truncate attention, change precision or introduce another cache.
The final group avoids a singleton because cuBLAS can select a different
arithmetic path for batchCount=1.

Compatibility

The full key dimension and original arithmetic are retained. Calls below the existing 256 MiB workspace threshold keep their prior dispatch.

Practical impact

The primary benefit is making long-context image requests fit by reducing temporary attention workspace. The earlier 18.13-to-2.33 GiB reproducer below demonstrates that capacity improvement; the 3.85–6.13% timings cover individual large attention calls. Neither is a whole-model speedup over f62ca29. Smaller calls retain their dispatch, and the standalone fixture checks the large-call arithmetic and memory boundaries.

Combined runtime context

A stock/full-stack engine comparison measures pristine f62ca29 against the complete integrated runtime on GB10 CUDA and M3 Ultra Metal, including full 262,144-token capacity. The original full sweeps used one process pair per machine and include additional changes beyond these 13 PRs. A targeted M3 follow-up did not reproduce the original short-decode slowdown, using a reversed-order pair, same-process controls and an exact original-executable replay. These results do not isolate this PR’s marginal gain, establish universal speed or quality, or newly benchmark its server/cache workflows. See the per-PR assessment for its supported benefit and limits.

Validation

The build and runtime checks below ran before upstream’s README-only update. The full PR diff and every other tracked file are byte-identical at this head; git diff --check passed again. The documentation-only rebase was not rebuilt.

One standalone commit, 78dd9deb5ea7, directly on upstream f62ca29a3087 (2026-09-07). The focused CUDA fixture checks independent double-precision references, wrapped raw KV, image boundaries, masks, incomplete head groups, scratch growth and output guards. Large/fractional/boundary modes retain complete-output comparison support.

On the preceding b6af0ad base, clean default Metal and CPU builds, make test-frontends, session-state, TP-command and vision-image tests, and git diff --check passed on M3 Ultra / Darwin 27.0 / Apple clang 21. These portability and host checks load no model or GPU fixture.

Earlier GB10/CUDA 13/sm_121a evidence: 128 full-output cases across two input families, both quality modes and cutoff/262K boundaries were byte-identical to the unchanged backend. Fractional boundary comparisons passed memcheck and synccheck. The scratch fixture fell from 584.56 MiB to 87.99 MiB, excluding initialization. A production-shaped Vision-Exp IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8 comparison matched 4,266,240 full-vocabulary floats after a 36K image prefill and 32 teacher-forced steps; four snapshot/suffix restores per arm also matched. The protected 140K reproducer and these model results predate this rebase. Balanced GPU-event call gains were 3.85–6.13% for the measured large attention shapes; no whole-model or decode gain is inferred.

Before the Metal-only upstream update, the same patch passed checks on NVIDIA DGX Spark / Ubuntu / Linux 6.17.0-1032-nvidia, GCC 13.3 and CUDA 13.0.88: clean make -j4 cuda-spark (sm_121a), make CUDA_ARCH=sm_121a test-frontends, session-state, TP-command and vision-image tests, and a clean make -j4 cpu build. These are host tests and compile/link coverage; CUDA execution is recorded separately. The later upstream updates are a five-line GLM SSD change inside a Metal-only guard and a README edit. CUDA source and the PR patch are unchanged; native Linux preprocessing of ds4.c was byte-identical for CUDA, CPU and sanitizer builds.

The standalone visual-attention fixture passed its default, fractional, boundary and memory modes. The default mode also passed Compute Sanitizer memcheck and synccheck with zero errors. These are independent-reference and guard checks; the earlier complete-output baseline comparisons above remain separately dated evidence.

No full aggregate model-suite, other-GPU, other-quantization or multi-GPU result is claimed.