Skip to content

perf: keep raw dit image tensors on cpu until vae preprocess. - #2261

Open
chuanhaoyu wants to merge 1 commit into
xLLM-AI:mainfrom
chuanhaoyu:feature/dit-keep-raw-images-cpu
Open

perf: keep raw dit image tensors on cpu until vae preprocess.#2261
chuanhaoyu wants to merge 1 commit into
xLLM-AI:mainfrom
chuanhaoyu:feature/dit-keep-raw-images-cpu

Conversation

@chuanhaoyu

@chuanhaoyu chuanhaoyu commented Aug 18, 2026

Copy link
Copy Markdown

Description

Keep raw DiT uint8 image tensors on CPU in DiTForwardInput::to(), and only move compute tensors to NPU.

VAEImageProcessor::resize() still uses a CPU Lanczos kernel, so moving images / images_list / mask_images / control_image / last_images to NPU early caused an extra NPU->CPU->NPU round-trip. This change avoids that copy for large and multi-image inputs such as Qwen Image Edit Plus.

Prompt embeddings, latents, masked_image_latents, and prompt_audio still follow the original .to(device) behavior.

Related Issues

N/A

Change Type

  • Bug fix
  • New feature
  • Performance improvement
  • Refactor
  • Documentation
  • Test
  • Build or CI

Pull Request Checklist

Thank you for contributing to xLLM. Before requesting review, please make sure the following items are complete.

PR Title and Commit Messages

  • The PR title and each commit message follow the xLLM commit format: <type>: <subject>.

Allowed types: feat, bugfix, docs, test, refactor, chore, style, revert, perf, model, build, release.
The subject should use clear English, start with a verb, include at least 4 words, and end with ..

Pre-commit Checks

  • I have installed pre-commit by running pip install pre-commit or an equivalent command.
  • I have installed the hooks with pre-commit install.
  • I have run pre-commit run --all-files and fixed any reported issues.

If you are unsure how to set up pre-commit, see the pre-commit documentation.

Self Review

  • I have self-reviewed the code according to .agents/skills/code-review/references/custom-code-style.md, especially code written or assisted by AI.
  • I have rebased this PR onto the latest main branch.

Build and Test Coverage

  • Tests have been added or updated as needed.
  • CUDA: python setup.py build test has passed on a CUDA machine.
  • NPU: python setup.py build test has passed on an NPU machine.
  • MLU: python setup.py build test has passed on an MLU machine.

Reviewer Notes

This PR only changes xllm/core/runtime/dit_forward_params.h.

Validated on NPU with Qwen Image Edit Plus:

  • single-image and multi-image requests succeed
  • raw images_list stays on CPU before VAE preprocess
  • condition_img / vae_img still land on NPU after preprocess
  • preprocess latency does not regress

NPU Lanczos resize is out of scope. Flux Fill / Flux Control / Wan I2V share the same raw-image fields, but dedicated end-to-end smoke was not run because those model services were not available in this environment.

@chuanhaoyu
chuanhaoyu force-pushed the feature/dit-keep-raw-images-cpu branch from c6097f3 to 892ab4e Compare August 18, 2026 11:07
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