Vllm Omni Npu Model Runner Upgrade

vLLM-Omni NPU Model Runner Upgrade Skill

Overview

This skill guides the process of upgrading vllm-omni's NPU model runners to align with the latest vllm-ascend codebase while preserving omni-specific enhancements. The NPU runners are designed to run omni multimodal models (like Qwen3-Omni, Bagel, MiMoAudio) on Ascend NPUs.

File Structure

NPU Model Runner Files

vllm-omni/vllm_omni/platforms/npu/worker/
├── __init__.py
├── npu_model_runner.py           # OmniNPUModelRunner (base class)
├── npu_ar_model_runner.py        # NPUARModelRunner (autoregressive)
├── npu_ar_worker.py              # AR worker
├── npu_generation_model_runner.py # NPUGenerationModelRunner (diffusion/non-AR)
└── npu_generation_worker.py      # Generation worker

GPU Reference Files (for omni-specific logic sync)

vllm-omni/vllm_omni/worker/
├── __init__.py
├── gpu_model_runner.py           # OmniGPUModelRunner
├── gpu_ar_model_runner.py        # GPUARModelRunner
├── gpu_ar_worker.py
├── gpu_generation_model_runner.py
├── gpu_generation_worker.py
├── mixins.py
├── base.py
└── gpu_memory_utils.py

vllm-ascend Reference Files

vllm-ascend/vllm_ascend/worker/
├── model_runner_v1.py            # NPUModelRunner (base class to copy from)
├── npu_input_batch.py
├── block_table.py
├── pcp_utils.py
└── worker.py

Inheritance Hierarchy

                    GPUModelRunner (vllm)
                         |
        +----------------+----------------+
        |                                 |
  OmniGPUModelRunner              NPUModelRunner (vllm-ascend)
  (vllm_omni/worker)              (vllm_ascend/worker)
        |                                 |
        +----------- OmniNPUModelRunner --+
                     (multiple inheritance)
                            |
            +---------------+---------------+
            |                               |
    NPUARModelRunner            NPUGenerationModelRunner
    (autoregressive)            (non-autoregressive/diffusion)

Omni-Specific Comment Markers

Omni-specific logic is marked with comment blocks:

# -------------------------------------- Omni-new -------------------------------------------------
# ... omni-specific code ...
# -------------------------------------- Omni-new -------------------------------------------------

Or simpler variations:

#  -------------------------------------- Omni-new -------------------------------------------------
#  ------------------------------------------------------------------------------------------------

Important:

• Always preserve and add these markers when modifying code.
• The reference documents (references/omni-specific-blocks.md) may not be up-to-date. Always grep for Omni-new in the GPU implementations to find the authoritative list of omni-specific blocks.
• When you discover new omni-specific code that is not documented in the references, please update the reference files.

Key Methods Requiring Attention

OmniNPUModelRunner (npu_model_runner.py)

NPUARModelRunner (npu_ar_model_runner.py)

NPUGenerationModelRunner (npu_generation_model_runner.py)

Upgrade Workflow

Step 1: Preparation

1. Identify target versions(Use gh cli to check):

   • We're using vllm-omni main branch
   • Check the last release of vllm-omni
   • Target vllm-ascend version(Just directly use the local latest vllm-ascend code)

2. Check GPU-side changes (since last release):

   cd /root/vllm-workspace/vllm-omni
   git log --oneline --since="<last-release-date>" -- vllm_omni/worker/

3. Read latest vllm-ascend code:

   • We don't track vllm-ascend changes - just directly use the latest code from /root/vllm-workspace/vllm-ascend/vllm_ascend/worker/model_runner_v1.py
   • Copy the relevant methods and re-insert omni-specific blocks

Step 2: Analyze Omni-Specific Logic

For each NPU model runner file:

1. Extract existing omni-specific blocks:

   grep -n "Omni-new" vllm_omni/platforms/npu/worker/npu_model_runner.py

2. Document each omni block:

   • Which method it belongs to
   • What functionality it provides
   • Dependencies on other omni code

Step 3: Update Base Class (OmniNPUModelRunner)

Note: Always check the GPU implementation gpu_model_runner.py for any new omni logic not yet documented in references.

1. Read the latest vllm-ascend NPUModelRunner.load_model

2. Copy the method, keeping the structure

3. Re-insert omni-specific logic (check GPU gpu_model_runner.py for authoritative list):

   • Replace CUDAGraphWrapper with ACLGraphWrapper
   • Keep talker_mtp initialization
   • Preserve buffer allocations for talker
   • Check for any new omni blocks added since last sync

4. Update _dummy_run:

   • Copy from vllm-ascend
   • Compare with GPU _dummy_run for omni-specific blocks
   • Re-insert all Omni-new marked code from GPU version

5. Update _model_forward:

   • Keep the omni wrapper logic
   • Update NPU-specific parts (graph params, SP all-gather)
   • Check GPU version for any new omni logic

Step 4: Update AR Model Runner

1. Compare with GPU gpu_ar_model_runner.py for any new omni features

2. Copy execute_model from vllm-ascend

3. Re-insert omni blocks (reference references/omni-specific-blocks.md, but note it may be incomplete):

   • IMPORTANT: Always check the GPU implementation gpu_ar_model_runner.py for all Omni-new marked code blocks
   • The reference doc may not include newly added omni logic - treat it as a starting point, not exhaustive
   • When discovering new omni code blocks, please update references/omni-specific-blocks.md
   • Common omni blocks include but are not limited to: KV transfer, multimodal outputs, sampling_metadata handling, etc.

4. Update sample_tokens (also compare with GPU implementation):

   • Compare with gpu_ar_model_runner.py's sample_tokens method
   • Identify all Omni-new marked code blocks
   • Ensure NPU version includes all omni-specific logic