[inureyes]

Summary

Follow-up to #2474 (Ref #1564) after validating weight materialization and staging on ROCm (gfx1151, AMD Radeon 8060S Graphics, ROCm 7.15):

• Schedule load_safetensors on the GPU stream on ROCm: src/lib/mlxcelverse/overlays/rocm/mlx/io/safetensors.cpp checked only cu::is_available() when selecting the default Load stream, so ROCm builds defaulted to Device::cpu (Load::eval_cpu), bypassing Load::eval_gpu and introducing cross-stream CPU-to-GPU fences whenever a loader derived tensors from Load arrays. It now checks (cu::is_available() || rocm::is_available()).
• Reused 4 × 64 MiB pinned StagingPool in mlx/backend/rocm/load.cpp: Replaced per-tensor hipHostMalloc(&out_ptr, nbytes, hipHostMallocDefault) + hipLaunchHostFunc(encoder.stream(), hip_host_free_callback, out_ptr) with four reused 64 MiB hipHostMallocPortable slots per HIP device and per-slot hipEventDisableTiming events, bounding pinned staging memory at 256 MiB per device and eliminating KFD page-pinning churn under eager eval_all. MLXCEL_ROCM_PINNED_LOAD=0 (or MLXCEL_CUDA_PINNED_LOAD=0) selects the synchronous pageable fallback.
• ROCm upload counters (LoadTransferStats): Added src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/load_staging.h and wired rocm::load_transfer_stats() into weight_load_transfer_stats in src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp under MLXCEL_BRIDGE_ROCM_BACKEND.

Measurements (gfx1151, warm page cache, --no-chat-template -p "The capital of France is" -n 16 --temp 0 --seed 0)

• Standalone HIP microbenchmark (16 × 128 MiB = 2 GiB):
  • Per-tensor hipHostMalloc + pread + hipMemcpyAsync + hipHostFree: 617.22 ms (hipHostMalloc=297.26 ms, pread=156.37 ms, hipMemcpyAsync=47.67 ms, hipHostFree=115.91 ms)
  • Reused 4 × 64 MiB pinned StagingPool: 162.98 ms (3.79× faster, read=160.20 ms, wait=0.04 ms, enqueue=1.05 ms)
• End-to-end model load & TTFT (MLXCEL_PROFILE_TTFT=1):

Model	Arm	Model load	First prefill (TTFT)	Load + first token	Decode	[TTFT] upload counters
Qwen3-0.6B-4bit (704 tensors, 0.34 GB)	eager + pinned	0.304 s (0.31 GB res.)	397.61 ms	0.702 s	303.35 tok/s	load eval=85.66ms pinned=704 memcpy=705, prefill weights=0
Qwen3-0.6B-4bit	lazy + pinned	0.169 s (0.00 GB res.)	470.86 ms	0.640 s	318.86 tok/s	prefill weights=704/0.34GB upload=68.34ms pinned=704
Qwen3-VL-2B-Instruct-4bit (704 text / 315 skipped vision, 0.97 GB)	eager + pinned	0.285 s (0.90 GB res.)	395.79 ms (-23.2%)	0.681 s	145.15 tok/s	load eval=108.40ms pinned=704 memcpy=706, prefill weights=0
Qwen3-VL-2B-Instruct-4bit	lazy + pinned	0.169 s (0.00 GB res.)	515.05 ms	0.684 s	143.63 tok/s	prefill weights=704/0.97GB upload=105.01ms pinned=704
Meta-Llama-3.1-8B-Instruct-4bit (743 tensors, 4.52 GB)	eager + pinned	0.555 s (4.21 GB res.)	444.79 ms (-46.2% vs lazy+pinned, -49.7% vs lazy+pageable)	1.000 s	43.42 tok/s	load eval=407.59ms staging=313.19ms pinned=743 memcpy=749, prefill weights=0
Meta-Llama-3.1-8B-Instruct-4bit	eager + pageable	0.610 s (4.21 GB res.)	449.13 ms	1.059 s	43.27 tok/s	load eval=459.69ms pageable=743 memcpy=743
Meta-Llama-3.1-8B-Instruct-4bit	lazy + pinned	0.154 s (0.00 GB res.)	826.21 ms	0.980 s	43.19 tok/s	prefill weights=743/4.52GB upload=373.87ms pinned=743 memcpy=749
Meta-Llama-3.1-8B-Instruct-4bit	lazy + pageable	0.151 s (0.00 GB res.)	883.64 ms	1.035 s	43.61 tok/s	prefill weights=743/4.52GB upload=479.22ms pageable=743
Qwen3-30B-A3B-4bit (1,351 tensors, 17.17 GB)	eager + pinned	9.409 s (16.00 GB res.)	634.01 ms (-94.0% vs lazy+pinned, -95.3% vs lazy+pageable)	10.043 s (-26.9% vs lazy+pageable)	61.37 tok/s	load eval=9175.45ms staging=8537.43ms pinned=1351 memcpy=1499, prefill weights=0
Qwen3-30B-A3B-4bit	eager + pageable	13.610 s (16.00 GB res.)	988.03 ms	14.598 s	61.06 tok/s	load eval=13424.21ms pageable=1351 memcpy=1351
Qwen3-30B-A3B-4bit	lazy + pinned	0.188 s (0.00 GB res.)	10,552.30 ms	10.740 s	59.27 tok/s	prefill weights=1351/17.17GB upload=9908.50ms pinned=1351 memcpy=1499
Qwen3-30B-A3B-4bit	lazy + pageable	0.185 s (0.00 GB res.)	13,555.83 ms	13.741 s	33.16 tok/s	prefill weights=1351/17.17GB upload=12985.53ms pageable=1351

Test plan

• cargo test --profile test-fast -p mlxcel-core --features rocm -- weight_materialize --test-threads=1 (10/10 passed, including pinned_upload_chunks_and_wraps_the_slot_ring on ROCm)
• cargo test --profile test-fast -p mlxcel --lib --features rocm -- ttft_profile (7/7 passed)
• cargo test --profile test-fast -p mlxcel --lib --features rocm -- loading:: (371 passed)
• make verify-rocm-overlay (records + drift + ownership mock passed)
• scripts/ci/rocm_smoke.sh (passed, 249.61 tok/s)
• cargo clippy --workspace --all-targets --features rocm -- -D warnings and cargo fmt --all -- --check
• Greedy output byte-identical across eager+pinned, eager+pageable, lazy+pinned, and lazy+pageable on Qwen3-0.6B-4bit, Qwen3-VL-2B-Instruct-4bit, Meta-Llama-3.1-8B-Instruct-4bit, gpt-oss-20b-MXFP4-Q4, and Qwen3-30B-A3B-4bit

Ref #1564, #2474