Fixed Piecewise

[wsantos]

rocm: get compiled PIECEWISE CUDA graphs working on the V620 stack

Follow-up to #17. --compilation-config '{"mode":3,"cudagraph_mode":"FULL_AND_PIECEWISE",...}'
now boots, captures piecewise graphs and serves correct output on the V620 stack
(Qwen3.8-Flash-Next W4A16 AutoRound, TP4/EP4, MTP2, PLE CPU offload, prefix caching,
4x V620 / gfx1030). Previously PIECEWISE was unreachable.

Why this is not duplicating an existing PR

Checked the open PRs against rdna_extras (#19, #18, #10, #9, #2) — none touches
CUDA graph capture, cudagraph_mode or the mRoPE/HC strides. #17 merged the
serving stack but kept mode 0 + FULL_DECODE_ONLY, so compiled piecewise capture
had never run on this stack.

What was wrong

1. cudagraph_mode: PIECEWISE with mode: 0 silently produced no graphs.
   VllmConfig.__post_init__ replaced the mode with NONE (piecewise capture needs
   Inductor splitting or breakable CUDA graphs) and only logged at INFO:
   Cudagraph mode PIECEWISE is not compatible with compilation mode 0. Overriding to NONE.
   The server answered correctly but every step ran eagerly.
2. HC injection stride. dai[:, lora_rank : lora_rank + hc_count] is a
   (num_tokens, 4) view with row stride 336; Inductor specializes on the trace-time
   stride and capture aborted with
   AssertionError: expected size 12==12, stride 4==336 at dim=0.
3. MTP speculator did not follow the ROCm FULL→piecewise redirect.
   rocm_full_executes_as_piecewise makes the target replay piecewise graphs instead
   of capturing FULL ones, but draft prefill still called run_fullgraph()
   (No cudagraph for ... FULL num_tokens=12 num_reqs=4 uniform_token_count=3) and the
   draft-decode manager wrongly skipped FULL capture because the skip decision read the
   global compilation config (No cudagraph for ... FULL num_tokens=3 num_reqs=3 uniform_token_count=1).
4. MRV1 mRoPE/XD-RoPE positions were non-contiguous
   (AssertionError: expected size 3==3, stride 2048==2049 at dim=0) — root-caused
   earlier but left unfixed; MRV2 already returned a contiguous copy.

Changes

File	Change
vllm/config/vllm.py	warning_once naming the remedy; FULL_AND_PIECEWISE downgrades to FULL_DECODE_ONLY (keeps its FULL graphs) instead of NONE.
vllm/models/qwen4_exp/amd/hyperconnection.py	_injection_slice() returns a contiguous HC injection; the non-fused split path too.
vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py	Draft prefill routes a FULL dispatch to _prefill(cudagraph_runtime_mode=PIECEWISE) when the redirect applies.
vllm/v1/worker/gpu/cudagraph_utils.py	FULL-capture skip also requires self.cudagraph_mode.has_piecewise_cudagraphs().
vllm/v1/worker/gpu_model_runner.py	Packed contiguous mRoPE/XD-RoPE workspaces + _contiguous_positions().
tools/rdna2/serve_v620_piecewise.sh	New launcher: same parameters as serve_v620_baseline.sh, mode 3 + FULL_AND_PIECEWISE, all env/paths self-contained and overridable.
docs/rdna2/V620-PIECEWISE-CUDAGRAPHS.md	Investigation, fixes, config and measurements.

Tests run

.venv/bin/python -m pytest --noconftest -q tests/compile/test_cudagraph_replay_inputs.py \
  -k "mrope or mrv1 or clone_activations_makes_index or copy_nested or copy_and_call"
# 5 passed

.venv/bin/python -m pytest --noconftest -q tests/compile/test_config.py \
  -k "piecewise_cudagraph_without_compilation or custom_op or copy_pass or version"
# 7 passed

tests/compile/test_config.py::test_splitting_ops_dynamic fails on this ROCm fork
both before and after the change (the ROCm block appends the TP collective ops to
splitting_ops); verified unchanged by stashing the branch.

Two new tests: test_piecewise_cudagraph_without_compilation and
test_mrv1_contiguous_positions_per_capture_size.

Model evaluation (4x V620, 16k / 1024 output, MTP2)

Same box, same harness (llm-context-bench, the #17 author's protocol), same
launcher parameters; the only difference is cudagraph_mode. Both tables are
valid (cold protocol: --profile pr17-rocm10-tuned-cold --suite all --lane performance --sizes 16k --max-retries 0 --input-size-tolerance-percent 11, fresh
boot, 0.0 % prefix-cache hit rate).

16k suite	Prompt tokens	Prefill tokens/s	Generation tokens/s	TTFT s	Valid
FULL_DECODE_ONLY regular prose	16,750	1,983.8	63.07	8.44	Yes
FULL_DECODE_ONLY coding	18,063	1,985.3	70.08	9.10	Yes
FULL_AND_PIECEWISE regular prose	16,750	1,972.4	25.74	8.49	Yes
FULL_AND_PIECEWISE coding	18,063	1,980.5	33.84	9.12	Yes

Prefill unchanged; mean inter-token latency 15.85 → 38.85 ms (regular) and
14.27 → 29.55 ms (coding). On ROCm the compiled path replays piecewise graphs with
eager GDN/attention breaks, so mode 0 + FULL_DECODE_ONLY remains the faster
decode configuration. This PR makes PIECEWISE available and correct, not faster;
a capture ladder covering prefill chunks is the follow-up that would change the
trade-off, and it needs a VRAM measurement against the fixed 3.75 GiB KV budget.

Boot evidence: ROCm FULL decode executes piecewise CUDA graphs +
Graph capturing finished in 2-3 secs, took 0.46 GiB; probes Paris,
1 + 1 = 2, 5 + 5 = 10, Berlin; MTP acceptance 80.8 %. A bare
"cudagraph_mode":"PIECEWISE" also boots and serves correct output.

Known limitations

• Breakable CUDA graphs (mode 0 + VLLM_USE_BREAKABLE_CUDAGRAPH=1 + PIECEWISE)
  capture but fault on replay (HSA_STATUS_ERROR_EXCEPTION in
  indexSelectSmallIndex<Half,long>); documented, not fixed.
• Ladders larger than [3,6,12] were not measured against the 3.75 GiB KV budget
  (a 1..256 ladder OOMs during KV allocation).

AI assistance

AI assistance was used to investigate, implement and test this change. The
submitting human reviewed every changed line and ran the commands above.

[BlivionIaG]

Integrated into rdna_extras via the local integration branch (tip 68a635e, which carries this PR's capture fixes together with #21's redirect change and the FA selection fix).

Measured on 4× V620 (gfx1030, TP4, W4A16 Flash-Next, 8× concurrency, seed 12345), merged tree:

• FULL_AND_PIECEWISE = FULL_DECODE_ONLY baseline: 62.56 vs 61.41 tok/s @ 16k/1k, TPOT 75.01 vs 76.63 ms; 166.03 vs 161.98 tok/s @ 1k/512 — no decode regression.
• PIECEWISE alone loses decode: TPOT 91.45 ms @ 16k (+19%) and 58.09 ms @ 1k/512 (+39%).
• Bare FULL (mode 3) still crashes at capture (expected size 8==4, stride 2560==2560; 'FULL is not supported with Qwen4ExpQSAFlashAttentionBackend').
• MTP2 still hangs at startup (pre-existing, unrelated to capture).

The V620 PIECEWISE launcher and the measured report from this PR are kept in the merged branch. Thanks — the capture fixes are what make the F&P combination possible.

[BlivionIaG]

Correction to the closing note above: the "bare FULL crashes at capture" finding was wrong - it was stale Inductor-cache poisoning in my test harness, not a capture bug.

With a fresh TORCHINDUCTOR_CACHE_DIR: cudagraph_mode=FULL + mode 3 boots, the runner resolves it to FULL_AND_PIECEWISE (warning: FULL is not supported with Qwen4ExpQSAFlashAttentionBackend, support UNIFORM_BATCH) and captures the FULL graphs. Measured on 4x V620 (8x concurrency, seed 12345): 62.24 tok/s / TPOT 75.12 ms at 16k/1k; 166.31 tok/s / TPOT 40.70 ms at 1k/512 - parity with F&P and the FULL_DECODE_ONLY baseline.

Root cause of the earlier crash: VLLM_DISABLE_COMPILE_CACHE=1 does not cover TORCHINDUCTOR_CACHE_DIR; my matrix reused a graph artifact specialized for batch 4 while capturing size 8 (expected size 8==4, stride 2560==2560). Operational rule: use a fresh inductor cache for every mode-3 configuration change.