From 7f4687a912f88669e5fe0124f25ec0c07343a0d7 Mon Sep 17 00:00:00 2001 From: Jeongkyu Shin Date: Tue, 20 Oct 2026 16:47:59 +0900 Subject: [PATCH] perf(rocm): stage weights through pinned slot ring and schedule safetensors on GPU Bring ROCm Load::eval_gpu and load_safetensors to parity with the CUDA path: - Check (cu::is_available() || rocm::is_available()) in the ROCm safetensors.cpp overlay so load_safetensors schedules Load on Device::gpu instead of falling back to Device::cpu. - Replace per-tensor hipHostMalloc + hipLaunchHostFunc(hipHostFree) in mlx/backend/rocm/load.cpp with four reused 64 MiB hipHostMallocPortable slots per device and per-slot hipEventDisableTiming events, with pageable fallback under MLXCEL_ROCM_PINNED_LOAD=0 (or MLXCEL_CUDA_PINNED_LOAD=0). - Export rocm::load_transfer_stats() in load_staging.h and wire it into weight_load_transfer_stats under MLXCEL_BRIDGE_ROCM_BACKEND. - Update ROCm overlay records and PR #2474 technical reports with gfx1151 measurements. --- ...474-pinned-staging-eager-weight-load.en.md | 52 ++- ...474-pinned-staging-eager-weight-load.ko.md | 52 ++- docs/environment-variables.md | 2 +- src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp | 6 + .../overlays/rocm/CORE_RESIDUAL.diff | 27 +- .../mlxcelverse/overlays/rocm/LOCAL_FIXES.md | 1 + src/lib/mlxcelverse/overlays/rocm/README.md | 2 +- .../overlays/rocm/mlx/backend/rocm/load.cpp | 303 +++++++++++++++--- .../rocm/mlx/backend/rocm/load_staging.h | 41 +++ .../overlays/rocm/mlx/io/safetensors.cpp | 5 +- 10 files changed, 412 insertions(+), 79 deletions(-) create mode 100644 src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/load_staging.h diff --git a/TECHNICAL_REPORTS/pr-2474-pinned-staging-eager-weight-load.en.md b/TECHNICAL_REPORTS/pr-2474-pinned-staging-eager-weight-load.en.md index 56a640cdb..a4eebff63 100644 --- a/TECHNICAL_REPORTS/pr-2474-pinned-staging-eager-weight-load.en.md +++ b/TECHNICAL_REPORTS/pr-2474-pinned-staging-eager-weight-load.en.md @@ -19,8 +19,9 @@ MLX `Load` arrays are lazy: `load_safetensors` returns unevaluated arrays, so a | Area | Change | |---|---| | `src/lib/mlxcelverse/overlays/cuda/mlx/backend/cuda/load.cpp`, `load_staging.h` | `Load::eval_gpu` overlay with four pinned 64 MiB slots per device | +| `src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/load.cpp`, `load_staging.h`, `mlx/io/safetensors.cpp` | ROCm `Load::eval_gpu` parity (four 64 MiB `hipHostMallocPortable` slots per device, upload counters, `load_safetensors` on GPU stream) | | `src/lib/mlxcel-core/src/weight_materialize.rs` | `MaterializationScope`, policy parsing, tower-tensor exclusion | -| `src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp`, `mlx_cxx_bridge.h`, `build.rs` | `array_weak_ref` / `array_weak_upgrade` bridge functions | +| `src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp`, `mlx_cxx_bridge.h`, `build.rs` | `array_weak_ref` / `array_weak_upgrade` and CUDA/ROCm `weight_load_transfer_stats` bridge functions | | `src/loading/materialize.rs`, `src/loading/mod.rs` | `load_model` and adapter variants open and finish a scope | | `src/lib/mlxcel-core/src/weights.rs` | weights loaders register returned tensors with the open scope | | `generate.rs`, `engine/direct.rs`, `ttft_profile_tests.rs` | `[TTFT] load` line; upload counters appended to the first-token line | @@ -90,25 +91,50 @@ OLMoE-1B-7B (raw bf16, all 3,072 expert tensors restacked by the loader): peak 1 Correctness: greedy output (`-n 24 --temp 0`) is byte-identical across before, eager + pinned, eager + pageable and lazy; an image request is identical on both binaries. +### 4.1 Measured Results on ROCm (`gfx1151`, AMD Radeon 8060S Graphics) + +Two ROCm-specific gaps were found and resolved while validating this PR on `gfx1151`: +1. **`mlx/io/safetensors.cpp` in the ROCm overlay checked only `cu::is_available()`** when choosing the default `Load` stream, so `load_safetensors` fell back 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. Checking `cu::is_available() || rocm::is_available()` schedules `Load` on `Device::gpu`. +2. **Per-tensor `hipHostMalloc` + `hipLaunchHostFunc(hipHostFree)` in `mlx/backend/rocm/load.cpp`.** Under eager `eval_all`, per-tensor page pinning paid 413 ms per 2 GiB in KFD page-pinning churn on `gfx1151` (standalone HIP benchmark for 16 × 128 MiB = 2 GiB: 617.22 ms total — `hipHostMalloc` 297.26 ms, `pread` 156.37 ms, `hipMemcpyAsync` 47.67 ms, `hipHostFree` 115.91 ms — vs **162.98 ms** with the 4 × 64 MiB `StagingPool`, **3.79× faster**) and kept duplicate pinned staging buffers alive alongside the device weights until batch completion. `mlx/backend/rocm/load.cpp` and `load_staging.h` now use the same 4 × 64 MiB pinned `StagingPool` (`hipHostMallocPortable` + `hipEventDisableTiming`) and `LoadTransferStats` counters as CUDA, with `MLXCEL_ROCM_PINNED_LOAD=0` (or `MLXCEL_CUDA_PINNED_LOAD=0`) selecting the pageable fallback. + +End-to-end warm-cache measurements on `gfx1151` (`--no-chat-template -p "The capital of France is" -n 16 --temp 0 --seed 0 --profile`, 5 prompt tokens): + +| Model | Arm | Model load | First prefill (TTFT) | Load + first token | Decode | Upload counters (`[TTFT]`) | +|---|---|---|---|---|---|---| +| `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%) | 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% vs lazy+pinned, -50% 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` | + +Greedy output (`-n 16 --temp 0 --seed 0`) is byte-identical between `eager` and `lazy` across all tested checkpoints on `gfx1151` (`Qwen3-0.6B-4bit`, `Qwen3-VL-2B-Instruct-4bit`, `Meta-Llama-3.1-8B-Instruct-4bit`, `gpt-oss-20b-MXFP4-Q4`, `Qwen3-30B-A3B-4bit`). + ## 5. Unverified and Open -- **No V100 or PCIe device.** The V100 figures in the issue were not re-measured. GB10 is coherent memory; the DMA cost is identical for pinned and pageable here, so the gain on a PCIe device may differ. -- **Metal and ROCm.** Eager is the default there too (it matches mlx-lm's default load behavior) but was not run. The pinned overlay is CUDA only. -- **Multi-GPU and tensor-parallel.** Not run on this one-GPU host. One process-wide mutex covers every device's pool, so a process driving several GPUs would serialize uploads (LOW). -- **Failure injection.** Pinned-allocation failure was not injected; the fallback was exercised only through `MLXCEL_CUDA_PINNED_LOAD=0`. `ensure_ready` relies on `cudaHostAlloc` leaving `slot.host` null on failure before `cudaFreeHost` (LOW). If `cudaEventRecord` fails after a successful copy the slot is not marked pending, reachable only after a CUDA error that already fails the eval (LOW). +- **No V100 or PCIe device.** The V100 figures in the issue were not re-measured. GB10 and `gfx1151` are coherent/unified memory; the gain on a discrete PCIe device may differ. +- **Metal.** Eager is the default there too (it matches mlx-lm's default load behavior) but was not run. The pinned `Load::eval_gpu` overlay covers CUDA and ROCm. +- **Multi-GPU and tensor-parallel.** Not run on these one-GPU hosts. One process-wide mutex covers every device's pool, so a process driving several GPUs would serialize uploads (LOW). +- **Failure injection.** Pinned-allocation failure was not injected; the fallback was exercised only through `MLXCEL_CUDA_PINNED_LOAD=0` / `MLXCEL_ROCM_PINNED_LOAD=0`. `ensure_ready` relies on `cudaHostAlloc` / `hipHostMalloc` leaving `slot.host` null on failure before freeing (LOW). If `cudaEventRecord` / `hipEventRecord` fails after a successful copy the slot is not marked pending, reachable only after a GPU error that already fails the eval (LOW). - **Remaining first-token cost.** `event_signal_kernel` fires 1,847 times on the first token in every arm, eager included: weights derived from `Load` arrays run on the generation stream at the first token. This is the next candidate inside the remaining 0.53 s and is not addressed here. -- **Largest checkpoint.** The 15 GB qwen checkpoint is the largest that loads on `main` on this host; `DeepSeek-V2-Lite` (30 GB) and `gpt-oss-20b` do not load on either branch. -- **Minor.** Slot events use `cudaEventDisableTiming` without `cudaEventBlockingSync`, so a slot wait spins a core; `MLXCEL_CUDA_PINNED_LOAD` is case-sensitive and does not accept `no`. Pre-existing upstream, not this PR: the pageable path's `malloc` is unchecked and leaks on a read exception, and `safetensors.cpp` adds `offset + data_offsets[1]` without an overflow check. -- Why the pageable read is cheaper at load (eager, 0.95 s) than inside the first `eval` (2.20 s) with the same code was not investigated; it does not affect the default. +- **Minor.** Slot events use `cudaEventDisableTiming` / `hipEventDisableTiming` without blocking-sync flags on the event itself, so a slot wait spins a core on CUDA; `MLXCEL_CUDA_PINNED_LOAD` / `MLXCEL_ROCM_PINNED_LOAD` are case-sensitive and do not accept `no`. Pre-existing upstream, not this PR: the pageable path's `malloc` is unchecked and leaks on a read exception, and `safetensors.cpp` adds `offset + data_offsets[1]` without an overflow check. +- Why the pageable read is cheaper at load (eager, 0.95 s) than inside the first `eval` (2.20 s) on GB10 with the same code was not investigated; it does not affect the default. ## 6. Verification -- `weight_materialize_tests` and `ttft_profile_tests` on CUDA, including a dropped-tensor-is-freed test and `pinned_upload_chunks_and_wraps_the_slot_ring` (352 MiB tensor, six chunks, four-slot ring), added in `8b5179819` after review noted the chunking path was covered only by the 636 MB embedding in the byte-identical output. -- `cargo test --release --features cuda -p mlxcel --lib -- loading::`: 372 pass. -- `cargo clippy` with `-D warnings` on mlxcel-core and mlxcel lib and tests, `cargo fmt`, license header check. -- Reviews: implementation review (one HIGH, fixed in `4a23cfbfd`); security and performance review found no CRITICAL or HIGH items. Checked clean: slot bounds, untrusted checkpoint offsets (the overlay reads the same byte range as upstream), exception safety (`ParallelFileReader::read` joins every batch before throwing), no deadlock from `cudaEventSynchronize`. +- `weight_materialize_tests` (10/10) and `ttft_profile_tests` (7/7) on both CUDA (GB10) and ROCm (`gfx1151`), including a dropped-tensor-is-freed test and `pinned_upload_chunks_and_wraps_the_slot_ring` (352 MiB tensor, six chunks, four-slot ring), added in `8b5179819` after review noted the chunking path was covered only by the 636 MB embedding in the byte-identical output. +- `cargo test --release --features cuda -p mlxcel --lib -- loading::` (372 pass) and `cargo test --profile test-fast --features rocm -p mlxcel --lib -- loading::` (371 pass; 1 expected native-`nvfp4` unsupported test on ROCm). +- `scripts/ci/rocm_smoke.sh` (passed, 234.92 tok/s) and `make verify-rocm-overlay` (records + drift + ownership mock passed). +- `cargo clippy` with `-D warnings` on mlxcel-core and mlxcel lib and tests (CUDA and ROCm), `cargo fmt`, license header check. +- Reviews: implementation review (one HIGH, fixed in `4a23cfbfd`); security and performance review found no CRITICAL or HIGH items. Checked clean: slot bounds, untrusted checkpoint offsets (the overlay reads the same byte range as upstream), exception safety (`ParallelFileReader::read` joins every batch before throwing), no deadlock from `cudaEventSynchronize` / `hipEventSynchronize`. ## 7. Related - Issue #1564 (closed by this PR); `docs/benchmark_results/volta-ttft-fixed-cost-2026-09-01.md` (#1545) measured the V100 baseline and named the three candidates. -- New environment variables: `MLXCEL_WEIGHT_MATERIALIZATION` (`eager` default, `lazy`), `MLXCEL_CUDA_PINNED_LOAD` (`0` disables). `MLXCEL_PROFILE_TTFT` gains a `[TTFT] load` line and upload counters. +- New environment variables: `MLXCEL_WEIGHT_MATERIALIZATION` (`eager` default, `lazy`), `MLXCEL_CUDA_PINNED_LOAD` (`0` disables), `MLXCEL_ROCM_PINNED_LOAD` (`0` disables). `MLXCEL_PROFILE_TTFT` gains a `[TTFT] load` line and upload counters on both CUDA and ROCm. diff --git a/TECHNICAL_REPORTS/pr-2474-pinned-staging-eager-weight-load.ko.md b/TECHNICAL_REPORTS/pr-2474-pinned-staging-eager-weight-load.ko.md index f570e5347..bcb3592a1 100644 --- a/TECHNICAL_REPORTS/pr-2474-pinned-staging-eager-weight-load.ko.md +++ b/TECHNICAL_REPORTS/pr-2474-pinned-staging-eager-weight-load.ko.md @@ -19,8 +19,9 @@ MLX `Load` 배열은 lazy다. `load_safetensors`는 평가되지 않은 배열 | 영역 | 변경 | |---|---| | `src/lib/mlxcelverse/overlays/cuda/mlx/backend/cuda/load.cpp`, `load_staging.h` | 디바이스당 64 MiB pinned slot 4개를 쓰는 `Load::eval_gpu` overlay | +| `src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/load.cpp`, `load_staging.h`, `mlx/io/safetensors.cpp` | ROCm `Load::eval_gpu` 동등 구현 (디바이스당 64 MiB `hipHostMallocPortable` slot 4개, 업로드 카운터, `load_safetensors` GPU stream 배정) | | `src/lib/mlxcel-core/src/weight_materialize.rs` | `MaterializationScope`, 정책 파싱, tower 텐서 제외 | -| `src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp`, `mlx_cxx_bridge.h`, `build.rs` | `array_weak_ref` / `array_weak_upgrade` bridge 함수 | +| `src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp`, `mlx_cxx_bridge.h`, `build.rs` | `array_weak_ref` / `array_weak_upgrade` 및 CUDA/ROCm `weight_load_transfer_stats` bridge 함수 | | `src/loading/materialize.rs`, `src/loading/mod.rs` | `load_model`과 adapter 변형이 scope를 열고 종료 | | `src/lib/mlxcel-core/src/weights.rs` | weights loader가 반환 텐서를 열린 scope에 등록 | | `generate.rs`, `engine/direct.rs`, `ttft_profile_tests.rs` | `[TTFT] load` 줄, 첫 토큰 줄에 업로드 카운터 추가 | @@ -90,25 +91,50 @@ OLMoE-1B-7B(raw bf16, expert 텐서 3,072개 전부를 loader가 재적층): pea 정확성: greedy 출력(`-n 24 --temp 0`)은 변경 전, eager + pinned, eager + pageable, lazy 사이에서 바이트 단위로 동일하고, 이미지 요청도 두 바이너리에서 같다. +### 4.1 ROCm (`gfx1151`, AMD Radeon 8060S Graphics) 측정 결과 + +`gfx1151`에서 PR을 검증하는 과정에서 ROCm 전용 격차 2건을 확인하고 함께 해결했다. +1. **ROCm overlay의 `mlx/io/safetensors.cpp`가 기본 `Load` stream 선택 시 `cu::is_available()`만 검사.** ROCm에서는 `load_safetensors`가 `Device::cpu`(`Load::eval_cpu`)로 빠져 `Load::eval_gpu`를 우회하고, loader가 `Load` 배열로부터 파생 텐서를 만들 때 CPU→GPU stream 간 fence가 발생했다. `cu::is_available() || rocm::is_available()`로 수정해 `Load`를 `Device::gpu`에 스케줄한다. +2. **`mlx/backend/rocm/load.cpp`의 텐서별 `hipHostMalloc` + `hipLaunchHostFunc(hipHostFree)`.** eager `eval_all` 아래에서는 텐서마다 페이지 pinning이 반복되면서 `gfx1151` 기준 2 GiB당 413 ms의 KFD 페이지 pinning 비용이 발생했고(16 × 128 MiB = 2 GiB 단독 HIP 벤치마크: 총 617.22 ms — `hipHostMalloc` 297.26 ms, `pread` 156.37 ms, `hipMemcpyAsync` 47.67 ms, `hipHostFree` 115.91 ms — 대 4 × 64 MiB `StagingPool` 재사용 시 **162.98 ms**, **3.79배 단축**), 배치 완료 전까지 디바이스 가중치와 중복되는 pinned 스테이징 버퍼가 메모리에 남았다. `mlx/backend/rocm/load.cpp`와 `load_staging.h`도 CUDA와 동일하게 4 × 64 MiB pinned `StagingPool`(`hipHostMallocPortable` + `hipEventDisableTiming`)과 `LoadTransferStats` 카운터를 사용하며, `MLXCEL_ROCM_PINNED_LOAD=0`(또는 `MLXCEL_CUDA_PINNED_LOAD=0`)으로 pageable fallback을 선택할 수 있다. + +`gfx1151` 종단간 warm-cache 측정 결과 (`--no-chat-template -p "The capital of France is" -n 16 --temp 0 --seed 0 --profile`, 프롬프트 5토큰): + +| 모델 | Arm | 모델 로드 | 첫 prefill (TTFT) | 로드 + 첫 토큰 | Decode | 업로드 카운터 (`[TTFT]`) | +|---|---|---|---|---|---|---| +| `Qwen3-0.6B-4bit` (텐서 704개, 0.34 GB) | `eager + pinned` | 0.304 s (resident 0.31 GB) | **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 (resident 0.00 GB) | 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개 / 비전 제외 315개, 0.97 GB) | `eager + pinned` | 0.285 s (resident 0.90 GB) | **395.79 ms** (-23%) | 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 (resident 0.00 GB) | 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개, 4.52 GB) | `eager + pinned` | 0.555 s (resident 4.21 GB) | **444.79 ms** (lazy+pinned 대비 -46%, lazy+pageable 대비 -50%) | 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 (resident 4.21 GB) | 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 (resident 0.00 GB) | 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 (resident 0.00 GB) | 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개, 17.17 GB) | `eager + pinned` | **9.409 s** (resident 16.00 GB) | **634.01 ms** (lazy+pinned 대비 -94.0%, lazy+pageable 대비 -95.3%) | **10.043 s** (lazy+pageable 대비 -26.9%) | 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 (resident 16.00 GB) | 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 (resident 0.00 GB) | 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 (resident 0.00 GB) | 13,555.83 ms | 13.741 s | 33.16 tok/s | prefill `weights=1351/17.17GB upload=12985.53ms pageable=1351` | + +`gfx1151`에서 테스트한 모든 체크포인트(`Qwen3-0.6B-4bit`, `Qwen3-VL-2B-Instruct-4bit`, `Meta-Llama-3.1-8B-Instruct-4bit`, `gpt-oss-20b-MXFP4-Q4`, `Qwen3-30B-A3B-4bit`)에서 `eager`와 `lazy`의 greedy 출력(`-n 16 --temp 0 --seed 0`)은 바이트 단위로 동일하다. + ## 5. 미검증 및 미해결 -- **V100/PCIe 디바이스 없음.** 이슈의 V100 수치는 재측정하지 않았다. GB10은 coherent memory라 pinned와 pageable의 DMA 비용이 같으므로 PCIe 디바이스의 이득은 다를 수 있다. -- **Metal, ROCm.** 기본값이 eager인 것은 같고(mlx-lm 기본 로드 동작과 일치) 실행하지 않았다. pinned overlay는 CUDA 전용이다. -- **Multi-GPU, tensor-parallel.** GPU 1개 호스트라 실행하지 못했다. 프로세스 전역 mutex 하나가 모든 디바이스 pool을 감싸므로 여러 GPU를 구동하는 프로세스는 업로드가 직렬화된다(LOW). -- **실패 주입.** pinned 할당 실패는 주입하지 않았고 fallback은 `MLXCEL_CUDA_PINNED_LOAD=0`으로만 실행했다. `ensure_ready`는 `cudaHostAlloc` 실패 시 `slot.host`가 null로 남는다는 점에 의존한다(LOW). 복사 성공 후 `cudaEventRecord`가 실패하면 slot이 pending으로 표시되지 않는데, 이미 eval을 실패시키는 CUDA 오류 이후에만 도달한다(LOW). +- **V100/PCIe 디바이스 없음.** 이슈의 V100 수치는 재측정하지 않았다. GB10과 `gfx1151`은 coherent/unified memory이므로 개별 PCIe 디바이스의 이득은 다를 수 있다. +- **Metal.** 기본값이 eager인 것은 같고(mlx-lm 기본 로드 동작과 일치) 실행하지 않았다. pinned `Load::eval_gpu` overlay는 CUDA와 ROCm을 모두 지원한다. +- **Multi-GPU, tensor-parallel.** 단일 GPU 호스트라 실행하지 못했다. 프로세스 전역 mutex 하나가 모든 디바이스 pool을 감싸므로 여러 GPU를 구동하는 프로세스는 업로드가 직렬화된다(LOW). +- **실패 주입.** pinned 할당 실패는 주입하지 않았고 fallback은 `MLXCEL_CUDA_PINNED_LOAD=0` / `MLXCEL_ROCM_PINNED_LOAD=0`으로만 실행했다. `ensure_ready`는 `cudaHostAlloc` / `hipHostMalloc` 실패 시 `slot.host`가 null로 남는다는 점에 의존한다(LOW). 복사 성공 후 `cudaEventRecord` / `hipEventRecord`가 실패하면 slot이 pending으로 표시되지 않는데, 이미 eval을 실패시키는 GPU 오류 이후에만 도달한다(LOW). - **남은 첫 토큰 비용.** `event_signal_kernel`이 eager 포함 모든 arm의 첫 토큰에서 1,847회 실행된다. `Load` 배열에서 파생된 가중치가 첫 토큰에서 generation stream으로 실행되기 때문이며, 남은 0.53 s 안의 다음 후보로 이번 PR에서는 다루지 않았다. -- **최대 체크포인트.** 이 호스트에서 `main`이 로드하는 가장 큰 체크포인트는 15 GB qwen이다. `DeepSeek-V2-Lite`(30 GB)와 `gpt-oss-20b`는 두 브랜치 모두 로드되지 않는다. -- **기타.** slot event는 `cudaEventBlockingSync` 없이 `cudaEventDisableTiming`만 써서 slot 대기가 코어 하나를 spin한다. `MLXCEL_CUDA_PINNED_LOAD`는 대소문자를 구분하고 `no`를 받지 않는다. 이 PR 이전부터 upstream에 있는 문제: pageable 경로의 `malloc`은 검사하지 않고 read 예외 시 누수되며, `safetensors.cpp`는 `offset + data_offsets[1]`을 오버플로 검사 없이 더한다. -- 같은 코드인데 pageable read가 첫 `eval` 안(2.20 s)보다 로드 시점(eager, 0.95 s)에 더 싼 이유는 조사하지 않았다. 기본값에는 영향이 없다. +- **기타.** slot event는 event 자체의 blocking-sync 플래그 없이 `cudaEventDisableTiming` / `hipEventDisableTiming`만 써서 CUDA에서는 slot 대기가 코어 하나를 spin한다. `MLXCEL_CUDA_PINNED_LOAD` / `MLXCEL_ROCM_PINNED_LOAD`는 대소문자를 구분하고 `no`를 받지 않는다. 이 PR 이전부터 upstream에 있는 문제: pageable 경로의 `malloc`은 검사하지 않고 read 예외 시 누수되며, `safetensors.cpp`는 `offset + data_offsets[1]`을 오버플로 검사 없이 더한다. +- 같은 코드인데 GB10에서 pageable read가 첫 `eval` 안(2.20 s)보다 로드 시점(eager, 0.95 s)에 더 싼 이유는 조사하지 않았다. 기본값에는 영향이 없다. ## 6. 검증 -- CUDA에서 `weight_materialize_tests`와 `ttft_profile_tests`: 버려진 텐서가 해제되는지 검증하는 테스트와 `pinned_upload_chunks_and_wraps_the_slot_ring`(352 MiB 텐서, chunk 6개, slot 4개 ring)를 포함한다. 후자는 chunking 경로가 바이트 동일 출력의 636 MB embedding으로만 검증된다는 리뷰 지적 후 `8b5179819`에서 추가됐다. -- `cargo test --release --features cuda -p mlxcel --lib -- loading::`: 372개 통과. -- mlxcel-core와 mlxcel lib/tests에 `cargo clippy -D warnings`, `cargo fmt`, 라이선스 헤더 검사. -- 리뷰: 구현 리뷰(HIGH 1건, `4a23cfbfd`에서 수정), 보안/성능 리뷰는 CRITICAL/HIGH 없음. 확인된 항목: slot 경계, 신뢰할 수 없는 체크포인트 offset(overlay는 upstream과 같은 바이트 범위를 읽음), 예외 안전성(`ParallelFileReader::read`가 던지기 전에 모든 batch를 join), `cudaEventSynchronize`로 인한 deadlock 없음. +- CUDA(GB10)와 ROCm(`gfx1151`) 양쪽에서 `weight_materialize_tests`(10/10)와 `ttft_profile_tests`(7/7): 버려진 텐서가 해제되는지 검증하는 테스트와 `pinned_upload_chunks_and_wraps_the_slot_ring`(352 MiB 텐서, chunk 6개, slot 4개 ring)를 포함한다. 후자는 chunking 경로가 바이트 동일 출력의 636 MB embedding으로만 검증된다는 리뷰 지적 후 `8b5179819`에서 추가됐다. +- `cargo test --release --features cuda -p mlxcel --lib -- loading::`(372개 통과) 및 `cargo test --profile test-fast --features rocm -p mlxcel --lib -- loading::`(371개 통과, ROCm 미지원 native `nvfp4` 1건 제외). +- `scripts/ci/rocm_smoke.sh`(통과, 234.92 tok/s) 및 `make verify-rocm-overlay`(records + drift + ownership mock 통과). +- mlxcel-core와 mlxcel lib/tests(CUDA 및 ROCm)에 `cargo clippy -D warnings`, `cargo fmt`, 라이선스 헤더 검사. +- 리뷰: 구현 리뷰(HIGH 1건, `4a23cfbfd`에서 수정), 보안/성능 리뷰는 CRITICAL/HIGH 없음. 확인된 항목: slot 경계, 신뢰할 수 없는 체크포인트 offset(overlay는 upstream과 같은 바이트 범위를 읽음), 예외 안전성(`ParallelFileReader::read`가 던지기 전에 모든 batch를 join), `cudaEventSynchronize` / `hipEventSynchronize`로 인한 deadlock 없음. ## 7. 관련 항목 - 이슈 #1564(이 PR이 종결). `docs/benchmark_results/volta-ttft-fixed-cost-2026-09-01.md`(#1545)가 V100 기준선을 측정하고 세 후보를 제시했다. -- 새 환경 변수: `MLXCEL_WEIGHT_MATERIALIZATION`(기본 `eager`, `lazy`), `MLXCEL_CUDA_PINNED_LOAD`(`0`이면 비활성). `MLXCEL_PROFILE_TTFT`에 `[TTFT] load` 줄과 업로드 카운터가 추가됐다. +- 새 환경 변수: `MLXCEL_WEIGHT_MATERIALIZATION`(기본 `eager`, `lazy`), `MLXCEL_CUDA_PINNED_LOAD`(`0`이면 비활성), `MLXCEL_ROCM_PINNED_LOAD`(`0`이면 비활성). `MLXCEL_PROFILE_TTFT`에 `[TTFT] load` 줄과 업로드 카운터가 CUDA/ROCm 모두에 추가됐다. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index d38b6be09..fe4da34d3 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -770,7 +770,7 @@ recommended as normal deployment settings. | `MLXCEL_SAMPLING_GUMBEL` | `0`/`false`/`off`/`no` disable (case-insensitive); any other value or unset enables | on | Softmax-free Gumbel-max categorical sampling kernel (issue #900), on by default on Metal, CUDA and ROCm (the HIP port is lablup/mlxcel#2064). On the no-filter stochastic path (`temperature > 0` with no top-k, top-p, or min-p) the fused sampler replaces `mlx::core::random::categorical` with one index-carrying max reduction over the vocabulary: adding i.i.d. Gumbel(0,1) noise to `logits / temperature` and taking the argmax draws exactly from `softmax(logits / temperature)`, so the normalization pass over 32K-152K entries disappears. Noise comes from a stateless in-kernel Philox-4x32-10 counter keyed on `(row, element)`, so the sampled id is a pure function of the launch key and the logits and does not move when the launcher splits a row across threadgroups to keep a small batch wide. One launch covers the whole `[B, vocab]` batch. **This changes the RNG stream**: at an equal seed the sampled tokens differ from the pre-#900 `categorical` streams, because the two consume the shared MLX key sequence differently. The distribution is identical (chi-square goodness of fit against exact softmax is a committed test), and greedy (`temperature == 0`) output is byte-identical either way. Set to a falsy value to restore `categorical` exactly, which is the way to reproduce a token stream recorded before this landed. Filtered configurations (top-k / top-p / min-p) never reach the kernel and are bit-identical with the switch either way; token bias is an upstream `-inf` logit mask that the kernel reproduces exactly, and XTC-active requests never reach this kernel (XTC lives inside the stock filter chain). Read once per process, so set it before starting `mlxcel` or `mlxcel-server`. Inert on CPU-only builds and under `MLXCEL_DEVICE=cpu`. A/B with `examples/gumbel_sampling_microbench.rs`. | | `MLXCEL_SAMPLING_REJECTION` | `0`/`false`/`off`/`no` disable (case-insensitive); any other value or unset enables | on | Sorting-free top-k / top-p / min-p sampling by dual-pivot rejection (issue #901), on by default on Metal, CUDA and ROCm (the HIP port is lablup/mlxcel#2064). On the filtered stochastic path the fused sampler replaces the `argpartition` top-k, the `argsort` + `cumsum` top-p and the min-p mask, plus the trailing `random::categorical`, with the softmax passes and one custom kernel; since the untempered-chain fix the kernel reads two probability rows (the untempered row resolves every filter, the tempered row weights the draw), which coincide at `temperature 1`. **Routing is restricted to the configurations measured faster.** The kernel replaces a *sort*, so it is routed only where the stock chain sorts, which is when top-p is active; and when top-k is active as well, only at vocabularies up to 32768 (152064 on a ROCm build). Measured on M1 Ultra over three repetitions of vocab {32K, 64K, 152K} x batch {1, 4, 8}: top-p alone 1.28x to 2.35x at every cell, top-k alone 0.31x to 0.97x at every cell, min-p alone 0.47x to 0.88x at 152K, top-k with top-p 1.27x to 1.64x at 32K but 0.71x to 0.83x at 152K. gfx1151 re-measured the top-k plus top-p row at vocab {32K, 64K, 128256, 151936, 152064} (lablup/mlxcel#2157, [results](benchmark_results/rocm-rejection-joint-cap-gfx1151-2026-10-07.md)): the microbenchmark's batches 4 and 8 win at every vocabulary while its synthetic batch-1 cell does not clear 1.0 above 32768, but end-to-end decode with `--temperature 0.8 --top-k 40 --top-p 0.95` measured the kernel 1.15x faster on Meta-Llama-3.1-8B-Instruct-4bit (vocab 128256) and 1.12x on Qwen2.5-7B-Instruct-4bit (152064), so a ROCm build routes the combination up to 152064 (a build-flag value, not a runtime backend check). Metal and CUDA keep 32768. The mechanism is the kernel's round count, which the microbenchmark reports: top-p accepts in one or two rounds because a probability-weighted draw lands in the high-mass head, while top-k needs the draw to land in the top forty of the whole vocabulary and takes two to seven rounds, growing with the vocabulary. Every round is another full-row sweep, and the loop runs on one threadgroup per row because the shrinking interval is carried across sweeps and MLX custom kernels have no grid-wide barrier, so the kernel pays single-core bandwidth per round while `argpartition` and the stock min-p mask run across the whole GPU. top-k alone, min-p alone, top-k with min-p, and top-k with top-p above the cap therefore stay on the stock chain and are bit-identical to before; the decline is announced once at INFO with the numbers behind it. On M1 Ultra vocab 65536 has not been measured for the top-k plus top-p combination and is excluded rather than interpolated; on gfx1151 it was measured and routes. All three filters are threshold filters on the probability value, so the filtered support is `{p > low}` for a single scalar per row; the kernel finds that scalar without sorting, by rejection sampling on a shrinking interval. Each round draws a candidate from the current proposal with a fixed-order block scan, takes the candidate's probability as one pivot and the midpoint of the bracket as the other, reduces `(count, mass)` above both in one vocabulary sweep, and accepts the candidate as soon as it passes every active filter test. The bisection pivot is the midpoint of the IEEE-754 **bit patterns**, not the arithmetic mean: pure integer arithmetic that a fast-math backend with flush-to-zero cannot corrupt for subnormal probabilities, and it isolates a single float in at most 32 rounds at any magnitude. min-p is folded into the initial interval and costs no round at all. One launch covers the whole `[B, vocab]` batch, and the kernel reads `{top_k, top_p, min_p}` per row so rows with different values need no second launch. **This changes the RNG stream**: at an equal seed the sampled tokens differ from the pre-#901 streams, because the kernel consumes the shared MLX key sequence differently from `random::categorical`. **One semantic difference is intentional.** When top-k AND top-p are both active, the stock chain masks to the top-k set and *renormalises* before applying top-p, so its mass target is `top_p * Z_k` with `Z_k` the top-k mass; the kernel applies both tests to the untruncated distribution, so its target is `top_p * total` and its support is a superset. The two agree unless some token's exclusive cumulative mass falls in `(top_p * Z_k, top_p * total]`. Every other configuration (top-k alone, top-p alone, min-p alone, top-k+min-p, top-p+min-p) is exact, and support equality against the `argpartition` mask including ties, plus chi-square goodness of fit against the renormalised truncated distribution, are committed tests. Greedy (`temperature == 0` or `top_k == 1`) is byte-identical either way, and the no-filter path is unaffected (it belongs to `MLXCEL_SAMPLING_GUMBEL`). **The routed path never synchronizes.** Both decode drivers are software pipelines that build step n+1 and `async_eval` it before reading step n, so a sampler that evaluates anything internally drains the queue inside the caller's build phase and collapses the pipeline. The first cut of this kernel read the per-row converged flags back to host inside `fused_sample` and measured 1.7x slower end to end on Qwen3-0.6B, against an op-level matrix that scored the same configuration 1.14x to 1.17x faster. The production path therefore evaluates nothing: the round cap is unreachable by construction (the bit-space bracket bounds the loop at 31 rounds), and the converged flags are stashed and inspected on a later call once `array::is_available()` reports the launch has landed, which in a decode loop is by the next token. A row that somehow exhausted the cap returns its own argmax, which is always inside the filtered support, so the degradation would be one greedy draw rather than an invalid token; the event is counted in `mlxcel_core::rejection_cap_overflow_rows()` and announced once at INFO. Which path a run took is announced at INFO, once per distinct outcome kind, so a benchmark arm can be proven from a log. Set to a falsy value to restore the `argpartition` chain exactly, which is the way to reproduce a token stream recorded before this landed. Read once per process, so set it before starting `mlxcel` or `mlxcel-server`. Inert on CPU-only builds and under `MLXCEL_DEVICE=cpu`. A/B with `examples/rejection_sampling_microbench.rs`. | | `MLXCEL_CUDA_F16_NORMALIZE` | `1`/`true`/`on`/`yes` enable; `0`/`false`/`off`/`no` disable | **on** below sm_80, off (opt-in) at sm_80 and later | **CUDA only.** Load-time bf16 -> f16 normalization of weights for the single-dtype decode graph (issue #636). The default is split by compute capability because the hardware is. **Ampere and later (sm_80+)**: off, unchanged from #636. bf16 has native ALUs there, the merged `overlays/cuda/mlx/dtype.cpp` promotion patch already yields a 0-AsType single-dtype bf16 decode graph, and f16 offers no measured throughput gain (qwen2.5-0.5b-bf16 decode 208 vs 207 tok/s on GB10) while narrowing dynamic range. **Volta and Turing (sm_70, sm_75)**: on, and quantized checkpoints are included. Those parts have no bf16 ALU, no bf16 tensor-core MMA atom, and no cuBLAS bf16 GEMM, so every bf16 operand is converted before it can execute. Measured on a Tesla V100 with qwen3.8-27B-4bit: prefill 85.3 s to 4.8 s and decode 118.1 to 49.2 ms per token, with `qmm_naive` spending 199.1 s against 11.4 s over an identical 994 launches. On a quantized checkpoint the packed planes stay u32 and only the bf16 side-data moves, which is why quantized models are in scope here and out of scope above. Set to `0` to opt out on pre-Ampere and keep bf16. The f16-fragile exception differs by architecture. At sm_80 and later it is the full conservative list from #732 (gemma, cohere/command-r, apertus, gpt-oss, and any config with a nonzero softcap or `logit_scale`). Below sm_80 every entry on that list was measured on a V100 and only **apertus** stays excluded: its xIELU squares its input, and f16 produces NaN at all 508 scored positions from the first token where bf16 scores 24.27. Gemma (+0.09% perplexity over 10208 tokens, decode 6.12 to 9.79 tok/s), gpt-oss (token-weighted -0.02%, with the two window sizes disagreeing in sign) and Cohere (which ships F16 weights already and so runs in f16 regardless, scoring 20.25 and 13.15) all measured clean, as did the generic softcap and `logit_scale` triggers, which is what a cap bounding a value rather than growing it predicts. Metal / Apple Silicon is governed by the separate always-on policy and is untouched. Drafter checkpoints loaded through `mlxcel-core` (DFlash, Muse Glimmer, Inkling MTP, Qwen 3.5 MTP) follow the same rule (issue #1782): before that, every drafter was converted to f16 unconditionally, and on Ampere and later CUDA, where the target keeps bf16, MLX promoted the bf16 residual stream times the f16 drafter weights to float32, so the whole drafter forward ran in f32 and re-upcast all of its weights on every round (about 240 launches and 30 to 38 ms of device time per verify round on GB10 with `qwen3.5-4b-4bit` and its DFlash drafter, the largest fixed term in the round). | -| `MLXCEL_CUDA_PINNED_LOAD` | `0`/`false`/`off` disable; anything else or unset keeps the default | on | **CUDA only.** Uploads checkpoint tensors (`Load::eval_gpu`, `overlays/cuda/mlx/backend/cuda/load.cpp`) through four reused 64 MiB pinned host staging slots instead of one pageable `malloc` buffer per tensor (issue #1564). The read into one slot overlaps the copy out of the previous one, the copy runs at the pinned rate, and the per-tensor `malloc`, first-touch page faults and `free` are gone. A tensor larger than one slot is copied in slot-sized chunks, so `cudaMemcpyAsync` calls exceed the tensor count by the number of extra chunks. The 256 MiB of pinned memory is allocated on the first upload and kept for the process. If that allocation fails, mlxcel prints one warning and uses the pageable path for the rest of the process instead of failing the load. `0` forces the pageable path, which is MLX's upstream behaviour; the uploaded bytes are identical either way. | +| `MLXCEL_CUDA_PINNED_LOAD` | `0`/`false`/`off` disable; anything else or unset keeps the default | on | **CUDA and ROCm (`MLXCEL_ROCM_PINNED_LOAD` on ROCm, with `MLXCEL_CUDA_PINNED_LOAD` also honored as a fallback).** Uploads checkpoint tensors (`Load::eval_gpu`, `overlays/cuda/mlx/backend/cuda/load.cpp` and `overlays/rocm/mlx/backend/rocm/load.cpp`) through four reused 64 MiB pinned host staging slots instead of one per-tensor buffer (issue #1564). The read into one slot overlaps the copy out of the previous one, the copy runs at the pinned rate, and the per-tensor `malloc`/`hipHostMalloc`, first-touch page faults and `free`/`hipHostFree` callbacks are gone. A tensor larger than one slot is copied in slot-sized chunks, so `cudaMemcpyAsync`/`hipMemcpyAsync` calls exceed the tensor count by the number of extra chunks. The 256 MiB of pinned memory is allocated on the first upload and kept for the process. If that allocation fails, mlxcel prints one warning and uses the pageable path for the rest of the process instead of failing the load. `0` forces the pageable path; the uploaded bytes are identical either way. | | `MLXCEL_CUDA_F16_FRAGILE` | `1`/`true`/`on`/`yes` enable; unset keeps the exclusion | off (opt-in) | **CUDA below sm_80 only.** Converts a family on the f16-fragile exception list anyway. That list is inherited from #732, which adopted it for Ampere and later where f16 gave, in that commit's words, "no AsType reduction and no throughput gain": with no upside, excluding a family on a suspicion costs nothing, and that PR's test plan ran no fragile family in f16. Below sm_80 the list is down to apertus, which is the one entry measurement supported, so this flag now means converting a family that produced NaN at every scored position on this hardware. It remains for a checkpoint whose activations you have reason to believe differ, and for taking the measurement on a family this list does not yet know. | | `MLXCEL_KEEP_BF16` | set to any value to enable; unset disables | off | **Measurement instrument, not a supported configuration.** Skips load-time bf16 -> f16 conversion for every checkpoint, on every backend, so a bf16 model loads uniformly bf16. It exists to A/B a dtype policy against the mixture it replaces, and it is what established the Gemma3n result below. Gemma3n used to keep only its language MLP bf16 and convert the rest to f16, leaving one checkpoint holding both dtypes; MLX promotes f16 with bf16 to f32, so every boundary between the kept MLP and its f16 neighbours promoted and the decode path read promoted weights the whole way down. Loading uniformly bf16 measured, on M5 Max with `gemma3n-e4b-bf16` at a 273-token prompt with an image and 128 generated, decode 38.02 to 47.48 tok/s and prefill 2143 to 2371 tok/s, against mlx-vlm 0.6.17 at 48.72 and 1207: decode moves from 78% of the reference to 97% while prefill stays about 1.95x ahead. Gemma3n now loads uniformly bf16 by default (see `gemma3n_bf16_key`), so this flag no longer changes that family. **It is not evidence that bf16 beats f16.** The control run is the point: on families with no forced-bf16 subset the two dtypes land within noise of each other on the same hardware, `qwen2.5-0.5b-bf16` at 401.8 against 397.3 tok/s and `llama-3.1-8b-bf16` at 33.07 against 33.02. What costs is the mixture, not the dtype, so every other family keeps converting to f16 and setting this flag for them buys nothing while narrowing nothing. Drafter checkpoints loaded through `mlxcel-core` honor it the same way (issue #1782). Read per load. | diff --git a/src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp b/src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp index cb905c9f2..ed2533646 100644 --- a/src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp +++ b/src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp @@ -12,6 +12,8 @@ #include "kv_inplace_write.h" // in-place decode KV row write (#1959). #ifdef MLXCEL_BRIDGE_CUDA_BACKEND #include "mlx/backend/cuda/load_staging.h" // weight upload counters (#1564). +#elif defined(MLXCEL_BRIDGE_ROCM_BACKEND) +#include "mlx/backend/rocm/load_staging.h" // weight upload counters (#1564). #endif namespace mlx_cxx { @@ -419,8 +421,12 @@ std::unique_ptr array_weak_upgrade(const MlxWeakArray& w) { } bool weight_load_transfer_stats(rust::Slice out) { +#if defined(MLXCEL_BRIDGE_CUDA_BACKEND) || defined(MLXCEL_BRIDGE_ROCM_BACKEND) #ifdef MLXCEL_BRIDGE_CUDA_BACKEND const auto s = mlx::core::cu::load_transfer_stats(); +#else + const auto s = mlx::core::rocm::load_transfer_stats(); +#endif const uint64_t fields[] = { s.tensors, s.bytes, diff --git a/src/lib/mlxcelverse/overlays/rocm/CORE_RESIDUAL.diff b/src/lib/mlxcelverse/overlays/rocm/CORE_RESIDUAL.diff index 9e9c69b10..8c8efbe7c 100644 --- a/src/lib/mlxcelverse/overlays/rocm/CORE_RESIDUAL.diff +++ b/src/lib/mlxcelverse/overlays/rocm/CORE_RESIDUAL.diff @@ -30,7 +30,7 @@ # overlay-sha256: 26584a93a3720250aa59c6dd9f091863f8c047691ce55f9a988a01da390ab02a mlx/fast.cpp # overlay-sha256: b1a5366b31412246256ea56f9114bfe7df2fcb93329301ce7ec924f7200c7695 mlx/fast.h # overlay-sha256: 8839267c3bc469b6a3f7d67c6028701a1a3bd02d0f09b6ac31b15407204d2653 mlx/fast_primitives.h -# overlay-sha256: f5fccd165887c2860c04fecf3fcf6073397789081183c2bcaaacad8f8f4d8e9e mlx/io/safetensors.cpp +# overlay-sha256: 78008a41de160e14cb1e2963d9ce0fb005653298a2cbf0e121ae38d1c71ccf0f mlx/io/safetensors.cpp # overlay-sha256: 555e31c93e5acd1678862d53820ebafaa0dc4a9074c61501b355ce198a76b81a mlx/ops.cpp # overlay-sha256: 5af152117800d10f56baf0540418aa8a2d80804106563d4a7190a1a357531c5f mlx/primitives.cpp # overlay-sha256: 5223221a04503b7f9ab5215fc594e57884b1388a1ca84e6b2ad33edb451d2525 mlx/stream.cpp @@ -38,7 +38,7 @@ # note mlx/backend/common/compiled.cpp: LOCAL_FIXES.md item 1: resolves the pin/fork conflict, keeping upstream's negative_strides tracking and the fork's constant-input skipping. # note mlx/backend/gpu/primitives.cpp: LOCAL_FIXES.md item 24: DynamicSliceUpdate copies through copy_gpu (donates only when is_donatable holds) and drops the fork's graph_active() forced donation. # note mlx/fast_primitives.h: LOCAL_FIXES.md item 1: resolves the pin/fork conflict, keeping upstream's compile_options and the fork's output_input_aliases on fast::CustomKernel. -# note mlx/io/safetensors.cpp: LOCAL_FIXES.md item 1: resolves the pin/fork conflict, keeping upstream's empty-tensor guard (as an early continue) and the fork's ROCm staged_write. +# note mlx/io/safetensors.cpp: LOCAL_FIXES.md items 1 and 49: resolves the pin/fork conflict (empty-tensor guard + ROCm staged_write) and schedules load_safetensors on the GPU stream when rocm::is_available(). --- a/mlx/allocator.h +++ b/mlx/allocator.h @@ -196,7 +196,26 @@ } // namespace mlx::core::fast --- a/mlx/io/safetensors.cpp +++ b/mlx/io/safetensors.cpp -@@ -274,13 +274,11 @@ +@@ -6,6 +6,7 @@ + #include + + #include "mlx/backend/cuda/cuda.h" ++#include "mlx/backend/rocm/rocm.h" + #include "mlx/io.h" + #include "mlx/io/load.h" + #include "mlx/ops.h" +@@ -126,7 +127,9 @@ + throw std::runtime_error(msg.str()); + } + +- auto stream = cu::is_available() ? to_stream(s) : to_stream(s, Device::cpu); ++ auto stream = (cu::is_available() || rocm::is_available()) ++ ? to_stream(s) ++ : to_stream(s, Device::cpu); + + uint64_t jsonHeaderLength = 0; + // This is the same limit as in the original Rust Safetensors code. +@@ -274,13 +277,11 @@ out_stream->write(reinterpret_cast(&header_len), 8); out_stream->write(header.c_str(), header_len); for (auto& [key, arr] : a) { @@ -213,7 +232,7 @@ #ifdef MLX_ROCM_HOST_STAGE // On ROCm the array lives in BAR-mapped VRAM: writing straight from // arr.data() makes the CPU read it uncached at ~30 MB/s. staged_write -@@ -291,7 +289,6 @@ +@@ -291,7 +292,6 @@ } #endif out_stream->write(arr.data(), arr.nbytes()); diff --git a/src/lib/mlxcelverse/overlays/rocm/LOCAL_FIXES.md b/src/lib/mlxcelverse/overlays/rocm/LOCAL_FIXES.md index 33778e496..6f157a039 100644 --- a/src/lib/mlxcelverse/overlays/rocm/LOCAL_FIXES.md +++ b/src/lib/mlxcelverse/overlays/rocm/LOCAL_FIXES.md @@ -76,3 +76,4 @@ The fork branched from upstream on 2026-06-26 (`39886de4`). Moving its ROCm code 47. **Move-only allocation ownership.** Upstream changed `allocator::malloc` and `array::set_data` to transfer `allocator::Data`, removed the free wrapper and made donation require owned storage. The new core overlay `mlx/allocator.h` adds only the ROCm forward declaration and friendship corresponding to CUDA's existing asynchronous allocation entry point. `mlx/backend/rocm/allocator.h` and `mlx/backend/rocm/allocator.cpp` return genuinely owned `Data` from the public `rocm::malloc_async`; the underlying `RocmAllocator` still returns raw `Buffer` as required by the allocator interface. This preserves ordinary donation and the existing encoder's shared array retention until batch completion; destructing `Data` reaches the same ROCm cache/deferred-free logic exactly once. `mlx/backend/rocm/gemms/naive_gemm.hip` wraps slab views in borrowed `Data` with the existing shell-only deleter, so they cannot donate externally owned slab storage. `mlx/backend/rocm/moe_swiglu.cpp` and `mlx/backend/rocm/reduce/reduce_utils.hpp` remove the obsolete free/deleter arguments when transferring owned allocations. Retargeting `mlx/backend/common/compiled.{h,cpp}` brings in upstream's `Data`-returning allocation callback and avoids borrowed wrappers. The host mock regression in `scripts/mlxcelverse/allocator_ownership_test.{cpp,sh}` compiles the actual header and async handoff, checks move-only/owned/borrowed/empty allocations, pending-batch shared retention, and cleanup on exceptions before and after enqueue. It does not execute HIP or establish GPU completion correctness. 48. **New primitive and capability contracts.** `mlx/backend/rocm/fence.cpp` updates `Fence::wait` to take the target `uint32_t value` (`fence->event.wait(value)`) and `Fence::update` to return `fence->count`, matching upstream `mlx/fence.h` at `243bdf09`. `mlx/backend/rocm/scaled_dot_product_attention.cpp` accepts upstream's `has_sinks` argument and excludes sinks from WMMA admission, matching the existing execution guard; vector and scalar flash keep their sinks support, and unsupported sinks plus LSE requests take the graph fallback (or reject `force_fused`). `mlx/backend/rocm/primitives.cpp` defines `GatedDeltaUpdateVJP::use_fallback` and `eval_gpu`: the capability check always selects upstream's differentiable graph, like the existing forward fallback, while the explicit unsupported-kernel stub satisfies linking without claiming a native implementation. `mlx/backend/rocm/binary.hip` adds `FloorDivide::eval_gpu` using the existing binary dispatch. Its operation in `mlx/backend/rocm/device/binary_ops.hpp`, also shared by `DivMod`, now floors signed and floating negative quotients instead of truncating them, matching target CUDA semantics. `mlx/backend/rocm/compiled.cpp` makes the same correction for generated kernels and keeps integral division exact rather than rounding 64-bit values through float. `mlx/backend/rocm/fft.hip` documents the audited `odd_out` support: the real transform uses the exact output shape for the hipFFT plan length, cache key and inverse scale, so odd inverse-real lengths are not reconstructed as an even `2 * (input - 1)`. No FFT algorithm change is needed. Actual ROCm compile, graph-fallback execution, negative floor-division numerics and odd-length FFT execution remain to be checked on an AMD host; this retarget host has no HIP toolchain or AMD device. +49. **GPU-stream `load_safetensors`, reused pinned host staging pool, and upload counters in `Load::eval_gpu`.** In `mlx/io/safetensors.cpp`, `load_safetensors(Reader, StreamOrDevice)` checked only `cu::is_available()` when choosing the default `Load` stream, so ROCm 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()`. In `mlx/backend/rocm/load.cpp`, `Load::eval_gpu` previously called `hipHostMalloc(&out_ptr, nbytes, hipHostMallocDefault)` and queued `hipLaunchHostFunc(encoder.stream(), hip_host_free_callback, out_ptr)` for every tensor. Under eager weight materialization (lablup/mlxcel#1564), `eval_all` runs all kept checkpoint tensors back-to-back, so per-tensor `hipHostMalloc`/`hipHostFree` paid 413 ms per 2 GiB in KFD page-pinning churn on gfx1151 (617 ms total for 16 x 128 MiB vs 163 ms with reused slots, 3.8x faster) and kept up to a batch's worth of duplicate pinned staging buffers live alongside the device weights. `mlx/backend/rocm/load.cpp` now uploads through four reused 64 MiB pinned host slots per HIP device (`hipHostMallocPortable` + per-slot `hipEvent_t` with `hipEventDisableTiming`), bounding pinned staging memory at 256 MiB per device and eliminating the per-tensor host callback, with the synchronous pageable `malloc` + `hipMemcpy` path kept as fallback when `MLXCEL_ROCM_PINNED_LOAD=0` (or `MLXCEL_CUDA_PINNED_LOAD=0`) or when slot allocation fails. `mlx/backend/rocm/load_staging.h` exposes `rocm::load_transfer_stats()`, wired into `weight_load_transfer_stats` in `src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp` under `MLXCEL_BRIDGE_ROCM_BACKEND` so `MLXCEL_PROFILE_TTFT=1` and `weight_materialize_tests::pinned_upload_chunks_and_wraps_the_slot_ring` inspect the same upload counters on ROCm as on CUDA. Applies to the fork (for `safetensors.cpp`); staging pool kept in mlxcelverse (lablup/mlxcel#1564). diff --git a/src/lib/mlxcelverse/overlays/rocm/README.md b/src/lib/mlxcelverse/overlays/rocm/README.md index cf73a967b..08b9a46cb 100644 --- a/src/lib/mlxcelverse/overlays/rocm/README.md +++ b/src/lib/mlxcelverse/overlays/rocm/README.md @@ -6,7 +6,7 @@ mlxcelverse is the name for everything mlxcel builds on top of upstream MLX: per The tree mirrors the MLX source tree. Every file is a whole-file overlay, copied over the fetched MLX checkout by `mlx_apply_source_overlays` in `../../CMakeLists.txt`, and only when `MLX_BUILD_ROCM` is on. Metal and CUDA builds never see these files. A ROCm build also compiles the one overlay in `../cpu/` that every build gets, the CPU backend's `mlx/backend/cpu/quantized.cpp` (float accumulation in the scalar quantized matmul, lablup/mlxcel#2248). That file is not part of this overlay, the fork does not change it, and it is not counted below; its delta and refresh rule are in its header and in `../../CMakeLists.txt`. -- `mlx/backend/rocm/**`: the ROCm backend (109 files). Upstream MLX has no such directory, so these files never conflict with a pin bump; they only need to keep compiling against the API of the pinned commit. +- `mlx/backend/rocm/**`: the ROCm backend (110 files). Upstream MLX has no such directory, so these files never conflict with a pin bump; they only need to keep compiling against the API of the pinned commit. - 16 MLX core files (`CMakeLists.txt`, `mlx/CMakeLists.txt`, `mlx/backend/common/{buffer_cache.h,compiled.cpp,compiled.h}`, `mlx/backend/gpu/primitives.cpp`, `mlx/{compile,device,fast,ops,primitives,stream}.cpp`, `mlx/{allocator,fast,fast_primitives}.h`, `mlx/io/safetensors.cpp`): the upstream file at the pinned commit merged with the fork's ROCm hooks, which are mostly `#ifdef MLX_USE_ROCM` blocks. Files in this directory other than `CMakeLists.txt` and `mlx/` (this README, `UPSTREAM`, `LOCAL_FIXES.md`, `CORE_RESIDUAL.diff`) are records and are not copied. `CORE_RESIDUAL.diff` is generated by `scripts/mlxcelverse/rocm_overlay.py drift --write`: it is what the core files carry beyond the MLX pin plus the fork's own change to each file, with a note per file naming the `LOCAL_FIXES.md` entry behind it. `make verify-rocm-overlay` (part of `make verify` and `make verify-rocm`) checks that these records agree with the tree and with the MLX pin. diff --git a/src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/load.cpp b/src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/load.cpp index b3b3e0728..61e0d3133 100644 --- a/src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/load.cpp +++ b/src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/load.cpp @@ -1,11 +1,37 @@ // Copyright © 2025 Apple Inc. +// +// Modified by mlxcel (lablup/mlxcel#1564): +// - Weights are uploaded through a small, reused pool of four 64 MiB pinned +// host staging slots per device instead of calling `hipHostMalloc` and +// queuing a `hipLaunchHostFunc(..., hip_host_free_callback, ...)` per tensor. +// The read into a slot overlaps the host-to-device copy out of the previous +// one, per-tensor `hipHostMalloc`/`hipHostFree` KFD page-pinning churn is +// gone, and pinned staging memory stays bounded at 256 MiB per device even +// when `eval_all` materializes an entire checkpoint at once. +// - The synchronous pageable path is kept as the fallback: it is taken when +// `MLXCEL_ROCM_PINNED_LOAD=0` (or `MLXCEL_CUDA_PINNED_LOAD=0`), or for the +// rest of the process once a pinned slot allocation fails. An async H2D copy +// from pageable memory is avoided because driver-internal staging over a +// non-coherent link (TB5 eGPU) can stall the stream. +// - Process-wide counters (`load_staging.h`) report the host staging leg and +// the copy back-pressure separately, matching the CUDA backend. #include +#include +#include +#include +#include +#include +#include +#include +#include #include +#include #include "mlx/backend/rocm/allocator.h" #include "mlx/backend/rocm/device.h" #include "mlx/backend/rocm/kernel_utils.hpp" +#include "mlx/backend/rocm/load_staging.h" #include "mlx/backend/rocm/utils.h" #include "mlx/primitives.h" @@ -28,68 +54,253 @@ void swap_endianness(uint8_t* data_bytes, size_t N) { } } -void hip_host_free_callback(void* ptr) { - (void)hipHostFree(ptr); +void swap_endianness_inplace(void* data, size_t size, size_t itemsize) { + auto* bytes = reinterpret_cast(data); + switch (itemsize) { + case 2: + swap_endianness<2>(bytes, size); + break; + case 4: + swap_endianness<4>(bytes, size); + break; + case 8: + swap_endianness<8>(bytes, size); + break; + } +} + +using Clock = std::chrono::steady_clock; + +uint64_t elapsed_ns(Clock::time_point start) { + return static_cast( + std::chrono::duration_cast( + Clock::now() - start) + .count()); +} + +struct Counters { + std::atomic tensors{0}; + std::atomic bytes{0}; + std::atomic pinned_tensors{0}; + std::atomic pageable_tensors{0}; + std::atomic memcpy_calls{0}; + std::atomic read_ns{0}; + std::atomic wait_ns{0}; + std::atomic total_ns{0}; +}; + +Counters& counters() { + static Counters c; + return c; +} + +// Four 64 MiB slots keep two reads and two copies in flight. 64 MiB is two of +// `ParallelFileReader`'s 32 MiB batches, so a full-slot read is split across +// its reader pool instead of running on one thread. A tensor larger than a slot +// is uploaded in slot-sized chunks. +constexpr size_t kSlotBytes = size_t(64) << 20; +constexpr int kSlotCount = 4; + +bool pinned_load_enabled() { + static bool enabled = [] { + const char* v = std::getenv("MLXCEL_ROCM_PINNED_LOAD"); + if (!v) { + v = std::getenv("MLXCEL_CUDA_PINNED_LOAD"); + } + return !(v && (std::strcmp(v, "0") == 0 || std::strcmp(v, "false") == 0 || + std::strcmp(v, "off") == 0)); + }(); + return enabled; } } // namespace namespace mlx::core { +namespace { + +// One pool per HIP device: an event can only be recorded on a stream of the +// device it was created on. Slots are allocated on first use and live for the +// rest of the process, like MLX's own allocator; they are deliberately never +// freed, so no destructor runs after the HIP runtime tears down at exit. +class StagingPool { + public: + bool ensure_ready() { + if (state_ == State::Ready) { + return true; + } + if (state_ == State::Unavailable) { + return false; + } + for (int i = 0; i < kSlotCount; ++i) { + Slot slot; + const char* failed = nullptr; + if (hipHostMalloc(&slot.host, kSlotBytes, hipHostMallocPortable) != + hipSuccess || + slot.host == nullptr) { + failed = "hipHostMalloc"; + } else if ( + hipEventCreateWithFlags(&slot.done, hipEventDisableTiming) != + hipSuccess) { + failed = "hipEventCreateWithFlags"; + } + if (failed) { + (void)hipGetLastError(); + if (slot.host) { + (void)hipHostFree(slot.host); + } + release(); + state_ = State::Unavailable; + std::cerr << "[mlxcel] pinned weight staging unavailable (" << failed + << " failed for " << (kSlotBytes >> 20) + << " MiB slots); loading weights through pageable memory." + << std::endl; + return false; + } + slots_.push_back(slot); + } + state_ = State::Ready; + return true; + } + + void upload( + io::Reader& reader, + size_t offset, + size_t nbytes, + char* dst, + hipStream_t stream, + bool swap, + size_t itemsize) { + auto& c = counters(); + size_t done = 0; + while (done < nbytes) { + size_t n = std::min(kSlotBytes, nbytes - done); + Slot& slot = slots_[next_]; + next_ = (next_ + 1) % slots_.size(); + if (slot.pending) { + auto wait_start = Clock::now(); + CHECK_HIP_ERROR(hipEventSynchronize(slot.done)); + c.wait_ns += elapsed_ns(wait_start); + slot.pending = false; + } + auto read_start = Clock::now(); + reader.read(static_cast(slot.host), n, offset + done); + if (swap) { + swap_endianness_inplace(slot.host, n / itemsize, itemsize); + } + c.read_ns += elapsed_ns(read_start); + CHECK_HIP_ERROR(hipMemcpyAsync( + dst + done, slot.host, n, hipMemcpyHostToDevice, stream)); + CHECK_HIP_ERROR(hipEventRecord(slot.done, stream)); + slot.pending = true; + c.memcpy_calls += 1; + done += n; + } + } + + private: + enum class State { Unknown, Ready, Unavailable }; + struct Slot { + void* host{nullptr}; + hipEvent_t done{nullptr}; + bool pending{false}; + }; + + void release() { + for (auto& slot : slots_) { + (void)hipEventDestroy(slot.done); + (void)hipHostFree(slot.host); + } + slots_.clear(); + } + + State state_{State::Unknown}; + std::vector slots_; + size_t next_{0}; +}; + +struct StagingPools { + std::mutex mutex; + std::unordered_map> by_device; +}; + +StagingPools& staging_pools() { + // Leaked on purpose, see `StagingPool`. + static auto* pools = new StagingPools(); + return *pools; +} + +} // namespace + +namespace rocm { + +LoadTransferStats load_transfer_stats() { + auto& c = counters(); + return LoadTransferStats{ + c.tensors.load(), + c.bytes.load(), + c.pinned_tensors.load(), + c.pageable_tensors.load(), + c.memcpy_calls.load(), + c.read_ns.load(), + c.wait_ns.load(), + c.total_ns.load(), + }; +} + +} // namespace rocm + void Load::eval_gpu(const std::vector& inputs, array& out) { + auto start = Clock::now(); + auto& c = counters(); auto& encoder = rocm::get_command_encoder(stream()); auto size = out.size(); auto nbytes = size * out.itemsize(); out.set_data(mlx::core::rocm::malloc_async(nbytes, encoder)); - // Stage through PINNED host memory. An async H2D copy from pageable memory is - // unreliable on a discrete GPU over a non-coherent link (TB5 eGPU): the - // driver must internally stage it, which can stall the stream (queue stuck, - // GPU shows busy, the eval's sync never returns). Pinned memory DMAs directly - // and lets the copy actually run asynchronously. - void* out_ptr = nullptr; - if (hipHostMalloc(&out_ptr, nbytes, hipHostMallocDefault) != hipSuccess || - out_ptr == nullptr) { - // Fallback: pageable + synchronous copy (still correct, just slower). - out_ptr = malloc(nbytes); - reader_->read(static_cast(out_ptr), nbytes, offset_); - if (swap_endianness_) { - switch (out.itemsize()) { - case 2: - swap_endianness<2>(reinterpret_cast(out_ptr), size); - break; - case 4: - swap_endianness<4>(reinterpret_cast(out_ptr), size); - break; - case 8: - swap_endianness<8>(reinterpret_cast(out_ptr), size); - break; - } + c.tensors += 1; + c.bytes += nbytes; + + if (nbytes > 0 && pinned_load_enabled()) { + auto& pools = staging_pools(); + std::lock_guard lock(pools.mutex); + auto& pool = pools.by_device[encoder.device().hip_device()]; + if (!pool) { + pool = std::make_unique(); + encoder.device().make_current(); + } + if (pool->ensure_ready()) { + pool->upload( + *reader_, + offset_, + nbytes, + gpu_ptr(out), + encoder.stream(), + swap_endianness_, + out.itemsize()); + c.pinned_tensors += 1; + c.total_ns += elapsed_ns(start); + return; } - (void)hipMemcpy(gpu_ptr(out), out_ptr, nbytes, hipMemcpyHostToDevice); - free(out_ptr); - return; } + + // Fallback: pageable staging buffer + synchronous H2D copy. An async H2D copy + // from pageable memory is unreliable on a discrete GPU over a non-coherent + // link (TB5 eGPU): driver-internal staging can stall the stream. + auto read_start = Clock::now(); + auto out_ptr = malloc(nbytes); reader_->read(static_cast(out_ptr), nbytes, offset_); if (swap_endianness_) { - switch (out.itemsize()) { - case 2: - swap_endianness<2>(reinterpret_cast(out_ptr), size); - break; - case 4: - swap_endianness<4>(reinterpret_cast(out_ptr), size); - break; - case 8: - swap_endianness<8>(reinterpret_cast(out_ptr), size); - break; - } + swap_endianness_inplace(out_ptr, size, out.itemsize()); + } + c.read_ns += elapsed_ns(read_start); + if (nbytes > 0) { + CHECK_HIP_ERROR( + hipMemcpy(gpu_ptr(out), out_ptr, nbytes, hipMemcpyHostToDevice)); + c.memcpy_calls += 1; } - (void)hipMemcpyAsync( - gpu_ptr(out), - out_ptr, - nbytes, - hipMemcpyHostToDevice, - encoder.stream()); - (void)hipLaunchHostFunc(encoder.stream(), hip_host_free_callback, out_ptr); + free(out_ptr); + c.pageable_tensors += 1; + c.total_ns += elapsed_ns(start); } } // namespace mlx::core diff --git a/src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/load_staging.h b/src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/load_staging.h new file mode 100644 index 000000000..1565c38af --- /dev/null +++ b/src/lib/mlxcelverse/overlays/rocm/mlx/backend/rocm/load_staging.h @@ -0,0 +1,41 @@ +// Copyright © 2025 Apple Inc. + +// mlxcelverse: counters for the weight upload path in `Load::eval_gpu` +// (`load.cpp` in this directory, lablup/mlxcel#1564). Mirrors +// `mlx/backend/cuda/load_staging.h` for the ROCm backend so +// `weight_load_transfer_stats` in `src/lib/mlxcel-core/cpp/mlx_cxx_ext.cpp` +// reports the same host staging and copy back-pressure legs on ROCm builds. + +#pragma once + +#include + +namespace mlx::core::rocm { + +struct LoadTransferStats { + // `Load` primitives evaluated on a HIP stream, and their bytes. + uint64_t tensors; + uint64_t bytes; + // How many of those took the pinned staging path and how many the pageable + // one (disabled by `MLXCEL_ROCM_PINNED_LOAD=0` / `MLXCEL_CUDA_PINNED_LOAD=0`, + // or after a pinned allocation failed). + uint64_t pinned_tensors; + uint64_t pageable_tensors; + // `hipMemcpyAsync` / `hipMemcpy` calls issued. Equal to `tensors` on the + // pageable path; a pinned upload issues one per staging chunk, so a tensor + // larger than one chunk adds more than one. + uint64_t memcpy_calls; + // Host staging leg: malloc (pageable path) plus the read into the staging + // buffer, wall-clock nanoseconds on the evaluating thread. + uint64_t read_ns; + // Host time blocked waiting for a pinned staging slot to drain. + uint64_t wait_ns; + // Wall-clock nanoseconds inside `Load::eval_gpu`, both legs plus overhead. + uint64_t total_ns; +}; + +// A consistent-enough snapshot of the process-wide counters (each field is +// read atomically; the set is not read under one lock). +LoadTransferStats load_transfer_stats(); + +} // namespace mlx::core::rocm diff --git a/src/lib/mlxcelverse/overlays/rocm/mlx/io/safetensors.cpp b/src/lib/mlxcelverse/overlays/rocm/mlx/io/safetensors.cpp index b1f82cebd..bb8d59170 100644 --- a/src/lib/mlxcelverse/overlays/rocm/mlx/io/safetensors.cpp +++ b/src/lib/mlxcelverse/overlays/rocm/mlx/io/safetensors.cpp @@ -6,6 +6,7 @@ #include #include "mlx/backend/cuda/cuda.h" +#include "mlx/backend/rocm/rocm.h" #include "mlx/io.h" #include "mlx/io/load.h" #include "mlx/ops.h" @@ -126,7 +127,9 @@ SafetensorsLoad load_safetensors( throw std::runtime_error(msg.str()); } - auto stream = cu::is_available() ? to_stream(s) : to_stream(s, Device::cpu); + auto stream = (cu::is_available() || rocm::is_available()) + ? to_stream(s) + : to_stream(s, Device::cpu); uint64_t jsonHeaderLength = 0; // This is the same limit as in the original Rust Safetensors code.