[Refactor][Diffusion] Move single-model layers into their model directories

[congw729]

Purpose

Complete a bounded part of JiusiServe/vllm-omni-project-manage#27 by moving model-specific layers to their owning model directories.

• Move WanS2VRotaryPosEmbed and RotaryEmbeddingWanS2V to models/wan2_2/rope.py, and RMSNormVAE to models/wan2_2/norm.py.
• Keep RotaryEmbeddingWan in the shared RoPE module: both Wan and Lingbot-World now use it. RotaryEmbeddingWanS2V inherits this shared implementation. Shared RoPE helpers and RotaryEmbeddingS2VGrid also retain their existing locations.
• Move the Bagel MoT package, including linear layers, RMSNorm and Triton ops, to models/bagel/mot/. Update model/test imports, benchmark and documentation paths, and the existing lint allowlist entry.

Synchronized with main at 6a03e45c340a60553884432ff4fc4a31a6219d55 on 2026-09-20 (PR head 73ca73037d89fe7c6712b1dad734a0b9ceed01d2). The migration preserves this baseline's operator implementations, including NPU index_select, RMSNorm residual handling, vLLM 0.29.0 adaptation and CustomOp platform dispatch. The previous shared S2V indexing change is removed from the PR.

The moved internal import paths change without forwarding shims. Callers importing these internal modules must use their new locations. MoT's package-local tuning directory moves with the package; VLLM_TUNED_CONFIG_FOLDER remains available. CLI options, model registration names and operator signatures are unchanged.

Current value and scope

Reassessed on 2026-09-11. The original scope has been narrowed because RotaryEmbeddingWan is now shared. The remaining migration has modest maintenance value: it makes ownership of Wan S2V/VAE and Bagel MoT code clearer while retaining genuinely shared components in layers/.

Existing model functionality works on main; this is a normal-priority maintenance cleanup with no performance improvement claim. The scope stays limited to relocation and reference updates.

Validation

Latest refresh (2026-09-20)

Merged main 6a03e45c3 without conflicts. The migration scope is unchanged. This brings in #7426, which updates the realtime VAD test to expect response.output_audio.* events.

The previous CUDA CI run #15669 had 32 passed steps, one failed step and two skipped steps. Diffusion groups and both L2/L3 Wan22 E2E steps passed. Its only failing step was the H100 entrypoint suite: the Qwen3-Omni server-VAD test still expected the legacy response.audio.done name. The test and event producer were unchanged by this PR; the main synchronization incorporates the upstream correction.

For the refreshed head, AST equivalence, stale-path checks and git diff --check passed. Applicable PR-file hooks passed except the same eight mypy findings, all reproduced on untouched main 6a03e45c3; only that hook was skipped for the signed-off merge commit. CUDA L2/L3 CI has been triggered with ready, merge-test and cuda-test: fresh run. Results are pending.

Earlier targeted tests

The following SP_H200 results were recorded on 2026-09-11 for head 22d0fe72a, synchronized with main 43b8de9b0. They used a dedicated worktree with Python 3.12 and vLLM 0.29.0, matching that CI baseline. These historical results are separate from the CI rerun for the refreshed head. Verified the runtime imported this worktree. GPU tests used one reserved H200-class device; separate pytest processes isolated distributed state.

311 tests passed across the following groups; all final test invocations exited with status 0.

Test group	Result
Shared layers	100 passed, 9 deselected
Wan helpers / quantization configuration / transformer construction	31 passed
Lingbot attention	15 passed
Bagel unit suite	100 passed
MoT linear kernels, including und-mode	49 passed
MoT RMSNorm kernels	14 passed
MoT NPU fallback with mocked platform	2 passed

Bagel, Lance, Lingbot, Wan S2V and VAE norm imports passed. The MoT package-local configuration directory resolves to the new Bagel location.

Reproduce from the checkout with the matching development environment. Use separate pytest invocations; a visible CUDA device is needed for the layers and Bagel groups even though their tests carry CPU markers.

export PYTHONPATH="$PWD"
python -m pytest -q tests/diffusion/layers/ -m 'core_model and cpu' --run-level=core_model
python -m pytest -q tests/diffusion/models/bagel/ -m 'core_model and cpu' --run-level=core_model
python -m pytest -q \
  tests/diffusion/models/wan2_2/test_wan22_quant_config_propagation.py \
  tests/diffusion/models/wan2_2/test_wan22_pipeline_helpers.py \
  tests/diffusion/models/wan2_2/test_create_transformer_quant.py \
  -m 'core_model and cpu' --run-level=core_model
python -m pytest -q tests/diffusion/models/lingbot_world/test_lingbot_world_attention.py -m 'core_model and cpu' --run-level=core_model
python -m pytest -q tests/diffusion/kernels/mot/test_mot_norm_npu.py -m 'core_model and cpu' --run-level=core_model
python -m pytest -q tests/diffusion/kernels/mot/test_mot_linear.py -m 'core_model and gpu' --run-level=core_model
python -m pytest -q tests/diffusion/kernels/mot/test_mot_norm.py -m 'core_model and gpu' --run-level=core_model

Additional validation:

• AST comparison confirms relocated and retained RoPE/Norm definitions match the synchronized main baseline. All seven MoT files are AST-identical after import normalization.
• No stale migrated import/config paths remain in the repository, including lint configuration.
• Ruff, formatting, SPDX, forbidden-import and other applicable changed-file hooks passed; git diff --check passed.
• Mypy 3.10 reports eight existing errors in shared RoPE/Norm and two test files. The unmodified main baseline reproduces the same eight errors. This hook is skipped for the merge commit after that comparison; no new ignores or unrelated typing changes are added.
• An initial layers run with GPUs hidden failed on CUDA tensors and NCCL setup. The CUDA-visible rerun is reported above. The old und-mode failures are not reproduced on this baseline.

The SP_H200 targeted suite did not run full-model inference; subsequent Wan22 E2E coverage is recorded above. Physical NPU/XPU hardware was not tested. CPU platform-fallback tests and preservation of main's implementations provide the available cross-platform evidence.

[chatgpt-codex-connector]

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.

[hsliuustc0106]

lgtm

[congw729]

@Gaohan123 @hsliuustc0106 can be merged

[congw729]

@hsliuustc0106 The conflicts are resolved in 22d0fe7, and I've narrowed the migration: RotaryEmbeddingWan stays shared because Lingbot-World also uses it; only Wan S2V/VAE and Bagel MoT layers move. All 311 targeted tests passed with vLLM 0.29.0, and the PR description now includes the current scope, maintenance value, and validation details, including the pre-existing mypy findings. Could you take another look when you have a chance? 🙂

[congw729]

@hsliuustc0106 Now that #5907 has merged, could you give this PR's narrowed scope a final look and merge it once the remaining merge requirements are satisfied? 🙂

RotaryEmbeddingWan stays shared for Wan and Lingbot; only Wan S2V/VAE and Bagel MoT layers move. The current head (22d0fe72a) passed 311 targeted tests on September 11, and all currently reported CI checks are green. I also checked integration with the latest main (01a2f9325): the simulated merge is conflict-free and the relocated implementations remain equivalent. This latest integration check was static; the runtime suite has not been rerun against that combined tree. The PR description contains the scope, validation details and limitations.

[congw729]

@hsliuustc0106 Updated to main in 73ca73037, and the CUDA L2/L3 rerun is now green: Buildkite #15677 — 34 steps passed, 0 failed, 2 skipped. Wheel builds, pre-commit and DCO also passed. The narrowed migration scope is unchanged.

This looks ready to merge from the code and test side. Read the Docs is still showing pending; could you merge once that remaining check clears? 🙂