[jiejingzhangamd]

Motivation

Add an opt-in, non-USERPTR allocation path for ROCm HiCache host pools to avoid page-migration-related GPU stalls. Request the pool's byte count directly from HIP rather than allocating through PyTorch's caching host allocator.

ROCm 7.2 enables USERPTR for the relevant paged-memory allocations when HSA_USERPTR_FOR_PAGED_MEM is unset; only 0 disables it (default initialization). Its host-allocation code selects between anonymous-memory USERPTR mappings and driver-managed GTT allocations.

Setting the variable alone does not change HiCache's existing hipHostRegister path. This PR switches eligible pools to hipHostMalloc without changing the ROCm default.

Changes

The implementation is in pool_host/common.py.

• Use the new path only on ROCm, with pinned memory, the default HostTensorAllocator, and HSA_USERPTR_FOR_PAGED_MEM=0.
• Bind the already-loaded HIP runtime and request numel * dtype.itemsize bytes with hipHostMalloc.
• Expose the allocation as a tensor without copying. Free it after the last tensor/view reference is released; skip cudaHostUnregister for these pools.
• Report allocation failures with the GTT capacity reference and a suggestion to reduce --hicache-size; do not silently fall back to USERPTR.
• Leave custom allocators, unpinned allocations, non-ROCm platforms, and configurations without the opt-in unchanged.

The exact-size allocation change comes from JohnQinAMD's commit 1a5d028.

Usage

Set this before starting SGLang, so it takes effect before HIP initialization:

export HSA_USERPTR_FOR_PAGED_MEM=0

The GTT path remains subject to driver allocation limits. SGLANG_HUGEPAGE_SIZE does not apply to this path; the allocator logs a warning if it is set.

Tests

test_hicache_host_register.py covers exact allocation size, release after the last view, skipping host unregister, unchanged allocation paths, and allocation failure without fallback. The new allocator tests mock HIP calls; they are not end-to-end performance or model-accuracy tests.

python3 -m pytest -q -p no:cacheprovider test/registered/unit/mem_cache/test_hicache_host_register.py

Local result at 1a5d028: 15 passed, 17 subtests passed, using lmsysorg/sglang-rocm:v0.5.20-rocm720-mi35x-20260924 on MI355X.

Reproducing the migration stall

bench_host_memory_migration.py is included in this PR. This is a manual hardware test requiring two memory NUMA nodes, sufficient host/GTT memory, and permission to call move_pages(); it is not part of regular CI.

Run the two arms in separate processes on the same idle GPU:

export HSA_USERPTR_FOR_PAGED_MEM=0
python benchmark/hicache/bench_host_memory_migration.py --allocator registered --pool-gb 180
python benchmark/hicache/bench_host_memory_migration.py --allocator gtt --pool-gb 180

The registered arm calls the legacy mmap + host-register path; the GTT arm calls the patched allocator. Both use the same environment. Each iteration runs four 4096-by-4096 BF16 matrix multiplications, a 16 MiB host-to-device copy, and synchronization. A second thread attempts three 2 MiB migrations, with 1,000 GPU iterations after each attempt. The script checks page-migration status, copy correctness, CPU wall-clock latency, and /sys/class/kfd/kfd/proc/<pid>/stats_<gpu>/evicted_ms. For a container, --kfd-pid identifies the worker's host PID without requiring a shared PID namespace.

Measured on MI355X GPU0, crsuse2-m2m-168, using the PR checkout in the 0924 image above (Docker CPU quota: 8; memory limit: 256 GiB). Both arms exited successfully.

Allocation	Maximum iteration time, trials 1 / 2 / 3 (ms)	Largest per-GPU KFD eviction delta per trial (ms)
Registered USERPTR	64162.856 / 64141.420 / 64152.832	64159 / 64140 / 64149
Exact-size GTT	0.796 / 0.891 / 0.792	0 / 0 / 0

The USERPTR arm moved all 512 pages on every attempt. The GTT arm returned -EFAULT for all 512 page-status entries each time, with zero eviction. Control median iteration times were 0.723 ms and 0.722 ms, respectively.

These are induced GPU-workload stalls, not token latency or P90 interactivity. The result establishes the migration/eviction behavior for this setup, not a general end-to-end speedup.

Performance

A fresh same-node sequential A/B was run on MI355X node crsuse2-m2m-059 (SPUR job 196589). Both arms used lmsysorg/sglang-rocm:v0.5.20-rocm720-mi35x-20260924, GLM-5.2 MXFP4, TP4/EP4, MTP 5/6/1, concurrency 8, HiCache 180 GB per rank, random seed 42, and a configured 3,600-second AgentX profiling window. The baseline used the image code with HSA_USERPTR_FOR_PAGED_MEM unset; the patched arm used PR head c7453eced20b6f5ed7a6d7f4905ae8da02ef94c4 with HSA_USERPTR_FOR_PAGED_MEM=0.

Metric	0924 baseline	PR + GTT opt-in	Change
P90 interactivity (tok/s/user)	86.57154	89.76933	+3.69%
Total throughput (tok/s)	52,781.70695	55,481.20941	+5.11%
Output/generation throughput (tok/s)	270.66238	257.17171	-4.98%
Total throughput per GPU (tok/s/GPU)	13,195.42674	13,870.30235	+5.11%
Output throughput per GPU (tok/s/GPU)	67.66559	64.29293	-4.98%
Measured profiling window (s)	3,606.56699	3,624.68331	—
Profiled requests / errors	1,284 / 0	1,311 / 0	—

Both runs exited cleanly on the same node. Their MTP accept lengths were 3.59689 and 3.59415. The patched run also recorded one warmup-only dropped InvalidInferenceResultError; it had zero profiling errors and does not enter the reported profiling metrics.

---

CI States

Latest PR Test (Base): ✅ Run #38122180045
Latest PR Test (Extra): ✅ Run #38122179876
Latest PR Test (AMD ROCm 10): ❌ Run #38122179882

[jiejingzhangamd]

/rerun-failed-ci

[jiejingzhangamd]

AMD CI failure comparison with scheduled main

I compared the failed AMD jobs in this PR with the closest scheduled PR Test (AMD) run on main that started before the PR run:

• Scheduled main: run 37395620023
• This PR: run 37424269792

Of the seven substantive failures, five reproduce on main with the same test and error signature:

Failure	Scheduled main evidence
MXFP8 route: (1792, 5120, 1, nan)	Exact match
KDA verify/decode parity: head_dim=128 and output-register allocation error	Exact match
Fused KDA bit-exact cases [8] and [9]	Exact match
MiniMax indexer: 131072 bytes of shared memory required, 65536-byte hardware limit	Exact match
Trace level 0: six spans collected instead of zero	Exact match

The same five failures also reproduce in the later scheduled main run 37553745648, so they are existing main failures rather than regressions introduced by this PR.

HiCache/AITER failure triage

I investigated the HiCache failure in more detail. The evidence strongly indicates that it is not caused by this PR:

• The crash happens during decode CUDA-graph capture while AITER is JIT-compiling module_moe_ck2stages_b16_b16_preshuffle_off_b16_silu_no_mulWeightStage2. Three parallel clang compilations fail with malloc(): smallbin double linked list corrupted, a segmentation fault, and free(): invalid pointer.
• This occurs before build_kv_cache() -> init_hicache(). The failed log never reaches the 200 GB HiCache host-pool allocation.
• The CI job does not set HSA_USERPTR_FOR_PAGED_MEM=0. Therefore the new hipHostMalloc/GTT branch added by this PR is not selected; the test retains the existing USERPTR allocation path.
• The preceding scheduled-main shard used the same ROCm image digest and AITER revision, built the same AITER module, allocated the HiCache pool, and passed. The following scheduled-main shard also passed.

I also ran TestHiCacheMLA.test_mmlu locally at PR head c7453ec on MI355X with the matching ROCm 10 image date and AITER revision, leaving HSA_USERPTR_FOR_PAGED_MEM unset as in CI. A fresh JIT build of the same AITER module completed in 48.1 seconds; CUDA-graph capture completed, the 200 GB HiCache host pool was allocated, and the server became healthy. The accuracy phase was then blocked only by an HTTP 429 while fetching the pinned MMLU archive. Separately, all 15 allocator unit tests pass locally.

The local run is not an exact gfx942 reproduction, but together with the execution order, inactive PR code path, and the two passing adjacent-main controls, it supports treating this as a transient ROCm clang/AITER JIT runner failure rather than a regression from #42553.

The MORI failure is also a one-off clang exit-code-139 during dependency installation; its preceding-main shard passed, and it happens before the test suite can exercise this PR. The remaining gate failures are downstream consequences, not additional test failures.

Overall, the AMD CI evidence does not indicate a regression caused by this PR.

[1am9trash]

LGTM.

• The behavior change is guarded by _uses_hip_owned_host_memory (_is_hip and env var HSA_USERPTR_FOR_PAGED_MEM=0), so there is no impact on the nv side.
• The new tests are CPU mocks and do not use a GPU. They also run on nv CI, and the non-hip case keeps the previous code path.
• bench_host_memory_migration.py is only a manual reproducer.

Hi @jiejingzhangamd,
Only one suggestion, not a blocker. pool_host/common.py is shared by every backend, and this PR adds some hip-specific helpers there. I think we could move them to a separate file in a follow-up PR?

[jiejingzhangamd]

LGTM.

• The behavior change is guarded by _uses_hip_owned_host_memory (_is_hip and env var HSA_USERPTR_FOR_PAGED_MEM=0), so there is no impact on the nv side.
• The new tests are CPU mocks and do not use a GPU. They also run on nv CI, and the non-hip case keeps the previous code path.
• bench_host_memory_migration.py is only a manual reproducer.

Hi @jiejingzhangamd, Only one suggestion, not a blocker. pool_host/common.py is shared by every backend, and this PR adds some hip-specific helpers there. I think we could move them to a separate file in a follow-up PR?

Thanks,

I've start a follow up PR at #42991 to move hip related helper in a separate file.

[jiejingzhangamd]

/rerun-failed-ci

[jiejingzhangamd]

/rerun-failed-ci

[jiejingzhangamd]

/rerun-failed-ci

[jiejingzhangamd]

/rerun-failed-ci