Skip to content

Keep HF cuFile/hipFile weights on device when memory permits - #2

Merged
cyyever merged 1 commit into
trimfrom
perf/hf-direct-device-load
Oct 22, 2026
Merged

cyyever merged 1 commit into
trimfrom
perf/hf-direct-device-load

Conversation

@cyyever

@cyyever cyyever commented Oct 22, 2026

Copy link
Copy Markdown
Collaborator

Overview

Keep oversized original HF checkpoint tensors on the device during cuFile/hipFile compatibility-mode loading when memory permits. The existing model loader still handles names, slicing, fusion and layout conversions; no deployment-specific checkpoint conversion is required.

Claims

  • The shared direct-storage iterator tries a device buffer before assembling a large tensor on CPU. Successful device staging avoids the previous GPU-to-CPU-to-GPU round trip.
  • Retain bounded device staging with CPU assembly when free memory is insufficient or allocation raises OutOfMemoryError. Only allocation failures trigger fallback; read failures still abort.
  • Leave headroom of max(256 MiB, one quarter of the incoming tensor size). This is a conservative heuristic, not a guarantee that downstream model transformations cannot run out of memory.
  • Existing destination callbacks from Add shared cuFile/hipFile weight destinations #1 remain preferred. P2P behavior is unchanged.

Validation

  • 18 tests passed across test_direct_io.py and test_hipfile.py, including sufficient memory, low memory, allocation races, byte correctness and tail chunks.
  • Command: PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 uv run --no-project .venv/bin/python -m pytest --confcutdir=tests/model_executor/model_loader -o addopts= tests/model_executor/model_loader/test_direct_io.py tests/model_executor/model_loader/test_hipfile.py -q.
  • Unit tests use CPU-backed reader/allocator doubles; they do not establish GPU correctness or speedup.
  • Relevant pre-commit hooks, including mypy, run on commit.
  • Hardware benchmark and model-output validation pending. Keep this PR draft until the actual HF Qwen path is measured.

Details

Large compatibility-mode weights previously always passed through a CPU tensor, even with spare VRAM. This incremental path reuses existing backend readers and chunked writes without adding a model-specific mapping or dependency.

This is device staging, not zero-copy parameter adoption: the model loader may still copy or transform the loaded GPU tensor. It does not implement safetensors-native DLPack integration or arbitrary final-destination mapping. No end-to-end acceleration is claimed yet.

Review must consider consumer memory spikes, fragmentation and the possibility that asynchronous model loading retains source tensors longer than the generator iteration. Allocation fallback does not catch OOM inside model transformations.

Duplicate searches found no open BakeLab PR and no matching upstream model-loader PR; upstream results concerned KV offload and unrelated kernels. This builds on #1.

Remaining work

  • Compare baseline and changed HF load times on the same hardware/cache conditions.
  • Record peak CPU/GPU memory and verify loaded values on real hardware.
  • Validate cuFile and hipFile paths and model output before merge.
  • Decide from measurements whether further native-buffer/DLPack or final-destination work is warranted.

Pull Request Checklist
  • I used vLLM's /pr-checklist skill.
  • AI assistance was used during creation (OpenAI Codex).
  • Design Fit: shared iterator, existing readers, no model-specific mapping.
  • Testing and Validation: unit tests passed; hardware/model validation pending.
  • Code Quality and Style: applicable hooks run, allocation policy documented.
  • Pull Request Contents: scope, limitations and pending evidence explicit.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: Yuanyuan Chen <cyyever@outlook.com>
@cyyever
cyyever marked this pull request as ready for review October 22, 2026 04:55
@cyyever
cyyever merged commit d47b84e into trim Oct 22, 2026
5 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant