[codex] cuda GH200 KV cache optimizations

[DominikBucko]

Summary

This PR adds CUDA-side KV cache storage options intended for large-context DS4 inference on Hopper/GH200-class GPUs.

• Add selectable CUDA KV cache storage via DS4_CUDA_KV_CACHE_TYPE=f32|f16|fp8 and legacy DS4_CUDA_KV_CACHE alias.
• Store raw, compressed attention, and index KV cache rows through cache-aware allocation/view helpers instead of assuming float row layout.
• Add f16 and fp8 cache load/store paths, including fp8 per-64-value scaling metadata.
• Update context memory estimates so they reflect the selected CUDA KV cache type.
• Use managed KV cache only for very large allocations where it is needed to start long-context runs.
• Fix the raw-only single-token decode path in the online mixed-attention kernels so layers without compressed KV attend the visible raw KV rows correctly.

Why

On a 96GB GH200, DS4 was limited by CUDA KV cache memory when pushing single-user long-context inference. The original CUDA path stored KV cache rows as f32 everywhere, which made large contexts unnecessarily expensive and prevented practical tests near the expected memory envelope.

The new f16 cache mode roughly halves KV cache memory while preserving coherent model behavior in the GH200 tests. fp8 mode is also available for experimentation, though it is currently slower/lower quality than f16 in my runs.

GH200 validation

Environment used for the main long-context profile:

DS4_CUDA_KV_CACHE_TYPE=f16 DS4_METAL_PREFILL_CHUNK=512 DS4_CUDA_NO_Q8_F16_CACHE=1

Observed on GH200 / sm_90:

• make cuda CUDA_ARCH=sm_90 completed.
• CUDA long-context regression passed: cuda long-context regression: OK.
• ds4-eval --questions 20 -n 4096 with f16 KV passed 18/20 questions.
• 32k context benchmark after the fix:
  • f16 KV: prefill ~1157 tok/s, generation ~27 tok/s, context buffers ~442 MiB.
  • fp8 KV: prefill ~860 tok/s, generation ~22 tok/s, context buffers ~163 MiB.
• f16 mode could start and serve with ctx=1800000 on the 96GB GH200, with context buffers reported around 13.28 GiB.

Notes

The env var name currently says CUDA because these paths are CUDA-specific. The surrounding function names still use the existing metal_graph_* naming convention from the repo.

[DominikBucko]

Closing in favor of a replacement PR with neutral branch and title naming.