Repository navigation
perf(rocm): pipelined B=1 decode on Gemma 3 is 1.3 percent below the pre-epic loop at 2048 tokens on gfx1151 #2239
Description
Problem / Background
Found during the ROCm verification of the unified engine (#2192, Phase 7 of epic #2166). mlxcel-bench-decode at pp2048/tg128 on models/mlx/gemma-3-4b-it-4bit, Radeon 8060S (gfx1151), three interleaved rounds under scripts/rocm_gpu_guard.sh, gives three arms that do not overlap:
- Pre-epic
CxxGeneratorbinary (4c44e317): 62.66, 62.82, 62.61 tok/s. main(b2fda840,mlxcel_core::engine::DirectEngine, pipelined B=1 decode from refactor: run mlxcel generate on the engine and retire CxxGenerator #2176): 61.93, 61.83, 61.67 tok/s. Median -1.3 percent, every round below every base round; the null arm spread is -0.48 to -0.02 percent, so the gap is outside noise.- Same
mainbinary withMLXCEL_FORCE_SYNC=1(synchronous steps): 63.62, 63.69, 63.50 tok/s, +1.5 percent over base, every round above every base round.
At pp512 the pipelined arm sits inside ADR 0007's 1.0 percent threshold (-0.56 percent, null -0.14 percent). The ten other (model, context) cells of the same session, covering Qwen3-0.6B, Llama-3.1-8B, Qwen3-30B-A3B, granite-4.0-h-tiny, Nemotron-3-Nano, gpt-oss-20b and Mixtral-8x7B, are all within the threshold. The regression is therefore specific to the B=1 lookahead pipeline on Gemma 3's model-owned rotating cache once the context passes its 1024-token sliding window: with the window full, every pipelined step overwrites a valid slot, so record_decode_write (src/lib/mlxcel-core/src/cache/decode_undo.rs:154) copies that slot's K/V rows into an undo entry for the rewind that rewind_decode_writes (decode_undo.rs:232) replays, and flush_decode_undo_rows (decode_undo.rs:395) evaluates those copies ahead of the overwrite. The synchronous arm never records a row, which is consistent with it being the fastest of the three. The eligibility and dispatch live in Engine::can_unwind_lookahead (src/lib/mlxcel-core/src/engine/lookahead.rs:68) and DirectEngine::lookahead_params / decode_pipelined (src/lib/mlxcel-core/src/engine/direct_decode.rs:163, :211).
Evidence: docs/benchmark_results/rocm-unified-engine-gfx1151-2026-10-08.md (Performance section) and docs/benchmark_results/data/rocm-unified-engine-gfx1151-2026-10-08/gemma3-pp2048/, both landing with the PR for #2192.
Proposed Solution
- Repeat the same three arms (pre-epic
4c44e317,mainpipelined,mainwithMLXCEL_FORCE_SYNC=1) on CUDA (GB10) withmlxcel-bench-decodeat pp2048/tg128 on gemma-3-4b-it-4bit, three interleaved rounds plus a null arm, to establish whether the gap is ROCm-only. The pipeline policy belongs toDirectEngineand is shared with CUDA, so no policy change may be decided from the gfx1151 numbers alone. - Depending on the result, take one of two paths and measure it on both backends: (a) make the undo-row bookkeeping cheaper on the wrapped window, for example by recording only the overwritten slot's rows lazily or batching the records of one pipelined step into a single copy and a single flush, keeping
rewind_decode_writesbyte-identical in effect; or (b) haveDirectEngine::lookahead_paramsselect the synchronous step for model-owned rewind families once their rotating window is full (can_unwind_lookaheadtrue but the cache has wrapped), leaving the pipeline on for the pre-wrap phase and for trimmable families. - Whichever path is chosen, keep greedy output identical: the
make engine-parityarms must stay byte-identical, andMLXCEL_FORCE_SYNCkeeps its current meaning as a user override.
Acceptance Criteria
- The three-arm measurement (pre-epic, pipelined
main,MLXCEL_FORCE_SYNC=1main) is repeated on CUDA (GB10) at pp2048/tg128 on gemma-3-4b-it-4bit and recorded underdocs/benchmark_results/with the raw per-round data. - gemma-3-4b-it-4bit pp2048 B=1 decode on
mainis within 1.0 percent of the pre-epic loop on gfx1151, measured as medians of three interleaved rounds underscripts/rocm_gpu_guard.sh, with the null arm reported alongside. - The pp512 cell and the ten other (model, context) cells from the test(rocm): verify the unified engine on ROCm after the batch-native engine epic #2192 session remain within the 1.0 percent threshold after the change.
- Greedy output is unchanged: the
make engine-parityarms stay identical. -
make verify-rocmis green on the gfx1151 host.
Verification
# ROCm host (gfx1151): run each arm three times, interleaved (base, main, sync, base, main, sync, ...), plus a null arm (main vs main)
cargo build --release --bin mlxcel-bench-decode
scripts/rocm_gpu_guard.sh -- ./target/release/mlxcel-bench-decode -m models/mlx/gemma-3-4b-it-4bit --prompt-tokens 2048 -n 128
MLXCEL_FORCE_SYNC=1 scripts/rocm_gpu_guard.sh -- ./target/release/mlxcel-bench-decode -m models/mlx/gemma-3-4b-it-4bit --prompt-tokens 2048 -n 128
# The pre-epic arm is the same command with a binary built from 4c44e317
# CUDA host (GB10): the same three arms without the guard
./target/release/mlxcel-bench-decode -m models/mlx/gemma-3-4b-it-4bit --prompt-tokens 2048 -n 128
# Output parity and the ROCm gate
make engine-parity MODEL=models/mlx/gemma-3-4b-it-4bit PARITY_ARGS=--expect-identical
make verify-rocmA pass is the main pipelined median within 1.0 percent of the pre-epic median on gfx1151 with the null arm spread reported, identical engine-parity outputs, and a green verify-rocm.
Activity
test(rocm): verify the unified engine on ROCm (#2250)
update(perf): record the wrapped-window Gemma 3 decode measurement (l…