Skip to content

perf(rocm): pipelined B=1 decode on Gemma 3 is 1.3 percent below the pre-epic loop at 2048 tokens on gfx1151 #2239

Closed

Description

@inureyes

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 CxxGenerator binary (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 main binary with MLXCEL_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

  1. Repeat the same three arms (pre-epic 4c44e317, main pipelined, main with MLXCEL_FORCE_SYNC=1) on CUDA (GB10) with mlxcel-bench-decode at 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 to DirectEngine and is shared with CUDA, so no policy change may be decided from the gfx1151 numbers alone.
  2. 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_writes byte-identical in effect; or (b) have DirectEngine::lookahead_params select the synchronous step for model-owned rewind families once their rotating window is full (can_unwind_lookahead true but the cache has wrapped), leaving the pipeline on for the pre-wrap phase and for trimmable families.
  3. Whichever path is chosen, keep greedy output identical: the make engine-parity arms must stay byte-identical, and MLXCEL_FORCE_SYNC keeps its current meaning as a user override.

Acceptance Criteria

  • The three-arm measurement (pre-epic, pipelined main, MLXCEL_FORCE_SYNC=1 main) is repeated on CUDA (GB10) at pp2048/tg128 on gemma-3-4b-it-4bit and recorded under docs/benchmark_results/ with the raw per-round data.
  • gemma-3-4b-it-4bit pp2048 B=1 decode on main is within 1.0 percent of the pre-epic loop on gfx1151, measured as medians of three interleaved rounds under scripts/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-parity arms stay identical.
  • make verify-rocm is 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-rocm

A 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.

Refs #2192, #2166, #2182, #2176.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    area:inferenceGeneration, sampling, decoding (incl. speculative, DRY)platform:linuxLinux (CUDA / packaging) specificpriority:mediumMedium prioritystatus:readyReady to be worked ontype:performancePerformance improvements

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions