perf(ds4): elide intermediate prefill logits and tune gfx1151 mmq tiles

[cheese-cakee]

Summary

Improves DeepSeek4 prompt processing by removing unnecessary intermediate logits work and tuning the Strix Halo MMQ path.

The patch delivers matched current-main gains at chunk 512 on both supported Lucebox placements:

• +4.1% on R9700 gfx1201 + Strix Halo gfx1151
• +5.7% on Strix Halo gfx1151 alone

It keeps terminal logits and stateful execution on their existing correctness paths, while limiting MMQ tile changes to RDNA3.5/gfx1151.

Why this is useful

Intermediate prefill steps do not consume vocabulary logits. Before this change, those steps could still execute the full 150k x 4096 LM-head projection, output RMSNorm, and synchronous device-to-host logits readback.

This PR allows those intermediate steps to stay on a no-logits graph while preserving logits generation wherever the server actually needs it. The result is less unnecessary compute and synchronization during prompt processing, with measurable gains on both dual-GPU and Strix-only hardware.

Changes

• Skip the LM-head projection, output RMSNorm, and synchronous logits readback on intermediate prefill steps where logits are not required.
• Preserve terminal logits, snapshots, speculative-drafter feature extraction, decode, and verification on their required-output paths.
• Retain the dominant no-logits graph topology while executing terminal logits-needed variants transiently.
• Tune RDNA3.5/gfx1151 MMQ tile bounds to 48x64 with four warps.
• Leave RDNA3.0/gfx1100, RDNA4/gfx1201, and NVIDIA CUDA tile paths unchanged.
• Remain backward-compatible with single-token autoregressive decode and speculative verification.

Matched performance

Evidence label: maintainer-measured. The maintainer reran PR #633 on Lucebox4 against current main with the same ROCmFP2 model, hardware placement, environment, and chunk size on both sides of every A/B comparison.

Candidate head: a2878db6ec0d03a6252a6f8222169425e8c62298.

Hardware / placement	Chunk size	Current main	PR #633	Matched delta
R9700 gfx1201 + Strix Halo gfx1151	512	19.50 s, 107.4 tok/s	18.70 s, 112.0 tok/s	+4.1%
Strix Halo gfx1151 only	512	19.10 s, 109.6 tok/s	18.00 s, 116.3 tok/s	+5.7%
R9700 gfx1201 + Strix Halo gfx1151	2048	16.12 s, 129.9 tok/s	16.50 s, 126.9 tok/s	-2.3%

The matched chunk-512 results show that the optimization benefits both the heterogeneous Lucebox path and the standalone Strix Halo path.

At chunk 2048, absolute throughput remains higher than at chunk 512, but this PR is 2.3% slower than its matched current-main baseline. That wide-chunk interaction is a focused follow-up tuning opportunity

Correctness and compatibility

The optimization changes when intermediate logits are materialized; it does not intentionally change attention, HC residual streams, MoE routing, KV-cache updates, compressor pooling, or terminal sampling behavior.

The final logits-needed step remains explicit, and snapshot and speculative-drafter feature paths continue to request the outputs they consume.

Known limitation

• Dual-GPU chunk 2048 currently regresses by 2.3% against matched current main. The chunk-512 paths improve on both tested hardware placements.

Verification

• Full test_deepseek4_unit on Lucebox dual GPU, gfx1201 + gfx1151, ROCm 7.2.4: 53/53 passed.
• Full test_deepseek4_unit on CUDA SM 8.9: 53/53 passed.
• git diff --check: clean.
• GitHub CI is green, including the gfx1201 and gfx1151 ROCm GPU jobs.

[cubic-dev-ai]

All reported issues were addressed across 4 files

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

Re-trigger cubic</sub>

[cubic-dev-ai]

All reported issues were addressed across 3 files (changes from recent commits).

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

Re-trigger cubic</sub>

[cubic-dev-ai]

All reported issues were addressed across 2 files (changes from recent commits).

<sub>Tip: Review your code locally with the cubic CLI to iterate faster.

Re-trigger cubic</sub>

[Graffioh]

does this mean single-chunk prompts never populate the cache? eg --chunk 512, a 5–512 token prompt goes straight to the terminal/logits path, with logits_needed always true, so can_cache_dominant is false and every request rebuilds the layer-major graph

repeated single-chunk prompts with the same shape rebuild the layer-major graph on every request...was that case benchmarked, or is the build cost small enough that it does not matter?

extra: A snapshot boundary also sets need_logits to true because restoring that snapshot later requires the logits from that exact position.