perf(rocm): find why batched server decode is far below single-stream #2156

Part of #1801

Problem / Background

Server decode with four concurrent requests on gfx1151 is far slower per request than single-stream decode, with or without the paged-attention kernels. From docs/benchmark_results/rocm-paged-attention-gfx1151-2026-10-05.md (PR #2103, Meta-Llama-3.1-8B-Instruct-4bit, mlxcel-server --parallel 4 --ctx-size 131072, 128 decode tokens per request, medians of 3), per-request decode tok/s:

Single-stream bench_decode on the same model is about 35 to 37 tok/s (rocm-memory-gfx1151-2026-09-30.md:116). The ~1K figure means about 165 ms between tokens per request, against about 27 ms single-stream, and the rate rises with context (about 72 ms at ~16K), which a per-step compute cost cannot explain. The page attributes it to "time spent waiting while the other three requests prefill" without measuring it; three 1K prefills at about 1000 tok/s account for about 3 s of a roughly 21 s decode window, so that explanation is incomplete.

Current Behavior

scripts/bench_serving_concurrency.py:229-233 computes decode_tok_s = (completion_tokens - 1) / (total_s - ttft), so the window includes any scheduler time spent on other requests' prefill chunks. The scheduler loop is src/server/batch/scheduler/run_loop.rs, decode ticks are in decode_tick.rs, chunked prefill is in prefill.rs. /metrics already exposes mlxcel_batch_decode_steps_total, mlxcel_batch_decode_tokens_total, mlxcel_batch_mixed_steps_total, mlxcel_batch_prefill_chunks_total and the llama-compatible tokens_predicted_seconds_total (src/server/routes/metrics.rs:174-194, 440-500). The Metal production run (docs/benchmark_results/paged-decode-production-m1ultra-2026-08-02.md) publishes only after/before ratios, so no Metal absolute figure exists to compare.

Proposed Solution

Measure, attribute, then fix or file.

1. Separate decode from interleaved prefill in the harness. Extend bench_serving_concurrency.py (it already snapshots /metrics around each level for the path counters) to print, per level, the deltas of the batch counters above and the derived mean decode step time (tokens_predicted_seconds delta over decode steps) and mean batch occupancy (decode tokens over decode steps). This makes the real batched step cost visible in every later run.
2. Profile the decode-only window. On gfx1151, run the batch-4 ~1K case under rocprofv3 --kernel-trace --stats (as scripts/rocm_decode_profile.sh does for bench_decode), and attribute step time to GPU kernels versus host gaps. Check these candidates explicitly and report each as confirmed or ruled out with the evidence:
   • the GEMM route of the batched projections: this checkpoint is f16, so select_qmm_route (src/lib/mlx-cpp/patches-rocm/mlx/backend/rocm/quantized/qmm.hip:889-930) cannot take the bf16-only WMMA path; if the batch-4 shapes reach DequantGemm, each step rematerializes f16 weights that the 256 MiB dequant cache (MLX_ROCM_QMM_DEQUANT_CACHE_MAX_BYTES, qmm.hip:997-1001) cannot hold (re-test with MLX_ROCM_QMM_DEQUANT_M_THRESHOLD=64);
   • host-side scheduler work per tick (paged metadata, masks, evals and syncs per step) visible as gaps between kernels;
   • prefill chunks interleaved into decode steps (mixed_steps).
3. Get a reference ratio. Run the same matrix with the same model on one Metal host and record batch-4 decode step time over single-stream step time on both backends.
4. Act. Fix a cause that is a defect in the issue's PR when it is contained (for example a route or threshold); file a separate sub-issue of epic: AMD GPU (ROCm) backend on Linux via mlxcelverse #1801 for anything larger, with the evidence.

Scope

• In scope: the harness extension, the profile, the Metal reference run, a results page docs/benchmark_results/rocm-batched-decode-gfx1151-<date>.md, and a contained fix.
• Out of scope: the paged-attention kernels (both arms show the same gap), prefill throughput.

Acceptance Criteria

• bench_serving_concurrency.py prints decode step ms and mean batch occupancy per level from /metrics; a new tests/test_bench_serving_concurrency.py covers the delta arithmetic with a canned /metrics text, and the existing harness tests still pass.
• The results page gives, for batch 1 and batch 4 at ~1K and ~16K on gfx1151: per-request decode tok/s, decode step ms, occupancy, the GPU-kernel versus host-gap split, and the Metal ratio from step 3. (Metal ratio not measured: no Metal host on the gfx1151 machine; recorded as such on the page.)
• Each candidate in step 2 is marked confirmed or ruled out with the measurement that decided it.
• If a contained fix lands: batch-4 ~1K decode step ms before and after, medians of 3 under scripts/rocm_gpu_guard.sh, single-stream decode not regressed beyond 3%, and make verify-rocm passes. Otherwise the follow-up issue numbers are linked here.

Verification

cargo build --release --features rocm
MODEL=models/mlx/Meta-Llama-3.1-8B-Instruct-4bit scripts/benchmark_paged_decode_production.sh   # inside scripts/rocm_gpu_guard.sh
python3 -m pytest tests/test_bench_serving_concurrency.py
make verify-rocm

Related: #2103, #2068, #899, #2085, #2112.

Progress

• Cause found and fixed in the ROCm overlay (batched [B, 1, K] projections ran as B one-row products; batch fold plus a small-row bf16 route): batch-4 ~1K decode step 166.7 to 65.9 ms, single-stream unchanged. Results: docs/benchmark_results/rocm-batched-decode-gfx1151-2026-10-08.md.
• The ~16K figure is set by interleaved long-context prefill, 78% flash SDPA: filed as perf(rocm): flash SDPA dominates long-context chunked prefill on gfx1151 #2251.
• make verify-rocm is run by the orchestrator before merge.