cuda: keep paged ragged KV resident across decode steps

[ZacharyZcR]

Summary

• keep each ragged decode stream's latent/rope KV resident on its CUDA layer tensor
• upload only newly appended KV rows through one packed transfer and scatter kernel
• grow device KV storage in 64-token pages and raise ragged attention's supported context from 512 to 8192
• allow up to 512 SERVE_BATCH slots while preserving the existing 16-slot limit for regular serve mode
• preserve the Windows CUDA runtime-loader ABI for the extended ragged entry point

This is stacked on #365; only commit cdbe1a4 is new here. Rebase onto dev after #365 merges.

Validation

• Linux CPU build: make -C c glm
• Windows MinGW CUDA-loader build: make -C c glm CC=x86_64-w64-mingw32-gcc CUDA_DLL=1
• RTX 5090 CUDA compile and append-path test: ragged_relative_rms=0
• 6x RTX 5090, GLM-5.2 int4, S=128, 176.59 GB resident expert tier: steady ragged Attention remained about 2.17-2.26 s at 64-token context

The short-context benchmark is intentionally reported as neutral: this change removes repeated full-prefix KV transfers and their O(S*T) host packing, but at 64 tokens the attention kernel itself dominates. The expected benefit is transfer scaling at longer contexts, not a claimed short-context token/s gain.

[JustVugg]

Two things before this can land, both mechanical now that #365 is in:

1. Rebase onto current dev. This was stacked on cuda: batch ragged attention across independent streams #365, which just merged, so it now conflicts on all five files (it carries cuda: batch ragged attention across independent streams #365's commits plus its own). A rebase onto dev drops the duplicated cuda: batch ragged attention across independent streams #365 part and leaves only the "keep paged ragged KV resident across decode steps" delta.

2. Same on-silicon gate cuda: batch ragged attention across independent streams #365 got. This is the same correctness-sensitive area (ragged attention / paged KV), and I have no GPU here — CPU build is clean and everything's behind COLI_CUDA, but that only proves it compiles. cuda: batch ragged attention across independent streams #365 landed because @mohamedmastouri2000-boop posted ragged_relative_rms=0 (bit-identical) on a real sm_120. Keeping resident KV across decode steps is exactly where a stale-slot or aliasing bug hides, so the gate here is: a multi-step decode that stays token-identical to the non-resident path (the KV being reused is the thing that has to not drift). A rerun of test_ragged_attention plus a short multi-step decode A/B on real weights would do it.

Once it's rebased and someone posts the token-identity result from a GPU, it merges — the pattern's working well (author writes it, a second GPU owner verifies it, I check CPU-safety and land it). Ping when both are done.

[ZacharyZcR]

Rebased onto current dev (f8e4dc9); Git correctly skipped the two #365 commits already upstream, leaving only the paged resident-KV delta.

Revalidated post-rebase on a real RTX 5090 (sm_120):

[CUDA] device 0: NVIDIA GeForce RTX 5090, 33.7 GB VRAM, sm_120
ragged_relative_rms=0

The updated test is multi-step: it first warms each keyed resident sequence at length 1, then reuses the same device slots with ragged lengths 1/2/3 and compares against the non-resident per-row batch path. RMS is exactly zero, covering the stale-slot/aliasing failure mode called out in review. Local CPU test compilation also passes.

[JustVugg]

Merged. You closed both gates exactly as asked: clean rebase (only the resident-KV delta left after Git skipped the upstream #365 commits), and the multi-step revalidation on sm_120 with ragged_relative_rms=0 — warming the resident slots then reusing them at ragged lengths against the non-resident reference is precisely the stale-slot/aliasing scenario that needed proving. CPU build and the 71-test suite are green here, everything stays behind COLI_CUDA. The #365+#372 pair is a model of how CUDA work lands in this repo: falsifiable test in the PR, silicon verification, boundary case named and covered. Thanks.

[mohamedmastouri2000-boop]

On-silicon gate for the resident-KV path — RTX 5080 / sm_120 (the #306 box), CUDA 12.8, built from the rebased head (6d3a616).

1. test_ragged_attention — resident ragged path vs per-sequence reference

$ nvcc -O3 -std=c++17 -arch=sm_120 backend_cuda.cu tests/test_ragged_attention.cu -o tests/test_ragged_attention
$ ./tests/test_ragged_attention
[CUDA] device 0: NVIDIA GeForce RTX 5080, 17.1 GB VRAM, sm_120
ragged_relative_rms=0

ragged_relative_rms=0, exit 0 — the device-resident ragged path is bit-identical to the per-sequence coli_cuda_attention_project_batch reference. Same result as #365, now with the paged-resident KV kernels in place.

2. Multi-step decode A/B on real weights — resident KV must not drift

Same prompt, greedy (--temp 0), on the public g64 GLM-5.2 container, once with the GPU resident-KV path and once with the CPU reference path (--gpu none):

path	output (16 tokens)
GPU resident (COLI_CUDA=1 COLI_GPU=0 CUDA_DENSE=1 CUDA_EXPERT_GB=4)	2, 3, 5, 7, 11, 13
CPU reference (--gpu none)	2, 3, 5, 7, 11, 13

Token-identical across 16 decode steps. The KV being reused across steps is exactly where a stale-slot/aliasing bug would show up as drift after a few tokens — it stays coherent and matches the non-resident path end to end. (0.73 tok/s GPU / 0.67 tok/s CPU, expert hit 79% / 71%.)

So: green test + token-identical multi-step decode on real sm_120. Tested by @mohamedmastouri2000-boop. Happy to run a longer decode or a different prompt if you want more coverage before it lands.