[RFC]: Support NVIDIA Alpamayo 1.5 (Reasoning VLA for Autonomous Driving) in vLLM-Omni

Motivation

Alpamayo 1.5 is a 10B-parameter Vision-Language-Action (VLA) driving-reasoning model released by NVIDIA in 2026 (weights: nvidia/Alpamayo-1.5-10B). Its core properties make it a natural fit for vLLM-Omni:

1. VLM backbone + Action Expert Mixture-of-Transformers (MoT) — structurally very close to the already-supported BAGEL (BagelForConditionalGeneration / BagelPipeline). Both follow the pattern "AR backbone generates tokens and produces a KV-cache, a second expert/DiT reuses the same KV-cache to run a diffusion denoising loop".
2. The backbone is Qwen3-VL-8B-Instruct (Qwen3VLForConditionalGeneration), which is already natively supported by upstream vLLM. The Action Expert is built from a deep-copied Qwen3 text config with embed_tokens removed — essentially the same recipe BAGEL uses for Qwen2MoTForCausalLM / Qwen2MoTDecoderLayer.
3. It uses Flow Matching (10-step Euler integration) to denoise in action space and produces a 64-waypoint, 10 Hz, 6.4 s horizon trajectory (xyz + rotation), instead of image/video/audio. This is a new output modality for vLLM-Omni's DiffusionOutput.
4. It supports multi-camera (variable camera count) inputs, navigation conditioning, Chain-of-Causation (CoC) reasoning, classifier-free guidance, and general VQA — all of which can be layered onto mechanisms already available in BAGEL / Qwen3-Omni today.

vLLM-Omni has shipped stable BAGEL support since 0.16.0, and the SIG-Omni roadmap explicitly pushes vLLM-Omni toward being a general "any-to-any multimodal" serving engine. Adding Alpamayo 1.5 brings the following benefits:

• Extends vLLM-Omni coverage from image / video / audio / text to embodied / autonomous-driving trajectories, paving the way for a generic DiffusionOutput(trajectory=...) / OmniOutputModality.TRAJECTORY abstraction that future VLA models (π0, RT-2, OpenVLA, ...) can reuse.
• Reuses existing infrastructure from BAGEL: dual-KV / CFG KV-cache injection, the MoT decoder layer, DiffusionPipelineProfilerMixin, AutoWeightsLoader, DiffusersPipelineLoader. The estimated new code is modest (< 2k LOC).
• Gives NVIDIA's PhysicalAI ecosystem and the vLLM community a high-throughput, disaggregation-ready reference implementation of a driving VLA.

---

Proposed Change

1. Architectural comparison: Alpamayo 1.5 vs. BAGEL

Dimension	BAGEL (supported)	Alpamayo 1.5 (this RFC)
Backbone	Qwen2 (MoT via Qwen2MoTDecoderLayer)	Qwen3-VL-8B-Instruct (Qwen3VLForConditionalGeneration)
Vision encoder	SigLIP-NaViT (SiglipNaViTWrapper)	Built-in Qwen3-VL ViT (multi-camera packed sequence)
"Expert"	gen-expert layers inside the same MoT	Separate AutoModel.from_config(text_config), embed_tokens removed, shares backbone KV-cache
Diffusion space	VAE latent (image, 16 channels)	Action space (xyz + rotation, not an image)
Diffusion algorithm	Rectified Flow / FM, ~50 steps	Flow Matching, 10-step Euler
CFG	cfg_text_scale + cfg_img_scale, multi-KV-cache	inference_guidance_weight + nav-CFG, multi-KV-cache
Output	PIL.Image.Image	(pred_xyz, pred_rot) + optional CoC text
Main entry	BagelPipeline.forward(req)	sample_trajectories_from_data_with_vlm_rollout / generate_text

Conclusion. We can fully reuse BAGEL's AR + Diffusion two-stage design. The only new pieces are:

• A new AR runner inheriting from Qwen3-VL that handles trajectory special-token placeholders (<|traj_history|>, <|traj_future|>, ...);
• A new diffusion pipeline that runs Flow Matching + the Action Expert;
• A new output modality (trajectory) plus its post-processing and OpenAI-compatible response wiring.

2. Directory layout

Following docs/contributing/model/adding_diffusion_model.md and aligned with BAGEL:

vllm_omni/
├── model_executor/models/
│   ├── alpamayo1_5/
│   │   ├── __init__.py
│   │   └── alpamayo1_5.py                # Alpamayo15ForConditionalGeneration
│   │                                     #   (vLLM-Omni-style AR runner + mm processor)
│   └── registry.py                       # +"Alpamayo15ForConditionalGeneration"
└── diffusion/
    ├── registry.py                       # +"Alpamayo15Pipeline" and pre/post process
    └── models/alpamayo1_5/
        ├── __init__.py
        ├── pipeline_alpamayo1_5.py       # Alpamayo15Pipeline (forward / diffuse / CFG-nav)
        ├── alpamayo1_5_transformer.py    # Alpamayo15ActionExpert (vLLM-ized Qwen3 layer stack
        │                                 #   with embed_tokens removed, KV-cache injection)
        ├── action_in_proj.py             # action_in_proj (noise + timestep -> hidden)
        ├── action_out_proj.py            # action_out_proj (hidden -> vector field)
        ├── action_space.py               # wraps ActionSpace.action_to_traj
        ├── flow_matching.py              # FlowMatching.sample (Euler, CFG)
        └── delta_tokenizer.py            # discrete trajectory tokenizer (history/future)

• The upstream code relies on hydra.utils.instantiate + Omegaconf dicts. We will not introduce a hydra dependency. Instead, we construct the submodules explicitly from config.json fields, exactly like vllm_omni/diffusion/models/bagel/pipeline_bagel.py does with _VaeCfg / _VitCfg.
• The Qwen3-VL backbone is reused directly from vllm.model_executor.models.qwen3_vl.Qwen3VLForConditionalGeneration — we do not fork it.

3. Registration

Add to vllm_omni/diffusion/registry.py under _DIFFUSION_MODELS:

"Alpamayo15Pipeline": (
    "alpamayo1_5",
    "pipeline_alpamayo1_5",
    "Alpamayo15Pipeline",
),

And to _DIFFUSION_POST_PROCESS_FUNCS:

"Alpamayo15Pipeline": "get_alpamayo15_post_process_func",

The post-process function packages (pred_xyz, pred_rot) into a DiffusionOutput(output=..., modality="trajectory").

Add to vllm_omni/model_executor/models/registry.py under _OMNI_MODELS:

"Alpamayo15ForConditionalGeneration": (
    "alpamayo1_5",
    "alpamayo1_5",
    "Alpamayo15ForConditionalGeneration",
),

This must match the architectures field of HF config.json. The upstream model_type is alpamayo_reasoning_vla; if architectures is not set in the released config, we add an alias in vllm_omni/patch.py.

4. Pipeline design (Alpamayo15Pipeline)

Mirroring the upstream sample_trajectories_from_data_with_vlm_rollout:

class Alpamayo15Pipeline(nn.Module, DiffusionPipelineProfilerMixin):
    _PROFILER_TARGETS = ["vlm.generate", "expert.forward", "diffuse", "action_decode"]

    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
        # 1. Download weights from HF; read config.json / generation_config.json
        # 2. Build Qwen3VLForConditionalGeneration (vllm backbone) as self.vlm
        # 3. Deep-copy self.vlm.config.text_config and apply expert_cfg overrides;
        #    build Alpamayo15ActionExpert (self.expert); do NOT copy embed_tokens
        # 4. Build FlowMatching / ActionSpace / action_in_proj / action_out_proj
        # 5. Set up weights_sources (DiffusersPipelineLoader.ComponentSource)

    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
        # 0. Parse req.multi_modal_data:
        #    - images: N multi-view camera frames (variable)
        #    - ego_history: {"xyz": [...], "rot": [...]}
        #    - navigation: Optional[str]  (reverse of nav_utils.remove_nav_text)
        #    - extra_args: num_traj_samples, num_traj_sets, use_cfg, top_p, temperature, ...

        # 1. Trajectory fusion: replace <|traj_history|> placeholders in input_ids
        #    with the output of hist_traj_tokenizer.encode(...) (vLLM-Omni port of
        #    fuse_traj_tokens).

        # 2. AR rollout. Two execution modes:
        #    a) In-process: call self.vlm.generate(..., return_dict_in_generate=True)
        #       -> past_key_values, sequences, rope_deltas
        #    b) Disaggregated (recommended for online serving): run AR inside
        #       Alpamayo15ForConditionalGeneration, ship KV-cache handles to the
        #       diffusion worker through OmniConnector, same as BAGEL.

        # 3. Build expert position_ids / attention_mask (Qwen3-VL's 3-component RoPE),
        #    equivalent to _build_expert_pos_ids_and_attn_mask upstream.

        # 4. Define step_fn: x, t -> action_in_proj -> expert(past_key_values=prompt_cache)
        #                            -> action_out_proj -> vector field
        #    CFG variant: maintain two KV caches (nav / unguided) and follow _guided_v.

        # 5. self.flow_matching.sample(batch_size, step_fn, ...) -> 10-step Euler.

        # 6. action_space.action_to_traj -> (pred_xyz, pred_rot)
        # 7. Optionally run generate_text to extract CoC / answer.
        return DiffusionOutput(
            output={"pred_xyz": pred_xyz, "pred_rot": pred_rot, "reasoning": extra},
            modality="trajectory",
            stage_durations=getattr(self, "stage_durations", None),
        )

The diffusion pipeline owns only expert + flow matching + action decode. The AR stage lives under vllm_omni/model_executor/models/alpamayo1_5/, which means we can reuse vLLM-Omni's async stage disaggregation (the exact same path BAGEL uses) without running Qwen3-VL a second time inside the diffusion worker.

5. AR runner (Alpamayo15ForConditionalGeneration)

• Inherits Qwen3VLForConditionalGeneration from upstream vLLM and overrides:
  • get_multimodal_embeddings / _parse_and_validate_multimodal_inputs to support multi-camera packed sequences (num_cameras ∈ [1, N]);
  • get_input_embeddings, replacing <|traj_history|> placeholders with embeddings produced by hist_traj_tokenizer;
  • Adds OmniAlpamayoProcessor (analogous to OmniBagelProcessor) which merges ego_history_xyz/rot, route_xyz, multi-camera images, and the text prompt into a single BatchFeature;
  • get_supported_mm_limits -> {"image": 6, "trajectory_history": 1, "route": 1}.
• Custom PromptReplacements required:
  • <|image_pad|> × num_image_tokens (reuses Qwen3-VL's logic);
  • <|traj_history_start|> <|traj_history|>*T <|traj_history_end|>;
  • <|route_start|> <|route_pad|>*K <|route_end|>.
• Vocabulary expansion: vocab_size = base_vocab + traj_vocab_size (=768) + len(TRAJ_TOKEN), matching the upstream _build_processor. load_weights must shard the resized embed_tokens / lm_head correctly under TP.

6. New output modality: trajectory

vLLM-Omni's DiffusionOutput today mainly carries PIL.Image | np.ndarray (video) | np.ndarray (audio). We add a small extension:

1. In vllm_omni/diffusion/data.py, extend DiffusionOutput with an optional field modality: str = "image" and reuse output (or add a payload: dict[str, torch.Tensor] companion field).
2. In vllm_omni/outputs.py, add OmniTrajectoryOutput (fields: pred_xyz: list[list[list[float]]], pred_rot, num_traj_samples, reasoning_text).
3. In vllm_omni/entrypoints/openai/serving_chat.py and friends, add a serializer for "output_modality": "trajectory" that returns JSON instead of a base64 image.
4. Add examples/offline_inference/alpamayo1_5/end2end.py with flags like --model nvidia/Alpamayo-1.5-10B --modality trajectory --num-traj-samples 16 --use-cfg-nav. It saves .npz / .json and optionally calls viz_utils.visualize_trajectories_on_cam to produce a PNG.

Note: we do not introduce a new top-level task (Text-to-Trajectory) to the v1 CLI in this RFC. Initial exposure goes through /v1/chat/completions with extra_body={"output_modality": "trajectory"}. When more similar models land (RT-2, π0, ...) we can promote it to a first-class task.

7. Attention / KV-cache reuse

• The AR stage continues to use upstream vLLM's PagedAttention + FlashAttention backend unchanged.
• The Expert stage must consume the KV-cache produced by the AR worker and append n_diffusion_tokens = 64 new queries on top of the fixed prefix. This matches BAGEL's CFG mechanism of injecting a NaiveCache through the request:
  • BAGEL today threads KV through req.sampling_params.past_key_values / kv_metadata (see pipeline_bagel.py, around L354–L396). Alpamayo can reuse the same field names verbatim — no new API surface is needed.
  • The Expert's cross-attention needs a 4D attention_mask that masks the [offset, -n_diffusion_tokens) gap; this is the vLLM-Omni port of _build_expert_pos_ids_and_attn_mask.
• Because Qwen3VLForConditionalGeneration uses 3-component RoPE, the Expert must receive position_ids with shape [3, b*, n_diffusion_tokens].

8. Acceleration / parallelism plan

Feature	Phase 1 (initial PR)	Phase 2	Notes
Tensor Parallelism	✅	—	Reuses Qwen3-VL TP. Expert attention / MLP layers swap to QKVParallelLinear / RowParallelLinear (follow BAGEL).
CFG Parallelism	✅	—	Two KV caches (CFG-nav / unguided) are a natural fit for CFGParallelMixin; implement diffuse() per docs/design/feature/cfg_parallel.md.
Sequence Parallelism	⚠️	✅	The trajectory sequence length is only 64, so SP has marginal benefit. Leave _sp_plan = None initially.
torch.compile	✅	—	_repeated_blocks = ["Qwen3VLDecoderLayer", "Alpamayo15ActionExpertLayer"].
CPU offload	✅	—	Expert + VLM + action heads are all nn.Modules; model-level offload works out of the box.
Cache-DiT / TeaCache	❌	⚠️	Only 10 denoising steps — caching returns are small. Skip initially.
Quantization	—	✅	bf16 first; later follow BAGEL's ComponentQuantizationConfig(prefix="alpamayo.vlm") for W4A16.
Disaggregated serving	✅	—	Two-stage (AR worker + Diffusion worker) via OmniConnector, same as BAGEL.

9. Tests and examples

• examples/offline_inference/alpamayo1_5/end2end.py: reads one sample from the HF dataset nvidia/PhysicalAI-Autonomous-Vehicles, generates pred_xyz / pred_rot, and uses viz_utils to overlay the trajectory on the front camera.

• examples/online_serving/alpamayo1_5/:

  • run_server.sh — vllm serve nvidia/Alpamayo-1.5-10B --omni --port 8091.
  • openai_chat_client.py — sends base64 images + ego_history + output_modality=trajectory.

• tests/e2e/test_alpamayo1_5.py:

  • L4 functionality test: with a fixed seed on a single sample, assert pred_xyz.shape == (1, 16, 64, 3) and that key waypoints match a HuggingFace reference within a tolerance (same style as tests/e2e/test_bagel.py).

• Append to docs/models/supported_models.md:

  | `Alpamayo15ForConditionalGeneration` | Alpamayo-1.5 (VLA, trajectory) | `nvidia/Alpamayo-1.5-10B` | ✅︎ | ⚠️ | | |
  

10. Compatibility and risks

Risk	Mitigation
Upstream config.json may declare architectures as Alpamayo1_5 / AlpamayoReasoningVLA, which will not match our Alpamayo15ForConditionalGeneration.	Add an alias in vllm_omni/patch.py; in parallel, send an upstream PR asking NVIDIA to set architectures explicitly.
New dependencies: hydra, einops, physical_ai_av, transformers >= 4.57 (Qwen3-VL).	einops is already present; we avoid hydra by constructing submodules manually; physical_ai_av is demo-only and is not a hard requirement.
The HF model is released under a non-commercial license.	Call this out explicitly in docs/models/supported_models.md. It does not affect the Apache-2.0 code we land.
ActionSpace / delta_tokenizer are NVIDIA-specific implementations.	Port only the numerical logic into vllm_omni/diffusion/models/alpamayo1_5/, preserving SPDX headers and the original license notice.
Qwen3-VL TP is stable upstream, but the Expert's shared-KV shape may conflict with vLLM's PagedAttention metadata.	In Phase 1 the Expert uses vllm_omni.diffusion.attention (a FlashAttention-based bypass, same as BAGEL) — it does not go through PagedAttention.

11. Milestones

• Step 1 (Week 1–2): Skeleton PR — registry entries + directory layout + Alpamayo15Pipeline.__init__ can .from_pretrained and load weights (no forward yet).
• Step 2 (Week 2–3): Offline end-to-end (single-process, synchronous) for both generate_text and sample_trajectories_from_data_with_vlm_rollout; numerical parity with the upstream reference (L4 test).
• Step 3 (Week 3–4): Wire up AR + Diffusion stage disaggregation (reusing BAGEL's OmniConnector path); enable TP and CFG parallel; land examples and docs.
• Step 4 (Week 4–5): CI L4 test + docs/models/supported_models.md + examples/online_serving/alpamayo1_5/.
• Step 5 (Week 5+): Quantization / cache_dit / SP and other advanced features (follow-up PRs).

Feedback Period.

No response

CC List.

@Gaohan123 @yinpeiqi @hsliuustc0106 @ywang96 @fake0fan @princepride @TKONIY

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.