[Dong1017]

Purpose

Upstream main cannot run the LingBot-World realtime benchmark on Ascend NPU: the first failure is RotaryEmbeddingWan.forward_native receiving SP-sharded 2D cos/sin (dim-2 mismatch, 40 heads vs 1170 local tokens), and the AR-Diffusion paged-attention NPU path has no working attention call. This PR carries the NPU enablement for the realtime streaming stack (#7645) so main runs end-to-end on 4x Ascend 910B1 with Ulysses SP-4 in eager mode.

The change surface is kept minimal — four files: the rope broadcast fix, the paged-attention NPU path, the transformer routing, and the validated 4-card deploy config. All NPU branches are device-guarded (query.is_npu, device.type == "npu", lazy import torch_npu inside NPU-only branches); GPU/CUDA paths are untouched.

Changes (6 commits):

• NPU adaptation for LingBot-World realtime benchmark: ar_diffusion_paged_attention NPU branch gathers the visible window and calls npu_fusion_attention; _to_device_async covers NPU pinned async copies; the transformer direct-replay path routes through the paged attention entry on NPU.
• paged_attention NPU fixes: the contiguous-gather path masked rows of the unflattened gather along the wrong axis; npu_fusion_attention with input_layout="TND" now receives per-sequence actual_seq_qlen / actual_seq_kvlen derived from query_start_loc / seq_lens.
• benchmarks: npu_usp4_eager.yaml, the 4-card Ulysses SP-4 eager deploy config. Four cards are required because the single-card geometry (weights plus one session's KV window) does not fit in 64 GB; torch.compile / CUDA-graph paths are GPU-only, so enforce_eager is set.
• rope: broadcast 2D SP-sharded cos/sin over batch and heads in the Wan native path (fixes the startup failure above).
• rope tests: pin the 2D kernel-facing cos/sin layout in the Wan native path as a CPU regression test (the 2D result must equal the manually expanded form bit for bit).
• paged_attention staging: on NPU, return zero-size buffers instead of (None, None) when staging is disabled — NPU eager dispatch of the fused write+attend custom op rejects None for a tensor argument declared mutable, and the op treats an empty stage buffer like None (fresh allocation). Other devices keep the upstream None contract. Upstream's op-internal window narrowing ([Bugfix][Core] Preserve static KV staging inputs in AR-Diffusion #7921) is untouched, and the e2e check below confirms numerical equivalence on NPU.

Test Plan

• e2e realtime benchmark, benchmarks/lingbot_world/benchmark_lingbot_world_realtime.py, LingBot-World 2.0 14B, 4x Ascend 910B1 (64 GB), Ulysses SP-4, eager, 32 chunks/session, two independent server+bench runs:

vllm-omni serve <model-path> --omni \
    --deploy-config benchmarks/lingbot_world/configs/npu_usp4_eager.yaml --port 8200

python3 benchmarks/lingbot_world/benchmark_lingbot_world_realtime.py \
    --host 127.0.0.1 --port 8200 --model <model-path> \
    --num-chunks 32 --warmup-sessions 1 --sessions 1 \
    --save-video out.mp4 --output-json out.json

• CPU regression suite: pytest tests/diffusion/ar_diffusion/test_paged_attention.py tests/diffusion/ar_diffusion/test_kv_cache.py tests/diffusion/layers/test_rope_broadcast.py tests/benchmarks/test_lingbot_world_realtime.py → 166 passed, 14 skipped (includes the 2D-layout regression test above)

vLLM Version: vllm 0.29.0 / vllm-omni main + this branch, torch 2.10.0, torch_npu 2.10.0.post4, vllm-ascend 0.19.1rc2.dev2045, triton_ascend 3.2.2

vLLM-Omni Commit: 91498a3 (test-only commit on top of 9678c60, where the e2e runs below were validated)

Test Result

• main + this branch starts the server and serves the realtime benchmark end-to-end on 4x 910B1 (previously impossible on NPU).
• Deterministic output: the two independent runs produce byte-identical videos (same md5), i.e. NPU eager execution is reproducible; the output is also bit-identical to our pre-rebase development stack, confirming [Bugfix][Core] Preserve static KV staging inputs in AR-Diffusion #7921's staging rework did not change numerics on NPU.
• Performance: steady chunk-interval median 3016.0 / 3016.4 ms (steady RTF 5.84 / 5.89 at 16 fps), consistent with the pre-rebase baseline band (3034-3093 ms).

Eager mode only; compiled execution on NPU (torchair / ACL graph) and the framework-side memory-management follow-ups are tracked in #7382.

Performance: whole-block index_select KV gather on NPU (2026-10-14, commit 7c7d9d3)

ar_diffusion_paged_attention's NPU branch now gathers the visible KV window with one whole-block index_select + flatten + truncate instead of per-token advanced indexing. The row mapping is identical (same block_table), so outputs are byte-identical; 477-frame precodec capture and independent RGB/PTS decode match the per-token implementation exactly.

Measured on 4x Ascend 910C1 (SP4/Ulysses4, TP1, BF16 eager, 832×480, seed42, 4 DMD steps, 40 chunks/477 frames, fixed teaser input; OmniRSI 3× alternating AB/BA, trial spread < 0.4%):

• full-request latency median 108291 → 103890 ms (−4.06%)
• lossless: candidate vs per-token precodec bit patterns byte-equal for all 40 measured chunks; decoded RGB (480×832×3) and PTS (16 fps) equal

SP8 screening (8 dies on one 910C1 node, same workload, steady chunk inter-arrival, n=33):

• median 1713.9 → 1587.5 ms (−7.4%), P99 2002.6 → 1907.0 ms (−4.8%)

Environment note: the contiguous-KV behavioral patch in the NPU platform (requires_contiguous_pa_kv_cache) exists in vllm-ascend dev2045 (the CI image above) but not in newer dev builds (dev2552); environments on newer vllm-ascend need that helper restored or the patch guarded.

[vllm-omni-review-bot]

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

Module owners: @wtomin @david6666666 @Isotr0py

Routing: @wtomin via module of the changed files, module named in the PR description, CODEOWNERS; @david6666666 via module of the changed files, module named in the PR description; @Isotr0py via module of the changed files, module named in the PR description

@Dong1017, 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.

[Dong1017]

Self-review notes

Ran the project's precheck-pr skill (full mode) against this branch before taking it out of draft. Summary of what was checked and what changed as a result.

Issues found and fixed during self-review

1. Upstream staging contract violation (caught by the CPU test suite). test_staging_is_only_allocated_for_the_gather_path[False] asserts history_staging() returns (None, None) when staging is disabled. My initial zero-size-buffer fallback was unconditional and violated that. Fixed: the fallback is now NPU-only (k_pool.device.type == "npu"); all other devices keep the upstream None contract. The NPU-side behavior is byte-identical to what the e2e validation below exercised, so the md5/bit-exact evidence still applies.
2. ruff F841: unused local batch_seqs removed.
3. ruff format: two redundant-paren lines reformatted.

Checks performed

• ruff check + ruff format --check with the repo config on the three touched Python files: clean. (PLC0415 is not in the repo's select list, so the lazy import torch_npu inside the query_flat.is_npu branch is fine; it is not in the forbidden-imports list, and it keeps GPU-only machines from importing it.)
• CPU regression suite in a container at this exact commit: pytest tests/diffusion/ar_diffusion/test_paged_attention.py tests/diffusion/ar_diffusion/test_kv_cache.py tests/diffusion/layers/test_rope_broadcast.py tests/benchmarks/test_lingbot_world_realtime.py → 165 passed, 14 skipped, 0 failed. (CPU-marked tests guard the shared paths; the NPU attention branch itself is exercised by the 4-card e2e runs below, not by CPU tests.)
• Code-quality sweep on added lines (kwargs plumbing / broad excepts / Any hints / hot-path copies / event-loop blocking): no hits. The if query.is_npu → elif query.is_cuda → else restructuring in transformer.py preserves the existing dense-attention fallback for CPU.
• Scope: 4 files, +145/−3; the deploy config is the only new file (yaml; the SPDX hook covers source files, and sibling benchmark configs carry no SPDX header either).

Known limitations (disclosed)

• Eager mode only; compiled execution on NPU (torchair / ACL graph) is follow-up work.
• Steady-state VRAM peak was not recorded in the e2e runs; capacity is enforced by the KV-pool free-memory admission check at startup.
• The NPU attention path performs a few per-layer host readbacks (seq_lens.item() / .cpu()) to feed npu_fusion_attention's TND layout. We measured their e2e cost on this workload as noise-level (<20 ms over a 32-chunk session), but it is an obvious follow-up optimization target.

[vllm-omni-review-bot]

Omni ReviewBot triage note

Automated triage of commit 7c7d9d35fe77 produced:

• Priority: high. Prompt maintainer attention is suggested.

These are automated triage suggestions only — the final decision belongs to the maintainers.

[vllm-omni-review-bot]

Omni ReviewBot routing record

Assigned Direct under experiment vllm-omni-strict-5050-20260829.

[vllm-omni-review-bot]

Omni ReviewBot review

PR description

This PR enables LingBot-World realtime on Ascend by fixing Wan native RoPE to broadcast SP-sharded 2D cos/sin over batch and heads, adding an NPU gather + npu_fusion_attention path in AR-Diffusion paged attention (with NPU-only empty staging buffers for the mutable custom-op contract), and routing the transformer direct-replay oracle through that same entry on NPU. It also adds a 4×910B1 Ulysses SP-4 eager deploy YAML because the single-card weight+KV footprint exceeds 64 GB. User-visible result: the previously broken Ascend realtime stack can serve the benchmark end-to-end in eager mode; CUDA/ROCm attention math paths are intended to stay unchanged.

Change flow

flowchart TD
  A["[EXISTING] Realtime serve / SP-4 session"]:::existing --> B["[NEW] npu_usp4_eager.yaml"]:::new
  A --> C["[CHANGED] RotaryEmbeddingWan 2D cos/sin broadcast"]:::changed
  A --> D["[CHANGED] ar_diffusion_paged_attention NPU gather+fusion"]:::changed
  D --> E["[CHANGED] LingBot self-attn NPU direct-replay routing"]:::changed
  B --> F["[EXISTING] 4x 910B1 eager realtime chunks"]:::existing
  C --> F
  E --> F
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

CI at da549760647d (2026-09-30T20:14:11.573540+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

See inline comments below.

---

_{🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!}

[vllm-omni-review-bot]

[P1] Shared paged-attention entry uses Tensor.is_npu before CUDA/CPU dispatch

ar_diffusion_paged_attention now gates on query_flat.is_npu before any CUDA/ROCm/reference branch. Stock PyTorch does not define Tensor.is_npu unless torch_npu (or an equivalent PrivateUse1 rename) has registered it, so CUDA/CPU callers can raise AttributeError before reaching FlashAttention or _reference_paged_attention. The same file already uses the portable guard k_pool.device.type == "npu" for staging; the hot path should match (query_flat.device.type == "npu") so GPU-only machines stay on the existing paths. Author CPU/e2e evidence used a torch_npu-bearing stack, which masks this on non-Ascend CI images.

[vllm-omni-review-bot]

[P2] Direct-replay NPU guard can AttributeError on CPU tensors

The direct-replay condition is now (query.is_cuda or query.is_npu). CUDA short-circuits safely, but on CPU is_cuda is false so Python evaluates query.is_npu; without torch_npu that raises instead of falling through to the prior dense self.attn oracle path. Prefer query.device.type == "npu" (same portable check as staging) so CPU direct-replay keeps working on non-NPU installs.

[vllm-omni-review-bot]

Omni ReviewBot: finding feedback

[p1] Shared paged-attention entry uses Tensor.is_npu before CUDA/CPU dispatch — vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py:522

ar_diffusion_paged_attention now gates on query_flat.is_npu before any CUDA/ROCm/reference branch. Stock PyTorch does not define Tensor.is_npu unless torch_npu (or an equivalent PrivateUse1 rename) has registered it, so CUDA/CPU callers can raise AttributeError before reaching FlashAttention or _reference_paged_attention. The same file already uses the portable guard k_pool.device.type == "npu" for staging; the hot path should match (query_flat.device.type == "npu") so GPU-only machines stay on the existing paths. Author CPU/e2e evidence used a torch_npu-bearing stack, which masks this on non-Ascend CI images.

This finding appears in the bot's COMMENT review, but GitHub could not place it as an inline diff comment. If you are the PR author and disagree, react 👎 to this comment. The disagreement will be shown to the maintainer; it does not approve or merge the PR.

[vllm-omni-review-bot]

Omni ReviewBot: no human activity for 13 days

@Dong1017 this pull request has had no human commit, comment or review since 2026-09-24. Please confirm the current plan and next step. The author or a maintainer decides whether to change the PR state.

To keep it moving, any one of these is enough: push an update, reply to the open blocker, or post the current plan and timeline.

[Dong1017]

@tzhouam, could you review the Ascend enablement in this PR, especially the NPU paged-attention/staging path and the Wan RoPE broadcast fix? Thank you.

[vllm-omni-review-bot]

Omni ReviewBot review

Changes since the previous review

• 0 new inline finding(s); 0 finding(s) below.
• Fixed: Shared paged-attention entry uses Tensor.is_npu before CUDA/CPU dispatch — vllm_omni/experimental/ar_diffusion/kv_cache/paged_attention.py
• Fixed: Direct-replay NPU guard can AttributeError on CPU tensors — vllm_omni/diffusion/models/lingbot_world/transformer.py

CI at c31d46153438 (2026-10-08T03:58:42.806795+00:00): required check(s) blocking: pre-commit (pending), buildkite/vllm-omni (missing), docs/readthedocs.org:vllm-omni (pending).

Full review analysis

PR description

This PR makes the LingBot-World realtime stack runnable on four Ascend 910B1 cards in eager Ulysses sequence-parallel 4. Shared Wan native RoPE now expands 2D (local_tokens, rope_dim) cos/sin to (1, S, 1, D/2) so sequence-parallel tokens broadcast over batch and heads instead of lining up with the head axis. On NPU, paged self-attention gathers the valid prefix of the block table and calls npu_fusion_attention with TND lengths; direct replay uses that same entry, and disabled history staging passes empty tensors so the fused write+attend op can dispatch. A new deploy file pins four devices, eager execution, and the 480×832 KV fraction used for the reported end-to-end benchmark.

Change flow

flowchart LR
  cfg["[NEW] npu_usp4_eager.yaml"]:::new
  serve["[EXISTING] LingBot realtime forward"]:::existing
  rope["[CHANGED] Wan forward_native 2D RoPE"]:::changed
  route["[CHANGED] Direct-replay NPU guard"]:::changed
  attn["[CHANGED] NPU paged fusion attention"]:::changed
  video["[EXISTING] Eager Ulysses SP-4 chunks"]:::existing
  cfg --> serve
  serve --> rope
  serve --> route
  route --> attn
  rope --> video
  attn --> video
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

---

_{🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!}

[vllm-omni-review-bot]

Omni ReviewBot review

0 actionable finding(s).

CI at 7d7d9941a2af (2026-10-08T20:19:22.518876+00:00): required check(s) blocking: buildkite/vllm-omni (missing).

Full review analysis

PR description

This change makes the LingBot-World realtime benchmark runnable in eager mode on four Ascend 910B1 cards with Ulysses sequence parallel 4. Wan's native rotary path expands sequence-by-rope 2D cos/sin so those tables broadcast across batch and heads. On NPU, paged self-attention gathers the visible KV window and calls npu_fusion_attention; the direct-replay path uses that same entry; and disabled history staging passes empty tensors so the fused write-and-attend op can accept its mutable arguments. The new deploy file pins the four-card eager topology, 480×832 resolution, and KV memory fraction used for that run.

Change flow

flowchart LR
  serve["[EXISTING] LingBot realtime serve"]:::existing
  yaml["[NEW] npu_usp4_eager.yaml"]:::new
  rope["[CHANGED] Wan native RoPE broadcast"]:::changed
  attn["[CHANGED] NPU paged attention gather"]:::changed
  route["[CHANGED] NPU direct-replay routing"]:::changed
  bench["[EXISTING] Eager SP-4 chunk loop"]:::existing
  serve --> yaml
  yaml --> rope
  rope --> attn
  attn --> route
  route --> bench
  classDef existing fill:#e5e7eb,stroke:#6b7280,color:#111827
  classDef changed fill:#fef3c7,stroke:#d97706,color:#451a03,stroke-width:2px
  classDef new fill:#dcfce7,stroke:#16a34a,color:#052e16,stroke-width:2px
  classDef removed fill:#fee2e2,stroke:#dc2626,color:#450a0a,stroke-width:2px

Loading

No actionable findings.

---

_{🤖 This review was generated by InferMatrix Copilot, an open-source repo-maintenance agent for PR review, CI debugging and issue triage. Try it on your own repo, and ⭐ star it if it helped!}

[vllm-omni-review-bot]

Omni ReviewBot: three questions on the performance claim

@Dong1017 this PR reads as a performance or value claim:

• claim: - full-request latency median 108291 → 103890 ms (−4.06%)

Before the full evidence checklist, three short questions:

1. Bottleneck — what is the current bottleneck, and which profile, trace or per-stage measurement shows it?
2. Value — what does the change buy the user or the system (latency, throughput, memory, cost), and at which workload?
3. A/B or ablation — is there a same-workload, same-head/config comparison that isolates each main claim on its own? For stacked optimizations, one number per item rather than a blended delta.

When you answer, the evidence that settles it is: base and head SHA, hardware, model, workload, warm-up and repeat count, mean or percentiles with their spread, and a correctness/quality-equivalence signal; an end-to-end claim also needs stage attribution.

[tzhouam]

LGTM