perf(ds4): compute greedy argmax on GPU (monolithic and layer split)

[Graffioh]

Summary

Compute the DeepSeek4 greedy argmax in the GGML GPU graph and return a single token id instead of the full n_vocab (129280) logits vector. Applies to the all-hot monolithic backend and to local/remote layer-split execution. Sampling and penalty requests keep the full-logits path; the hybrid MoE path keeps its CPU fallback.

What changed

• greedy prefill/decode emit one int32 argmax token instead of ~505 KiB of logits
• cached, dynamic, fused, and layer-major output graphs build only the requested logits-or-argmax path
• CPU/CUDA/HIP argmax tie-breaking unified to the lowest token id
• prefix-cache snapshots record their terminal representation (logits vs argmax); an incompatible restore falls back to fresh prefill

Performance

Reduces the greedy return payload from ~505 KiB (full n_vocab logits) to 4 bytes (one token id) per token. The saving is small on this hardware — a unified-memory copy or a local IPC transfer was never the bottleneck — and scales with the transport cost:

• monolithic HIP (Strix Halo, unified memory): greedy 24.8 vs sampled 23.2 tok/s. Greedy edges ahead, though that gap also reflects greedy skipping the sampler, not the argmax read alone.
• layer split, 3090 CUDA[0,10) + Strix Halo HIP[10,43): greedy 12.4 vs sampled 12.2 tok/s — within noise; per-token time is compute-bound on the remote shard and the local IPC copy is cheap.
• scales up for networked shards (eliminates ~505 KiB/token of boundary traffic) and faster-compute setups where the fixed transfer is a larger share of per-token time.

Correct and non-regressive in every configuration measured.

Follow-up

DSpark speculative decode does not benefit from this change: its greedy verify step reads back the full per-position logits (all_logits_out) to check each drafted token, so it never touches the single-output GPU argmax path. A future PR could extend the same idea to the verify — compute the per-position argmax on the GPU and return the drafted token ids instead of q × full-logits — alongside skipping the output projection on non-final prefill positions.

Validation

• unit tests (CPU + CUDA + HIP): output-path selection, argmax tie-breaking, snapshot compatibility + fresh-prefill fallback, sampler history
• GPU e2e: monolithic HIP greedy/sampled; layer-split CUDA+HIP greedy/sampled — valid tokens, no invalid argmax token
• fixed during GPU validation: layer-split greedy prefill dropped the argmax on the compressor-split batch path for non-zero-layer_begin shards (returned -1); the last chunk's argmax is now propagated

Rebase notes

Rebased onto main after DSpark spec decode (#496), indexed layer-major prefill (#520), and sampler-history fixes (#527/#536). Greedy DSpark seeds from the retained argmax token; layer-major prefill satisfies argmax-only requests via a host-side scan; fused verify graphs require an argmax tensor only when built to compute one.

[cubic-dev-ai]

1 issue found across 13 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="server/src/deepseek4/deepseek4_graph.cpp">

<violation number="1" location="server/src/deepseek4/deepseek4_graph.cpp:5961">
P2: Greedy layer-major prefill still transfers the full vocabulary logits buffer and computes argmax on CPU, so prompts with 5..`DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS` miss the advertised 4-byte GPU argmax path. Consider adding an argmax output mode to `ds4_try_layer_major_prefill` rather than using `layer_major_logits_scratch`.</violation>
</file>

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[cubic-dev-ai]

P2: Greedy layer-major prefill still transfers the full vocabulary logits buffer and computes argmax on CPU, so prompts with 5..DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS miss the advertised 4-byte GPU argmax path. Consider adding an argmax output mode to ds4_try_layer_major_prefill rather than using layer_major_logits_scratch.

Prompt for AI agents

Check if this issue is valid — if so, understand the root cause and fix it. At server/src/deepseek4/deepseek4_graph.cpp, line 5961:

<comment>Greedy layer-major prefill still transfers the full vocabulary logits buffer and computes argmax on CPU, so prompts with 5..`DS4_MAX_LAYER_MAJOR_PREFILL_TOKENS` miss the advertised 4-byte GPU argmax path. Consider adding an argmax output mode to `ds4_try_layer_major_prefill` rather than using `layer_major_logits_scratch`.</comment>

<file context>
@@ -5887,16 +5952,31 @@ bool deepseek4_step_layer_range(
-        is_last_shard && out_logits && ds4_backend_is_gpu(backend)) {
+        is_last_shard && (out_logits || out_argmax) &&
+        ds4_backend_is_gpu(backend)) {
+        std::vector<float> layer_major_logits_scratch;
+        std::vector<float> & layer_major_logits =
+            out_logits ? *out_logits : layer_major_logits_scratch;
</file context>

[Graffioh]

this will be done eventually in another PR to not clutter this one