[Bugfix] Make _fold_seqlen_indptr cudagraph-safe (avoid scalar H2D copy)

[micah-wil]

Motivation

This PR fixes a bug in _fold_seqlen_indptr (introduced in #4964) in which the function is not cudagraph safe as it can trigger a H2D transfer.

Technical Details

_fold_seqlen_indptr does:

out = torch.empty(..., device=indptr.device)
out[0] = 0
out[1:] = torch.cumsum(folded_lens, dim=0).to(indptr.dtype)

out[0] = 0 is unsafe during cuda graph capture as it can be lowered to an H2D copy, resulting in operation not permitted when stream is capturing.

Test Plan

This was first exposed in vLLM CI when running the following test on MI300 HIP_VISIBLE_DEVICES=0 pytest -v -s tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py::test_eagle_correctness_light[ROCM_AITER_FA-deepseek_eagle]
https://buildkite.com/vllm/amd-ci/builds/12434/list?jid=01a05932-e4b1-4cd7-90df-6cefc7be92f9&tab=output#L1820

(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]   File "/usr/local/lib/python3.12/dist-packages/vllm/_aiter_ops.py", line 2731, in mla_decode_fwd
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]     torch.ops.vllm.rocm_aiter_mla_decode_fwd(
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]   File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 1275, in __call__
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]     return self._op(*args, **kwargs)
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]            ^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]   File "/usr/local/lib/python3.12/dist-packages/vllm/_aiter_ops.py", line 604, in _rocm_aiter_mla_decode_fwd_impl
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]     mla_decode_fwd(
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]   File "/usr/local/lib/python3.12/dist-packages/aiter/mla.py", line 862, in mla_decode_fwd
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]     qo_indptr = _fold_seqlen_indptr(qo_indptr, fold_factor)
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]   File "/usr/local/lib/python3.12/dist-packages/aiter/mla.py", line 410, in _fold_seqlen_indptr
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]     out[0] = 0
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374]     ~~~^^^
(EngineCore pid=7775) ERROR 08-31 19:15:40 [core.py:1374] RuntimeError: Cannot copy between CPU and CUDA tensors during CUDA graph capture unless the CPU tensor is pinned. Please use tensor.pin_memory() or allocate the tensor with pin_memory=True.

Test Result

Test passes with this PR:

============ 1 passed, 14 warnings in 696.39s (0:11:36) ==================

[github-actions]

🏷️ CI Guide

Runs automatically on every PR:

• ✅ Pre-checks (submodule verification, code formatting)
• ✅ Aiter op tests (gfx942 + gfx950)
• ✅ Triton tests on MI35X (only when aiter/ops/triton/** or related paths are changed)

Extended tests (opt-in via labels):

Label	Tests
ci:gfx1250-ffm-triton	Run the five-shard gfx1250 FFM Triton test suite
ci:triton-300x	Run an additional Triton test job on MI300X in PRs; main branch always runs both MI35X and MI300X
multigpu	Aiter multi-GPU tests on the 8-GPU runner
ci:sglang	SGLang integration tests: DeepSeek-R1-MXFP4 accuracy, Qwen 3.5 accuracy
ci:atom	ATOM benchmark: DeepSeek-R1-0528, GPT-OSS-120B
ci:atom_full	ATOM accuracy suite for PR and main models from ATOM models_accuracy.json
ci:vllm	vLLM benchmark: GPT-OSS-120B, DeepSeek-R1-0528, Kimi-K2.5
ci:all	All standard extended tests (excludes ci:atom_full)

Only add ci:atom_full for FlyDSL or Triton upgrades.
Add labels via the sidebar or gh pr edit 5202 --add-label <label>

PR title tags & labels:
Component tags ([Triton/Gluon], [HIP], [CK], [ASM], ...) are added to the PR title and as PR labels automatically from the changed files and re-synced on every push — change-type tags like [fix]/[Perf], op tags like [MLA], and human labels (ci:*) are left untouched. Add the no-auto-title label to opt this PR out.

[AndreasKaratzas]

Can we ensure any missing CI tests are also added for this issue? they can ofc be vibe coded.

[micah-wil]

@AndreasKaratzas good call, done.

[zufayu]

The fix is right. I checked old vs new on 30 combinations (batch ∈ {0,1,2,8,37,129} × fold_factor ∈ {1,2,3,8,16}) — identical shape, dtype and values, including the batch=0 edge where cumsum is empty. F.pad keeps everything on device and the scalar H2D is gone.

The only hard blocker is one line of formatting. Black wants the two adjacent string literals joined:

     aiter.logger.info(
-        "_fold_seqlen_indptr cuda-graph capture/replay: all passed "
-        "(fold_factor=%s)",
+        "_fold_seqlen_indptr cuda-graph capture/replay: all passed " "(fold_factor=%s)",

Worth stating the knock-on: check-signal reports Pre-checks did not pass ... FAILED: Check Code Style with Black, so 27 of 38 checks skipped. The new test has never executed on a runner. The red has nothing to do with correctness, but the green hasn't been earned yet either.

On the test itself — most of it cannot fail. Measured on gfx942 (ROCm 7.0, torch 2.9.1), capturing the pre-fix implementation:

  b=1   ff=3   old-> HIP error: operation not permitted when stream is capturing   new-> ok
  b=1   ff=8   old-> HIP error: operation not permitted when stream is capturing   new-> ok
  b=8   ff=3   old-> HIP error: operation not permitted when stream is capturing   new-> ok
  b=8   ff=8   old-> HIP error: operation not permitted when stream is capturing   new-> ok
  b=37  ff=3   old-> HIP error: operation not permitted when stream is capturing   new-> ok
  b=37  ff=8   old-> HIP error: operation not permitted when stream is capturing   new-> ok

Every combination raises the same error at the same place, so num_batches=(1, 8, 37) crossed with the two fold factors buys six graph captures where one would prove the same thing. The axes are orthogonal to the defect — out[0] = 0 is a constant copy, unaffected by batch size or fold factor.

The replay half is dead for the same reason: the exception fires inside with torch.cuda.graph(g):, so control flow never reaches graph.replay() or the torch.equal(out, ref) comparison. It was carried over from test_moe_sorting_flydsl_cuda_graph_capture, where it earns its place — that bug was .item() baking a capture-time count into the graph, so replaying with a different value is exactly the thing under test. Here there is no analogous failure mode.

Structure. The call site sits at module top level, while the test it cites as its model keeps everything inside main() (test_moe_sorting.py: main() at 486, parse_args() at 588, the capture call at 620, guard at 669). test_mla_persistent.py has no __main__ guard at all — args = parser.parse_args() is bare at 2117 — so any import of this module now runs argparse against the caller's sys.argv and performs six graph captures as a side effect. Pre-existing, but this raises the cost of it.

Minor: test_fold_seqlen_indptr_cuda_graph_capture starts with test_ and takes a required positional. Fine under python3 <file>, a fixture error the moment pytest collects it.

Comments are 33% of the added block (20 docstring + 7 # against 43 code lines), and the docstring restates the bug history that the PR description already carries, traceback included.

Suggested shape: one capture, one shape, pass if it doesn't raise — about ten lines — and move the call inside a function rather than running it at import. The capture-safety check is the right instinct; aiter's op_tests are eager-only, which is why this had to surface in vLLM CI instead of here.

[micah-wil]

@zufayu Fixed, thanks!