Skip to content

[Misc] Add AR and diffusion MoE backend acceptance tests - #9061

Open
shunwendongan wants to merge 2 commits into
vllm-project:mainfrom
shunwendongan:codex/moe-backend-runtime-acceptance
Open

shunwendongan wants to merge 2 commits into
vllm-project:mainfrom
shunwendongan:codex/moe-backend-runtime-acceptance

Conversation

@shunwendongan

Copy link
Copy Markdown
Contributor

Add AR and diffusion MoE backend acceptance tests

Contributor intent

  1. What outcome do i want to achieve?

    Complete the AR MoE and Diffusion MoE acceptance tests requested by the
    maintainer after #8092
    and submit them as an independent follow-up contribution.

  2. What is the highest-leverage variable for that outcome? I think

    Executable coverage of deploy-selected backend routing: real worker startup,
    scheduler-visible prefill/decode, and actual linear, attention, and MoE
    implementations rather than configuration-only assertions.

  3. How will this work move that variable?

    Add a local tiny random-weight Qwen3 MoE deploy test that starts a real Omni
    AR stage and serves a token-ID request, plus a native SenseNova-U1 MoE block
    test that follows deploy configuration into the real Triton expert forward
    and compares it with an independent expert-loop reference.

  4. Whose customer or user needs does this address, and what are those needs?

    It addresses vLLM-Omni maintainers and users who need deploy backend choices
    to reach the intended runtime implementation and to catch routing or
    prefill/decode regressions before model-specific deployment work.

Purpose

This is the MoE follow-up to the dense FLUX runtime coverage in
#9051. It adds coverage
for the two MoE follow-up targets from the #8092 maintainer feedback. The
deploy/YAML propagation fixes and dense FLUX tests remain in #9051; this
independent branch changes tests only.

This change adds executable runtime coverage for:

  • an AR MoE stage/worker starting from deploy YAML and serving both prefill and
    decode;
  • the selected AR linear, attention, and MoE methods being called by the real
    worker; and
  • a diffusion MoE block performing real expert math after deploy backend
    selection.

After this change:

  • tests/engine/test_backend_runtime.py creates a local tiny
    Qwen3MoeForCausalLM config, uses load_format: dummy, starts a real
    AsyncOmniRunner worker, sends 17 prompt token IDs, and verifies four output
    tokens plus scheduler-computed prefill/decode events and implementation
    probes;
  • tests/engine/backend_runtime_worker.py supplies a test-only worker class
    that adapts only the stock model's Omni metadata signature and observes the
    real backend calls;
  • tests/helpers/runtime.py adds a tokenizer-free token-ID request helper for
    this acceptance test; and
  • tests/diffusion/models/sensenova_u1/test_sensenova_u1_moe_cuda.py drives
    the existing native SenseNova-U1 MoE block from a deploy-selected
    VllmConfig, observes the real router GEMM and MoE-kernel apply calls,
    verifies that the selected expert implementation is TritonExperts, and
    checks an independent FP32 expert-loop reference.

No production code, CI YAML, model checkpoint, tokenizer, or new runtime
dependency is added. The tests use random tiny weights and establish runtime
and numerical wiring only; they make no trained-model quality, image-quality,
throughput, latency, or performance claim. The diffusion test is a native MoE
block forward test, not a complete SenseNova pipeline startup/request test.

Design documents and user guides are unchanged because the tests exercise
existing configuration projection, AR scheduling, and native diffusion model
contracts without changing production behavior, public options, or defaults.
The one-token prefill-tail assertion follows the existing computed-token
phase rule; it does not introduce a new scheduler policy.
Test Plan

vLLM Version: 0.31.0

vLLM-Omni Commit: 8174297afc59ea2060547ea3b705cd0a0630924c (base:
abff966489bd8011009d84d828417ce0c55af872)

Hardware: WSL2 Ubuntu 24.04.4, NVIDIA GeForce RTX 3080 10 GiB, compute
capability 8.6, driver 616.56, CUDA runtime 13.0, Python 3.12.3.

Packages: PyTorch 2.13.0+cu130, Triton 3.7.1, Transformers 5.14.1,
pytest 9.1.1.

Environment policy: existing /home/shunwendongan/codex-work/vllm-omni-backend-runtime/venv
and HF cache only; HF_HUB_OFFLINE=1 and TRANSFORMERS_OFFLINE=1 were set.
No new model or package was downloaded.

Test Result

GPU acceptance tests on the current branch:

1 passed, 16 warnings in 154.11s
1 passed, 16 warnings in 6.24s

Commands:

cd /path/to/vllm-omni
HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 \
VLLM_WORKER_MULTIPROC_METHOD=spawn \
../venv/bin/python -m pytest -s -v -n 0 \
  tests/engine/test_backend_runtime.py \
  -m "core_model and cuda and L4" --run-level core_model

HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 \
../venv/bin/python -m pytest -s -v \
  tests/diffusion/models/sensenova_u1/test_sensenova_u1_moe_cuda.py \
  -m "core_model and cuda" --run-level=core_model

The passing AR test asserts default_unquantized_gemm, TritonAttentionImpl.forward,
TritonExperts.apply, a spawned worker PID, finite hidden states, four output
tokens, and scheduler events (0,16,prefill), (16,1,prefill),
(17,1,decode), (18,1,decode), (19,1,decode).

The diffusion log records one router GEMM call and one MoE-kernel apply
call with TritonExperts selected, with
relative error 0.004687030799686909 against the independent reference and
0.7783500552177429 after swapping gate/up weights.

CPU and static checks:

187 passed, 1 skipped, 23 warnings
py_compile: passed
Ruff check: passed
Ruff format --check: passed
typos / debug-statements / end-of-file / line-ending / whitespace: passed
mypy-3.10 repository wrapper (isolated hook dependencies): 4 files passed
repository SPDX / forbidden-imports / torch.cuda / test-mark checks: passed
AR collect-only: 1 test collected
git diff --check: passed

CPU command:

PYTHONPATH="$PWD" ../venv/bin/python -m pytest -q \
  tests/helpers/tests/test_runtime_async_omni.py \
  tests/config/test_backend_selection.py \
  tests/engine/test_stage_engine_args.py \
  -m "core_model and cpu" --run-level=core_model

The CPU summary above was observed in the terminal; a separate log for that
187-test run was not retained. Both GPU final-run logs and exit codes were
retained. The AR log prints the temporary probe JSON path, not its contents;
the implementation/event claims above are checked by assertions in the passing
test. That temporary JSON is no longer available.

The repository's existing Buildkite L4 and Diffusion Model Test jobs already
collect these paths, so no CI YAML change is needed. Draft PR admission still
uses the project's normal ready label or maintainer trigger; this local
evidence is not an assertion that GitHub/Buildkite has already run.

The monolithic pre-commit aggregate was not completed: its Linux actionlint
environment was not cached. Applicable checks were instead run using existing
cached entrypoints, without installing packages. The repository's actual
tools/pre_commit/mypy.py 3.10 wrapper passed all four changed Python files in
the existing isolated lint environment with the hook's declared dependencies,
including types-PyYAML; no error-code suppression was needed. YAML,
shellcheck, markdownlint, TTS-ratchet, and Buildkite-schema hooks do not match
the four changed Python paths.

References

Signed-off-by: wu kui <wukui1030660585@163.com>
@chatgpt-codex-connector

Copy link
Copy Markdown

Codex Review Summary

This comment shows the latest Codex review activity on this pull request.

Review Status Commit Review trigger
📝 Code Review ⚠️ Failed 2026-10-18T14:13:11.095128Z 8174297 PR opened
ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

  • Open a pull request for review
  • Mark a draft as ready
  • Comment "@codex review" or "@codex security review".

Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings.

@chatgpt-codex-connector

Copy link
Copy Markdown

Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits.
Credits must be used to enable repository wide code reviews.

@vllm-omni-review-bot

Copy link
Copy Markdown

This PR appears to belong to: docs/design/module/diffusion/index.md.

Module owners: @david6666666 @Isotr0py @princepride

Routing: @david6666666 via module named in the PR description; @Isotr0py via module named in the PR description; @princepride via module named in the PR description

@shunwendongan, please review your own changes and leave a short self-review comment describing what you checked. PRs without author self-review may not be assigned a reviewer.

Please take a look when you have a chance. If you would like an automated review, mention @vllm-omni-review-bot in a comment.

@shunwendongan

Copy link
Copy Markdown
Contributor Author

Self-review: I checked the four changed test files and verified that the AR test exercises deploy-selected linear, attention, and MoE implementations by a real worker with scheduler-visible prefill/decode events. The diffusion test follows deploy configuration into the native Triton MoE forward and compares the output with an independent reference. The PR does not modify production code, CI configuration, model checkpoints, or public defaults. Local GPU, CPU, static checks, and the GitHub build/pre-commit checks are documented in the PR description.

@hsliuustc0106 hsliuustc0106 added the CI/CD codes related to changes to CI/CD label Oct 19, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI/CD codes related to changes to CI/CD

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants