Skip to content

Latest commit

 

History

History
175 lines (138 loc) · 7.45 KB

File metadata and controls

175 lines (138 loc) · 7.45 KB

GB10 (NVIDIA DGX Spark) CUDA decode notes

This fork extends antirez/ds4 with CUDA decode optimizations targeted at NVIDIA's GB10 / DGX Spark (sm_121, LPDDR5x). All changes are kernel-level inside ds4_cuda.cu; the upstream model, tokenizer, server, and CPU/Metal paths are untouched.

Result

Single-token decode at ctx 7047 on the fixed-imatrix DS4-Flash model:

State gen tok/s
Baseline (this fork's optimizations disabled) 13.66
With all this fork's optimizations enabled (default) 15.16
Δ +1.50 tps, -6.93 ms/token

3-run average per config, clean bench (no profilers). Correctness preserved: --logprob-vectors matches against the local fixture (tests/test-vectors/local.vec) and all source-hook tests pass.

Also in this fork: MTP α=0 correctness fix

Before any of the perf work, the GB10 session uncovered a silent correctness bug in the upstream CUDA backend: when both a base gguf and an --mtp gguf are loaded, the model-range fd-cache was resolving MTP weight offsets against the base model's host base, producing NaN activations and draft_accept_rate = 0.000 on every measurement. The MTP path appeared "working" — no errors — but every draft was rejected.

Fix: track the host base per fd in the CUDA model-range resolver, plumbed through cuda_model_range_ptr_from_fd and g_model_fd_host_base. After the fix the same bench yields draft_accept_rate ≈ 0.57 on DS4-Flash at draft=2, so MTP α can finally be measured on this backend.

Full repro, before/after numbers, and upstream-friendly patch summary in docs/mtp-fdcache-bug-report.md.

Performance levers

Two distinct kernel optimization patterns, landed across two commits:

F16 GEMV uint4 vectorization (7fed170)

The F16 pair / single GEMV kernels (matmul_f16_pair_ordered_chunks_kernel, matmul_f16_ordered_chunks_kernel) were doing scalar __half2float(wr[i]) reads with a shared-memory + serial-32 reduction. Both sat at ~44% of GB10's sequential-read bandwidth ceiling (~245 GB/s).

Added matmul_f16_pair_warp_vec8_kernel and matmul_f16_warp_vec8_kernel: uint4 (16-byte) weight loads = 8 halves per LSU op, __half22float2 unpack, warp_sum_f32 (warp-shuffle) reduction. Bit-equivalent up to a minor reduction-order difference within fp16-input rounding.

Achieved bandwidth on the compressor pair: 44% → 86% of sequential ceiling (~211 GB/s). The win compounds across every single-token F16 GEMV in the decode loop because the ordered_router dispatch branch covers them all.

This commit also ships an output_a/output_b/HC-expand sequential fuse (grouped_q8_0_a_preq_to_q8_kernel + ds4_gpu_attention_output_q8_fused_hc_tensor) that writes the low intermediate directly as Q8_0 layout, skipping one quantize_q8_0_f32 launch and the float low DRAM roundtrip.

q8_0 GEMV CTA-parallelism rewrite (1315d31)

The q8_0 pair and hc_expand kernels achieved only 131-132 GB/s on shared_gate_up / shared_down / attn_q_a-kv-pair shapes, despite the same kernel family achieving 192-213 GB/s on q_b and output_a. The gap is shape, not load width: those kernels' 8-output-per-CTA layout produces only 128-512 CTAs depending on shape, which is around or below GB10's ~288 concurrent-CTA capacity at this kernel's thread- bound occupancy (48 SMs × 1536 max threads/SM ÷ 256 threads/CTA = 6 CTAs/SM). The dp4a inner loop already reads 4 bytes at a time, and q8_0's 34-byte block layout (2-byte fp16 scale + 32 int8) breaks 16-byte vector alignment, so wider loads aren't the fix.

Added matmul_q8_0_pair_preq_warp1_quad_kernel<BLOCKS=128> and matmul_q8_0_hc_expand_preq_warp1_kernel<BLOCKS, WARPS_PER_OUT>: 1-output-per-CTA × N-warps-per-output (where N×32 = BLOCKS). Each lane reads exactly one block; cross-warp reduction via shared memory. CTA counts go from ~256-512 → 1024-4096 across affected shapes.

Naïve uint4 vectorization on q8_0 weights was tried and falsified (the 2-byte scale at offset 0 of every 34-byte block breaks alignment). A template-unroll-only variant of the pair kernel was tried and also falsified (null result, reverted).

How to repro

make ds4-bench ds4_test ds4

# Baseline = all this fork's optimizations disabled
DS4_CUDA_DISABLE_OUT_AB_FUSE=1 DS4_CUDA_NO_F16_PAIR_VEC8=1 \
  DS4_CUDA_NO_F16_VEC8=1 \
  DS4_CUDA_Q8_PAIR_NO_WARP1_QUAD=1 DS4_CUDA_Q8_HC_EXPAND_NO_WARP1=1 \
  ./ds4-bench --cuda -m ./ds4flash.gguf \
  --prompt-file bench/promessi_sposi.txt \
  --ctx-start 7047 --ctx-max 7047 --gen-tokens 128
# Expect ~13.65 tok/s

# Full optimizations on (defaults)
./ds4-bench --cuda -m ./ds4flash.gguf \
  --prompt-file bench/promessi_sposi.txt \
  --ctx-start 7047 --ctx-max 7047 --gen-tokens 128
# Expect ~15.16 tok/s

# Correctness gate (requires local.vec regenerated for your model)
DS4_TEST_MODEL=./ds4flash.gguf \
DS4_TEST_VECTOR_FILE=tests/test-vectors/local.vec \
  ./ds4_test --logprob-vectors

About local.vec: the fixture in this fork was generated against the fixed-imatrix model variant (DeepSeek-V4-Flash-IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-chat-v2-imatrix.gguf, revision b0c3326 on huggingface.co/antirez/deepseek-v4-gguf). If you ran download_model.sh q2 and got the default chat-v2.gguf (no -imatrix suffix), --logprob-vectors will fail not because the kernels are broken but because the model's top-1 logits differ at a few positions. Regenerate the fixture for your model with:

./tests/test-vectors/regen_local_vectors.py \
    -m ./ds4flash.gguf \
    -o tests/test-vectors/local.vec

The upstream cloud fixture (tests/test-vectors/official.vec) is the quality reference; local.vec is the code-regression gate.

The bench prompt (bench/promessi_sposi.txt) and source-hook gates (--bench-mtp-spec-source --bench-exact-replay-source --cuda-indexed-decode-heads8-source --decode-profile-source) are unchanged from upstream.

Disable flags

All optimizations are on by default and individually disablable:

Flag Disables
DS4_CUDA_DISABLE_OUT_AB_FUSE=1 output_a/b sequential fuse
DS4_CUDA_NO_F16_PAIR_VEC8=1 F16 pair vec8 kernel
DS4_CUDA_NO_F16_VEC8=1 F16 single vec8 kernel
DS4_CUDA_Q8_PAIR_NO_WARP1_QUAD=1 q8_0 pair warp1 kernel
DS4_CUDA_Q8_HC_EXPAND_NO_WARP1=1 q8_0 hc_expand warp1 kernel

Diagnostics

Several event-profile env flags are useful for substage attribution (all zero cost when off):

Flag Reports
DS4_CUDA_DECODE_EVENT_PROFILE=1 Top-level decode stage ms/token
DS4_CUDA_ATTENTION_EVENT_PROFILE=1 Attention substage breakdown (q_proj, kv_proj, compressor_indexer, attn_kernel, output_proj)
DS4_CUDA_INDEXER_EVENT_PROFILE=1 Compressor matmul, indexer scoring, indexer top-k splits
DS4_CUDA_MOE_EVENT_PROFILE=1 Aggregate MoE achieved GB/s/token

Further reading

Status

This is a personal fork; upstream PRs of the cleanest standalone pieces (F16 vec8 in particular) may be sent to antirez/ds4 in the future.