perf(rocm): remove per-call KV-cache copies from the MLA and paged-append paths

[demandal25]

Summary

BatchMLAPagedAttentionWrapper.run() concatenated ckv_cache and kpe_cache on every call to build the single buffer AITER's MLA kernels read. torch.cat copies the whole tensor it is handed, so the cost tracked the allocated page pool rather than the live pages — at a 4.5 GB pool with a 9 MB working set, 2.799 ms per decode step, per layer. This makes that copy a view when the two halves are adjacent, which is how vLLM and SGLang already allocate MLA caches.

Also collapses plan()'s two device→host syncs into one, and moves the AITER paged-KV append off AITER's per-call Python dispatch into a compiled shim.

What changed

• flashinfer/mla_rocm.py — _combined_kv_view() detects ckv/kpe as adjacent halves of one allocation and returns a zero-copy as_strided view; run() falls back to the copy with a one-time warning otherwise, and rejects 4-D caches outright. _gather_plan_inputs_on_host() fetches the three index tensors in a single host round-trip.
• flashinfer/csrc_rocm/page_aiter.cu — new shim calling aiter::reshape_and_cache_flash, with the FlashInfer→vLLM slot-index translation folded into one kernel (it was seven Python tensor ops). Validation matches csrc_rocm/page.cu: device, rank, head/dim, nnz, launch status.
• flashinfer/csrc_rocm/page_aiter_jit_pybind.cu, flashinfer/jit/page.py — TORCH_LIBRARY_FRAGMENT registration and gen_page_aiter_module(), following the norm_aiter pattern.
• flashinfer/page.py — routes the AITER append through the shim behind register_custom_op + register_fake_op; _try_get_page_aiter_module() keeps backend="auto" from raising when the shim cannot build; explicit backend="aiter" now enforces the same NHD/dtype gates auto applies. Adds the missing fake op for append_paged_mla_kv_cache.
• flashinfer/__init__.py — export append_paged_mla_kv_cache on HIP; it was CUDA-only.
• benchmarks/rocm_benchmarks/bench_mla_hip.py — new, four modes (attn/cat/plan/pool) so results stay attributable; --separate reproduces the pre-fix allocation layout.

Architecture / design notes

Detection rather than a new parameter, because production already has the layout. vLLM stores kv_c_and_k_pe_cache and SGLang stores K_Buffer, both sliced with torch.split — exactly the adjacency this detects — so those callers get the fast path with no code change and the signature stays identical to the CUDA wrapper. The rejected alternative, caching the concatenated buffer across calls, is silently wrong: the append op mutates the cache between decode steps, so a cached copy serves stale KV.

The old docstring told callers to pre-allocate a combined buffer and "pass it as both arguments (sliced)". That never worked — torch.cat copies regardless of how its inputs were carved — and the shape it named (4-D, with a size-1 axis) would then be unsqueezed to 5-D, which AITER rejects. Replaced with the 3-D form the code now honors; 4-D input raises instead of warning that a genuinely adjacent buffer is not adjacent.

plan() is bounded by round-trip count, not payload. A device→host sync costs ~70 µs on gfx942 regardless of size, and these tensors are batch+1 int32 elements, so batching the three into one copy is the whole win. The remaining ~70 µs is the H2D of the computed last_page_len; removing it means computing on-device while validating on host, which needs kv_len_arr on the device and turns circular when it is not.

The shim needs a fake op. Registering it through TORCH_LIBRARY_FRAGMENT makes it visible to Dynamo, which places it in the FX graph and fake-runs it during tracing — and the shim calls numel()/size(), which throw on symbolic shapes. The Python path it replaced was not a registered torch op, so Dynamo graph-broke on it instead. The norm/rope/activation AITER shims call their modules directly from the public entry points with no custom-op wrapper, so they carry the same exposure; not changed here.

backend="auto" must never raise. The shim links AITER's C++ symbols, so it can fail for reasons an import probe cannot see (no hipcc, unwritable cache dir, AITER signature drift). The None is cached deliberately — functools.cache does not memoize exceptions, so a raising getter would retry the whole AITER rebuild on every call. MissingJITCacheError is re-raised rather than caught, since conftest.py turns it into a skip and records the module for the prebuilt cache.

Benchmark results

MI300X (gfx942, 304 CU), bf16, page_size 1, torch 2.9.1+rocm7.2, amd-aiter 0.1.10.

MLA decode run() — active set fixed at b=8 / kv=1024 (9 MB), only the pool grows:

pool	before	after
0.01 GB	0.081 ms	0.071 ms
0.28 GB	0.281 ms	0.071 ms
1.12 GB	1.076 ms	0.072 ms
4.50 GB	2.799 ms	0.071 ms

Flat across a 450× range — cost now depends on the working set. Production pools fill HBM, well past the largest row here.

MLA decode, standard shapes:

batch	kv_len	heads=16	heads=128
8	32768	0.362 → 0.126 ms (2.9×)	0.573 → 0.313 ms (1.8×)
32	8192	0.360 → 0.130 ms (2.8×)	0.507 → 0.242 ms (2.1×)
32	32768	1.424 → 0.349 ms (4.1×)	1.806 → 0.756 ms (2.4×)

heads=16 is TP8, heads=128 is DP-attention + EP; AMD recommends the former below ~256 concurrent requests and the latter above ~512, so both are load-bearing.

plan(): 120–130 µs → 87–92 µs, flat in batch size (1..256).

AITER paged-KV append: 102.9 → 11.0 µs at nnz=32, 101.2 → 11.2 at 256, 101.8 → 12.1 at 4096 (~9×). Also removes the lazy build-on-first-call that otherwise stalls the first append in a server process.

Measured, but deliberately not changed

• auto routes the KV append to AITER, which is slower than the in-tree kernel — 7.4 vs 11.6 µs at nnz=4096, 593 vs 751 µs at 262144, converging to ~0.79×. The gap is AITER's reshape_and_cache_flash itself (2.86 vs 3.62 TB/s on identical work), not the shim's slot-mapping kernel, which is free at scale. This predates the PR and only gfx942 was measured, so the routing is left alone.
• The MLA decode kernel has no headroom at TP8 — against a measured bare-read floor for the same KV volume it runs at 1.01–1.09×, i.e. full attention for roughly the cost of reading the KV cache once. AITER's num_kv_splits heuristic is also already optimal (best 1.02× across ten shapes against fixed values).
• Neither append kernel is tuned for CDNA3 — AppendPagedKVMlaCache hardcodes vec_size = 2 (4-byte accesses for bf16, where the MHA path computes 16-byte) and AppendPagedKVCache derives it partly from HEAD_DIM / 32, a CUDA warp-size constant. Coalescing recovers most of it (68% of peak at nnz=65536 vs 75% for a plain copy_), so the impact is moderate.
• build_slot_mapping_kernel does unchecked page-table indexing where the torch.index_select it replaced was bounds-checked. This is parity with the native kernel, which is equally unchecked, so it should be fixed in both together rather than here.

CDNA3 impact

This is the CDNA3 path. Everything above was measured on gfx942. gfx950 is unmeasured and no behaviour is arch-gated, so CDNA4 gets the same code with unverified numbers.

Test plan

Verified on 8× MI300X (gfx942, ROCm 7.2.0, torch 2.9.1+rocm7.2.0, amd-aiter 0.1.10):

• Full tests/rocm_tests suite, -n auto --reruns 2 -m "not slow": 27,535 passed, 3,584 skipped, 0 failed, 0 errors.
• Combined-view and concatenating paths produce bitwise-identical output (torch.equal, not a tolerance) at three shapes including heads=128.
• Adjacency detection unit-tested on CPU tensors: accepts torch.split views and layer-indexed slices; rejects separate allocations, reversed halves, padded row strides, mismatched dtypes — each of which would otherwise alias the wrong memory.
• The pre-existing aiter-vs-native append tests assert bit-exactness across dtypes, page sizes and head configs, and cover the new kernel unchanged.
• The AITER append honors torch's current stream. amd-aiter 0.1.10 reads at::hip::getCurrentHIPStream(), so this passes today; newer AITER moved these ops to a thread-local stream, which would silently put the append on the null stream — and since torch's default stream is stream 0, every other test would still pass. This is the one that would not.
• test_torch_compile_without_custom_ops_fails un-skipped: its reason ("AITER routes append_paged_kv_cache away from the custom-op path") is obsolete now that the shim goes through register_custom_op, and it was being skipped on exactly the hardware that exercises it.
• pre-commit run -a

[Copilot]

Pull request overview

This PR improves ROCm MLA and paged KV-cache append performance by removing unnecessary copies and host synchronizations, while adding compiled AITER dispatch.

Changes:

• Adds zero-copy MLA cache handling and batched planning transfers.
• Adds validated, fallback-capable AITER append shims with torch.compile support.
• Adds tests, benchmarks, ROCm export coverage, and profiler metadata ignores.

Reviewed changes

Copilot reviewed 10 out of 11 changed files in this pull request and generated 6 comments.

Show a summary per file

File	Reviewed changes	Final review comments
tests/rocm_tests/test_torch_compile_hip.py	Updates compile-path coverage.	—
tests/rocm_tests/test_mla_aiter_hip.py	Tests MLA views, planning, and validation.	—
tests/rocm_tests/test_append_paged_kv_cache_aiter_hip.py	Tests append correctness, validation, fallback, and streams.	Nit (3 votes): Replace device-wide synchronization with an inter-stream event/dependency check.
flashinfer/page.py	Adds AITER dispatch, fallback, validation, and fake-op integration.	Critical (1 vote): Isolate module caching by architecture/device or reject cross-architecture reuse. Moderate (2 votes): Normalize accepted int64 metadata or support both dtypes in the shim.
flashinfer/mla_rocm.py	Adds zero-copy cache detection and batched planning.	Nit (2 votes): Add an explicit one-time guard for the fallback warning.
flashinfer/jit/page.py	Defines the AITER page JIT specification.	—
flashinfer/csrc_rocm/page_aiter.cu	Implements slot mapping and AITER append dispatch.	Critical (1 vote): Validate V-cache rank, shape, strides, and K/V stride compatibility. Critical (1 vote): Validate kv_indptr rank and length before slot mapping.
flashinfer/csrc_rocm/page_aiter_jit_pybind.cu	Registers the compiled append shim.	—
flashinfer/__init__.py	Exports the ROCm MLA append API.	—
benchmarks/rocm_benchmarks/bench_mla_hip.py	Adds MLA performance benchmarks.	—
.gitignore	Ignores profiler metadata outputs.	—

Suppressed comments (5)

benchmarks/rocm_benchmarks/bench_mla_hip.py:326

• _alloc_kv() returns adjacent views by default, so run() takes the zero-copy path and performs no pool-sized torch.cat; that copy exists only with --separate. This KernelConfig nevertheless always reports 2 * pool_pages * ... bytes, so the default --mode pool roofline/arithmetic-intensity output attributes GBs of nonexistent traffic and incorrectly grows with pool capacity. Make the byte count conditional on _bench_args.separate and include the active attention traffic.

        # The pool-sized cat dominates: read ckv+kpe, write the combined buffer.
        theoretical_bytes=2 * pool_pages * _QK_HEAD_DIM * 2,

flashinfer/csrc_rocm/page_aiter.cu:72

• The new shim rejects int64 index tensors, but the public API currently accepts them: the native path coerces batch_indices, positions, kv_indices, and kv_indptr with .int(), and the previous AITER path converted them with .long(). A caller using default torch.arange/int64 tensors therefore succeeds with backend='native' but now gets this error on the AITER auto/explicit path. Preserve the existing API by conditionally converting these inputs before the custom op or dispatching the shim for both integer widths.

  TORCH_CHECK(batch_indices.scalar_type() == at::kInt && positions.scalar_type() == at::kInt &&
                  kv_indices.scalar_type() == at::kInt && kv_indptr.scalar_type() == at::kInt,
              "batch_indices/positions/kv_indices/kv_indptr must be int32");
  TORCH_CHECK(batch_indices.numel() == positions.numel(),
              "batch_indices and positions must have the same length, got ", batch_indices.numel(),
              " vs ", positions.numel());

flashinfer/jit/page.py:34

• aiter_jitspec_flags() eagerly calls ensure_aiter_lib() before returning, so this generator builds the external AITER module even when FLASHINFER_DISABLE_JIT=1. That bypasses the intended JitSpec.build() MissingJITCacheError path: an uncached or failed AITER dependency raises a generic build error (explicit AITER tests fail instead of being skipped/recorded, while auto fallback can hide the missing page_aiter cache). Defer resolving/building these flags until JIT is allowed, or handle the disabled-JIT path before calling aiter_jitspec_flags().

    extra_include_paths, extra_ldflags = aiter_jitspec_flags("module_cache")

flashinfer/jit/page.py:42

• The page_aiter spec has a fixed name, while aiter_jitspec_flags() points at a version/architecture-specific cached AITER library. Once page_aiter/build.ninja exists, JitSpec.build() does not rewrite it, so upgrading AITER or changing the selected architecture leaves the shim linked/rpathed to the old library (or silently keeps using it). Include the AITER dependency identity in the spec/cache key or invalidate/rewrite the ninja file when these flags change.

    extra_include_paths, extra_ldflags = aiter_jitspec_flags("module_cache")
    return gen_jit_spec(
        "page_aiter",
        [
            jit_env.FLASHINFER_CSRC_DIR / "page_aiter.cu",
            jit_env.FLASHINFER_CSRC_DIR / "page_aiter_jit_pybind.cu",
        ],
        extra_include_paths=extra_include_paths,
        extra_ldflags=extra_ldflags,

flashinfer/page.py:247

• This fake implementation is the fix for torch.compile on the newly exported MLA append op, but the compile test exercises only append_paged_kv_cache; no test invokes append_paged_mla_kv_cache with custom ops enabled under torch.compile. A registration or signature regression here would therefore pass the suite. Add a minimal compiled MLA append case.

@register_fake_op("flashinfer::append_paged_mla_kv_cache")
def _fake_append_paged_mla_kv_cache_kernel(
    append_ckv: torch.Tensor,
    append_kpe: torch.Tensor,
    batch_indices: torch.Tensor,
    positions: torch.Tensor,
    ckv_cache: Optional[torch.Tensor],
    kpe_cache: Optional[torch.Tensor],
    kv_indices: torch.Tensor,
    kv_indptr: torch.Tensor,
    kv_last_page_len: torch.Tensor,
) -> None:
    pass

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.