cuda: DGX Spark / GB10 backend support — HBM-resident model

[TrevorS]

cuda: DGX Spark / GB10 backend support — HBM-resident model

Summary

Makes the 80 GB IQ2XXS DeepSeek V4 Flash checkpoint run well on NVIDIA DGX Spark (GB10, sm_121, 121 GiB UMA). On Spark the GPU can consume host-mmap'd weights via ATS, but at lower effective bandwidth than device-resident copies. This PR makes a budgeted set of hot tensor spans HBM-resident at startup, leaving the cold MoE routed experts ATS-mapped.

Standalone change: no MTP behavior, no kernel semantics change. Plain decode only. Adds GB10 to speed-bench/ next to m2_ultra.csv and m4_max.csv.

Speed — ds4-bench standard sweep (promessi_sposi.txt, gen=128, GB10)

ctx	upstream/main	this PR	Δ
2048	13.85	14.24	+0.39
8192	13.67	14.10	+0.43
16384	13.54	13.97	+0.43
18432	13.45	13.88	+0.43

Steady ~+0.4 t/s from the small-N kernel fuses (Q_A+KV_A pairing, hc_expand epilogue parallelization, head_rms_norm+rope_tail fusion).

UMA headroom (the bigger reason)

On current upstream/main, the startup cache copies the full 80.76 GiB model to device. On a 121 GiB Spark that leaves little headroom once KV cache and prefill activations grow: in my testing the standard ds4-bench 2048→65536 sweep did not complete past ~18k context on upstream.

This PR caps the cache at 24 GiB with a MoE filter, so only the hot spans (attention projections, shared experts, embedding, output head) become device-resident — ~8.2 GiB in practice — while cold routed experts stay ATS-mapped. The full 2048→65536 sweep completes with headroom to spare. (Worth confirming on your own Spark — UMA pressure depends on driver and allocator behavior.)

What's in the PR

• Startup HBM cache with a budget cap + MoE filter (hot spans device-resident, cold routed experts ATS-mapped).
• cuda_model_range_ptr helper: single hash-keyed lookup for device-resident pointers on the hot path.
• GPU argmax kernel (the prior fallback misused indexer scoring as argmax, double-paying dispatcher cost at N=1).
• Pair-fused Q_A + KV_A matmuls in qkv_rms_fused decode.
• Parallelized matmul_q8_0_hc_expand epilogue across n_hc lanes.
• HBM cache extends to the MTP support model (no behavioral change without --mtp).
• Drop cudaHostRegisterReadOnly (unsupported on GB10).
• speed-bench/gb10.csv from the standard sweep.

Tested

• make clean && make cuda-spark — clean
• make cpu — clean
• ./ds4_test --long-context, --tool-call-quality, --server, --metal-kernels — OK
• ./ds4-bench 2048→65536 sweep — completes; speed-bench/gb10.csv
• Two ds4_test checks fail identically on stock upstream/main (f91c12b) — not introduced here:
  • --logprob-vectors short_code_completion — the only divergence across all 4 steps is the case of the markdown code-fence language tag: the full-precision official API emits ```c, the IQ2XXS (2-bit) local model emits ```C; the generated code (return snprintf) is byte-identical. A near-tie the aggressive quant resolves to uppercase; reproduces on stock upstream/main.
  • --metal-tensor-equivalence — long-context only. Upstream's MoE routed-expert down-projection accumulates via float atomicAdd when n_tokens >= 128 (use_atomic_down); the order is scheduling-dependent, so two runs of the same config drift at ulp scale and occasionally flip a greedy argmax at long ctx. DS4_CUDA_MOE_NO_ATOMIC_DOWN=1 (an upstream flag) makes it bit-exact, confirming the cause. Stock upstream flakes identically (~2/5 runs); this is a plain-decode-only PR and doesn't touch that path.

Hardware: NVIDIA DGX Spark (GB10 / sm_121), driver 580.142, CUDA 13.0
Model: DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf

Note

Foundation for a stacked follow-up (#14) that adds MTP combined-forward speculative decode. This PR stands alone — it's what makes Spark usable at high context whether or not you care about MTP.