The tile8 gate/up kernel is latency-bound, not format-bound: 57 KiB of
staged-activation LDS pins residency to 1-2 blocks/WGP and each block
has only 8 warps to hide the serial block-loop latency. The down
kernel has the opposite shape - ~917k tiny blocks per layer-chunk at
4096 tokens, each re-staging the same 16.6 KiB of activations.
Three variants, all preserving the exact per-(token,row) accumulation
order (same 8-lane q8_K block split, same dp4a helpers, same
quarter-warp reduction, same epilogue - bit-exact by construction):
- DS4_ROCM_ENABLE_MXFP4_TILE4: 4-token expert tiles for gate/up
(templated tileN<TILE_TOKENS, QWARPS> kernel) - 29 KiB LDS at 7168
dims, 4 blocks/WGP.
- DS4_ROCM_ENABLE_MXFP4_ROW64: 512-thread gate/up blocks (64 rows x
8 tokens against one staged copy) - 2 blocks/WGP = 32 resident warps.
- DS4_ROCM_MXFP4_DOWN_RGROUP=1..8: down kernel loops over consecutive
32-row groups against one staged activation copy - staging traffic
and block count divide by R.
tests/test_mxfp4_rocm now snapshots the default path's out tensor at
n=128/512 and memcmp-asserts each variant against it (cross-path
bitwise equality, stronger than the tolerance checks).
Compile-verified on gfx1151; hardware A/B and the acceptance gate run
in the next test window. The WMMA rewrite is deferred until these
say whether occupancy or format is the binding constraint.
Production kernel trace of the 10.5k-token prefill (rocprofv3 on the
coordinator) showed the MXFP4 routed-MoE tile kernels dominate prefill:
gate/up 247.7 ms and down 174.9 ms per layer-chunk, bound not by
compute (~0.6% of fp32 peak) but by expert-weight DRAM re-reads - every
8-token tile re-reads its experts' full gate+up rows.
Changes (all bit-exact; per-element accumulation order untouched):
- gate/up tile8 kernel grid swapped to (tile, row-block): consecutive
blocks re-walk the same 32-row weight slice across an expert's token
tiles, so the slice stays L2-resident. Production: 247.7 -> 222.3 ms
per layer-chunk (-10.3%), 10.5k prefill 61.5 s -> 59.9 s.
- A-activation LDS staging threshold raised 16 -> 28 q8_K blocks (the
7168-dim case; 56.9 KiB dynamic smem with the >48 KiB opt-in).
Neutral in production (L2 was already absorbing the re-reads), large
in the isolated bench; staged and unstaged dot helpers are bit-equal
(integer bsum order is exact; verified in-tree and at 7168 dims).
- down-kernel grid swap measured NEGATIVE (174.9 -> 184.8 ms) and was
reverted.
- Two experimental kernels parked default-off behind env gates with
honest numbers: tile32 B-register reuse (DS4_ROCM_ENABLE_MXFP4_TILE32;
register spills, 2155 ms bench) and row8-ldsB B-staging
(DS4_ROCM_ENABLE_MXFP4_LDSB; 154 s production - scattered A reads).
- tests/bench_mxfp4_rocm: production-shape MoE bench harness. NOTE:
it is a poor proxy for production kernel time (isolated process
memory/L2 behavior differs ~5x); production traces are the arbiter.
- DS4_ROCM_MOE_PATH_DEBUG: launch-path diagnostics.
Validation: tests/test_mxfp4_rocm PASS (default dims + 7168/top-8
variant), acceptance gate df07199e5a292872 holds, decode 15.38 t/s,
10.5k probe 59.9 s (176 t/s avg).
The synthetic MoE test ran at 512/256 dims, which never reaches the
16-chunk activation shape the resident kernels see in production, and
its down-projection reference was built from the analytic mid. The GPU
quantizes its own gate/up mid before the down stage, and a legitimate
one-LSB rounding flip against the analytic mid moves one down dot by up
to d * 2^(e-127) * 6: at real-shape scales that exceeds any fixed
tolerance and failed tokens=3/4 despite both results being correct.
Default the test to the DS4 Flash routed shape (4096/2048, 256 experts,
6 used) and judge the down stage against a CPU requantization of the
GPU's own mid, so the comparison shares quantized inputs and the
tolerances stay tight enough to catch real kernel bugs. Expected rows
are cached per pattern and reused after a bitwise mid comparison, which
keeps the CPU reference cost flat across the 512-token sweep.
Coalesce MXFP4 weight reads across wave32 lanes and unpack nibbles with register permutations. Keep the gate/up and down reductions register-resident, route two-to-four-token batches through the direct kernels, and extend the ROCm CPU-oracle coverage across that dispatch boundary.
Implement resident MXFP4 gate/up and down paths for decode and batched routed MoE on ROCm. Add a synthetic CPU-oracle regression covering full, partial, and large expert tiles, report MXFP4 as a 4-bit routed quant, and make Strix host objects position independent for TheRock linking.
Integrate the MXFP4 Metal decode/prefill optimizations (rounds 7-9) as one
commit. Keep both the MXFP4 half_lut/tail_cull down-pipeline selection and
main's mpp_mask override for IQ2_XXS gate/up and Q2_K/IQ2_XXS down; they
target disjoint quant types and coexist.
Stop relying on the model to restart a call after an injected reasoning close. Wait for a complete call, keep the preceding prose as reasoning, return the call structurally, and keep its protocol text out of streaming output.
Closes#675
Track routed-expert cache entries referenced by encoded Metal command buffers and prevent cache eviction or buffer reuse until the owning command buffers have completed. This preserves layer-batched decode while avoiding stale buffer captures under SSD-streaming cache pressure.
Add a focused Metal SSD-streaming cache-pressure regression for issue #384.
The model sometimes opens a DSML stanza without closing its thinking
first. The generation loop deliberately ignores tool markers inside
<think>, so such turns ran to the token limit and the call was dropped
at parse time, stalling agent sessions (#318).
Recover forward instead of rewriting sampled context: when a complete
stanza opening appears inside an unclosed think block, force-feed
"</think>" plus a blank line and let the model continue. Measured on
the real model, that position predicts a fresh stanza opening strongly
enough that the call restarts cleanly on the executable side; the
dangling opening stays harmlessly inside reasoning. Re-emitting the
stanza opening ourselves was tried and rejected: the duplicated opening
reads as an already-made call and the model ends the turn.
Detection is on accumulated text, so marker tokenization does not
matter, and only a complete opening triggers: quoted fragments and a
lone "<" keep decoding untouched. DS4_SERVER_DISABLE_THINK_TOOL_RECOVERY
restores the old behavior.
The new think-tool-recovery test replays a teacher-forced malformed
prefix token by token, asserts the recovery triggers exactly when the
opening completes, and verifies the live model then emits a valid,
parseable call.
Add a small ds4-agent unit-test harness for anchored [upto] edits and include it in make test so regressions in old-text matching are caught by default.
metal_graph_verify_suffix_tops() asked ds4_gpu_indexer_topk_tensor() for the
top-`top_rows` tokens of a single logits row (n_tokens=1, top_k=top_rows)
instead of the top-1 of each of `top_rows` rows. The signature is
(selected, scores, n_comp, n_tokens, top_k), so row_tops[i>0] came back as
row 0 runner-ups and the verifier accepted or rejected drafts against the
wrong rows, committing tokens the model never produced. Visible only at draft
depth > 2: at depth 2 top_rows == 1 and the code takes the dedicated argmax
path above this branch, so the swapped call is never reached.
Adds tests/ds4_test.c --mtp-verify-depth: runs greedy speculative decode at
draft 4 over a verbatim-copy task, then teacher-forces the committed tokens
back through plain decode and asserts each is a (near-)argmax at its position.
That is the invariant speculative verify must preserve, and unlike comparing
whole token streams it tolerates the benign tie divergences of near-greedy
speculation. Self-skips unless DS4_TEST_MTP points at an MTP GGUF. Without the
fix a committed token sits ~21 logits below the argmax (worst gap 20.96 at
token 145); with the fix every committed token is the argmax (gap 0.00).
The shared test engine loads the MTP head only when DS4_TEST_MTP is set, and
only on the fast engine, so the default suite and the quality engine are
unaffected.
Verified on M4 Max / Metal, DeepSeek-V4-Flash q2-q4 imatrix base + MTP-Q4K-Q8_0
head:
make ds4_test
DS4_TEST_MODEL=<base> DS4_TEST_MTP=<mtp> ./ds4_test --mtp-verify-depth
Ran alongside --server --metal-tensor-equivalence --long-context: all pass.
The swapped call site is backend-shared (ds4.c), so CUDA is affected and fixed
identically, but this was not verified on a CUDA machine.
Refs https://github.com/antirez/ds4/issues/167.
DSML tool calls are executable only after the assistant closes the thinking stanza. The live decoder now skips tool-marker detection while thinking is open, and the final parser searches for executable DSML only after the last </think> when thinking mode is active. This prevents a DSML block inside reasoning from being duplicated as both thinking text and a structured tool call.
We were not able to prove from traces that the model actually made this exact mistake in the reported case, because the trace gives text rather than the original sampled-token/KV state. It is nevertheless possible, and it matches the reported shape of the bug. We did see evidence that the model can discuss tool usage and DSML syntax inside the thinking stanza, so this is a conservative protocol fix: it does not force the end of thinking, and it does not execute tool calls while still inside thinking.
It remains unclear what the model would do if it truly mis-emits a tool call before </think>: it might recover by closing thinking later, or it might not. A follow-up commit should address the likely root source by changing sampler behavior itself instead of relying only on protocol filtering.
antirez
Jonathan Yates
ivanfioravanti
Salvatore Sanfilippo
Andreas Spannagel
user
hexxyan
Fabio Malpezzi