[Pingzii]

What this fixes

The Attention-side gate reads the profile flag in plain Python:

if force_balanced_topk_ids_enabled() or _EXTRA_CTX.in_profile_run:
    topk_ids = _force_balanced_topk_ids(...)

Tracing specializes that branch on the value seen while profiling, so the compiled graph keeps applying profile-time balancing to live routing. Native Ascend avoids this because the same read sits inside the opaque moe_forward boundary.

Change

The decision moves into a mutating custom op:

def afd_balance_routing_ids_(topk_ids: torch.Tensor, num_experts: int, force_balance: bool) -> None:
    from vllm_ascend.ascend_forward_context import _EXTRA_CTX
    from afd_plugin.routing_balance import _force_balanced_topk_ids
    if force_balance or _EXTRA_CTX.in_profile_run:
        _force_balanced_topk_ids(topk_ids, num_logical_experts=num_experts)

• registered with direct_register_custom_op(mutates_args=["topk_ids"]) plus a fake implementation, so the flag is read at execution time;
• the gate GEMM and topk stay inside the compiled region, and nothing disables compilation;
• the two call sites keep their previous semantics (v4: profile only, v2: its existing explicit switch);
• the pure tensor helper moves to afd_plugin/routing_balance.py, so the compatibility layer no longer imports model code;
• registration happens in the existing worker-side patch installer.

No runner patch.

Evidence

CPU results at 110793d (earlier head): targeted tests 30 passed (recomputing the op through a fixture-scoped registration built from the production signature, Meta uses the production fake, a no-guards replay reuses the first compiled callable with the old plain-Python branch as a negative control, opcheck covers the mutating branch, Inductor case included); maintenance suite passes. A reversed check with mutates_args=() fails opcheck, so the mutation assertion is live.

NPU results below were collected on the integrated branch at e181f22, using the same
core routing-boundary implementation with the runner-level patch removed. This PR ports
that fix to main; the current head passes 30 targeted CPU tests and the maintained CPU suite.

• Standalone NPU check: profile mode writes balanced IDs; live IDs are preserved when profile mode is disabled.
• Graph capture completes 5/5 (Graph capturing finished in 8 secs).
• Runtime confirms first replay: ACL graph replay is active ... (logged once).
• Full GSM8K (1,319 samples): 90.14, with gate-on-A, quant_mode=0, and FULL_DECODE_ONLY.
• Same-day throughput: 309,221 median for the replacement (n=2), versus 311,978 for the patch-based state (n=1), a difference of -0.9% with limited samples.

Registration now also covers the standalone checkpoint-reference script after device
initialization. CPU gate tests use a shared, scoped fixture that registers the production op
and fake implementation. Against head 110793d, all 30 targeted tests pass, including the 13
existing DSV4 gate tests; the maintained CPU suite also passes.

Compiled-artifact reuse on this head (732d03a)

tests/npu/check_routing_boundary_reuse.py, real NPU, production op registration and a real
forward context. The stack is put into V2 extra-kwargs mode through its own entry point
(sync_v2_extra_kwargs), and the log confirms VLLM_USE_V2_MODEL_RUNNER=True is set; using Model Runner V2.

Three assertions establish the layout before the reuse checks run:

• the cache reports V2 mode (_USE_V2_EXTRA_KWARGS is True);
• the write lands in additional_kwargs (it raised KeyError on the legacy layout, so this is
  not a pass-by-default);
• the op's read path (_EXTRA_CTX.in_profile_run) agrees with what was written.

Then: compile once while the profile flag is set and assert that first compiled call produced
balanced IDs; clear the flag, reuse the same compiled callable for two distinct non-balanced
inputs (both keep their own IDs, backend reports one compilation); capture an ACL graph with the
flag cleared and replay it twice over one static buffer, checking each replay's IDs outside the
captured region.

eager compiled reuse: OK
acl graph reuse: OK
ALL CHECKS PASSED

Device results on this head were provided by the contributor; the reviewer did not re-run them.

Still open: the V2 gate regression needs a full vLLM install to execute, and three failures in
tests/unit/model_executor/models/test_forward_context.py are observed locally only as whole-file
skips, so their attribution is unconfirmed.

Notes

This supersedes the earlier approach in this PR, which replaced NPUModelRunner.profile_run and disabled compilation for the profile pass. That copy is gone.

Three failures in tests/unit/model_executor/models/test_forward_context.py are observed locally only as whole-file skips (no vLLM on this host). Their attribution is unconfirmed and is not claimed to be excluded here.

[jiangkuaixue123]

Could we fix the compilation boundary in AFD's Attention-side routing instead of replacing the upstream NPUModelRunner.profile_run?

In the pinned stack, native Ascend reads in_profile_run inside the opaque moe_forward / moe_forward_shared custom-op boundary. AFD bypasses AscendRoutedExperts and introduces that Python branch directly in its Attention-side routing path, where tracing can specialize it. This appears to be the boundary we should restore.

disable_compilation does restore the previous state on exit, so this does not permanently disable compilation. However, disabling compilation for the entire profile pass changes the upstream startup compilation timing and profiles the eager path. Copying the upstream method and adding installation, conflict detection, and global import state also makes a small AFD-specific fix harder to understand and maintain.

Please first check whether the profile-dependent expert-ID balancing can live inside an existing AFD opaque custom op; if there is no suitable boundary, consider a minimal routing custom op with the appropriate fake implementation. Keep the in_profile_run read inside that boundary so it is evaluated at execution time, while preserving the upstream profiling/compilation lifecycle. The boundary should stay narrow to limit lost compiler optimizations.

Validation should show that profiling still uses balanced expert IDs, subsequent compiled requests use live routing, and compilation remains enabled, with NPU correctness and performance evidence for the affected path.

[Pingzii]

Thanks, agreed. The fix should restore the compilation boundary in AFD's Attention-side routing rather than replace the upstream profiling lifecycle. I'll first check the existing custom-op boundaries and isolate the profile-dependent expert-ID balancing as narrowly as possible. I'll then remove the runner patch and validate balanced profiling, live routing with compilation enabled, and NPU correctness/performance.

[Pingzii]

Here is the plan we settled on for this PR.

Approach. Rather than replacing the upstream profiling lifecycle, we restore the compilation boundary in AFD's Attention-side routing. The profile-dependent balancing now runs behind an opaque custom op:

def afd_balance_routing_ids_(topk_ids: torch.Tensor, num_experts: int, force_balance: bool) -> None:
    from vllm_ascend.ascend_forward_context import _EXTRA_CTX
    from afd_plugin.routing_balance import _force_balanced_topk_ids
    if force_balance or _EXTRA_CTX.in_profile_run:
        _force_balanced_topk_ids(topk_ids, num_logical_experts=num_experts)

It is registered with direct_register_custom_op(mutates_args=["topk_ids"]) plus a fake implementation, and called from both gate paths. The flag is therefore read at execution time, the gate GEMM and topk stay inside the compiled region, and no part of the profile pass disables compilation. The wait for in_profile_run moved from Python into the op; the two call sites keep their previous semantics (v4: profile only, v2: its existing explicit switch).

What this PR becomes. The profile_run replacement, its install point, the disable_compilation wrapper, the copied upstream method, and the corresponding test and documentation entries are removed (~430 lines deleted). The op is registered during worker initialization through the existing patch installer, so no new import path is introduced.

Validation so far (CPU only). A regression loads the real op through a fixture-scoped registration (infer_schema from the production signature plus _scoped_library), registers the production fake as the Meta kernel, reuses the first compiled callable to emulate a no-guards replay, and keeps the old plain-Python branch as a negative control; opcheck exercises the mutating branch so a wrong mutates_args fails. Targeted tests pass (20) and the maintenance suite passes.

Still open. Device evidence is not collected yet: a balanced profile run, live routing under graph replay with compilation enabled, and correctness/performance. Our test node currently fails to initialize the MoE all-to-all HCCL group on both this candidate and the pre-change baseline, so this is an environment issue being resolved first; I will post the device results here once it is clear.

I will update this PR to the new implementation once the device checks are in. If you would rather see the boundary fix as a fresh PR against current main, say so and I will move it.

[Pingzii]

Device checks are in for the replacement implementation (opaque routing-id boundary, profile_run patch removed). Configuration: single host A5 4A4F, Attention DP4/TP1, compute_gate_on_attention=true, quant_mode=0, automatic routing, no DBO, cudagraph_mode=FULL_DECODE_ONLY, chunk 3072, concurrency 128, 8275 in / 32 out.

1. Profile balancing and live routing through the op

Standalone NPU check against the registered op, read outside any capture:

profile: [[0,1],[2,3],[4,5],[6,7],[0,1],[2,3],[4,5],[6,7]]      # i % num_experts
live[5]: [[5,6],[7,0],[1,2],[3,4],...]  unchanged: True

With in_profile_run=True the op writes the balanced sequence; with it false a non-balanced input is left untouched. The second line is the discriminating case (the first input happened to already be balanced, so it proves nothing on its own).

2. Compilation, capture and replay in the served model

Capturing CUDA graphs (FULL): 100%|...| 5/5 [00:07<00:00, 1.41s/it]
Graph capturing finished in 8 secs, took 0.18 GiB
ACL graph replay is active for the V2 target model (logged once)   # all four workers
POST /v1/completions HTTP/1.1" 200 OK

Capture completed 5/5, matching the five configured capture sizes. The first replay is confirmed by the runtime's message. Per-request replay is not instrumented and I did not add a probe, so I make no claim about it beyond that.

3. Correctness and performance

GSM8K (1,319): 90.14 on this configuration.

Same-day A/B, identical configuration and device layout, only the checkout differs:

Checkout	State	Runs	Throughput
ee9b0ed (branch + operator)	old profile_run patch present	1	311,978
e181f22 (this implementation)	patch removed, boundary op added	2	312,489 / 305,954 (median 309,221)

The medians are 0.9% apart and the best new run is 0.2% above the old one. No regression is observed in this round, but the sample is small (n=1 old, n=2 new) and this does not establish equivalence. Earlier sessions measured 326k-329k for this configuration; the same-day old-state run also lands at 312k, which points at an environmental difference rather than the change.

Limits: one node, one workload, n=1 old arm and n=2 new arm, and this round holds quant_mode=0; the MXFP8 transport path was not exercised here. The standalone check above shows the op's behaviour in both states, but it does not yet provide device evidence for two consecutive inputs reusing a single compiled artifact; that case stays open.

[jiangkuaixue123]

Re-reviewed head 110793d9fbab30ecbb15eee065700c5e8fa9c577. The earlier concern about replacing profile_run is resolved by the narrow routing custom-op boundary. Two items remain:

[P2] Register the op in the existing V2 gate regression test. The new unconditional call at deepseek_v2_attention_gate.py:169 is also reached by test_attention_gate_balances_only_profile_routing. That test directly calls the gate without initializing an AFD worker or requesting registered_routing_ids_op. In a fresh test process with the Ascend dependencies available, the op is therefore unregistered; hosts without those dependencies skip the test and hide the regression. Please use the shared CPU registration fixture for this test, as the V4 tests now do, and run the test individually. This finding is based on the call/registration paths; I did not execute the fork's tests locally.

Remaining NPU validation: the reported device results are from integrated commit e181f22, rather than this PR head, and the PR explicitly leaves reuse of one compiled artifact open. Please add a focused check on the current head and pinned stack: compile through the profile path with balancing enabled, then disable the profile flag and run two distinct, non-balanced live inputs through that same compiled artifact without recompilation. Verify that profiling produces balanced IDs and both live executions preserve their respective routing IDs. For ACL graph coverage, capture with profile mode disabled and replay with the two different inputs, checking routing output for each replay. Record the exact commit, environment, command, and assertions/results. The standalone eager-op check, successful capture, and first-replay log do not establish this specific reuse behavior.

[jiangkuaixue123]

Please fix DCO and pre-commit

[jiangkuaixue123]

Follow-up review of 2a8f1fa96e5dc65532366df568086fd4ced206a7: the V2 gate test now requests the registration fixture, resolving the previous registration omission. Remaining findings and layout suggestions:

1. [P2] Compile under the profile context in the NPU regression. At check_routing_boundary_reuse.py:79, _compile_once() runs with in_profile_run=False. The preceding True context only executes the eager op. This misses the original failure mode: specializing the balancing branch during profiling and reusing that artifact for live requests. Even the old plain-Python branch could pass this scenario. Compile and verify balanced IDs under the True context, then switch to False and reuse that exact callable for both distinct live inputs. Keep the ACL replay check, and report the corrected check's results on the current head and pinned stack.

2. [P2] Consolidate the duplicated registration fixtures. test_routing_ids.py:33 retains a local balance_op fixture that duplicates the schema inference, CPU/Meta registration, module replacement, and restoration in tests/unit/fakes/routing_ids_op.py. Future registration or cleanup changes would need to be applied twice, contrary to the shared-fake convention in AGENTS.md. Use one shared fixture, exposing the production module and op as needed by the different tests.

3. Fold the small balancing helper into the NPU routing module. afd_plugin/routing_balance.py has only one production caller, compat/npu/routing_ids.py, and the helper is short. Moving it out of model code avoids the reverse dependency, but adding a package-root module is unnecessary. Please keep the helper alongside the custom op in compat/npu/routing_ids.py; this preserves the dependency direction and keeps this NPU-specific behavior together without introducing another abstraction.

4. Complete the ownership and test-entry documentation. Keeping the op in compat/npu/routing_ids.py, its unit tests in tests/unit/compat/npu/, and the shared fixture in tests/unit/fakes/ is reasonable. Add the new production module to the relevant docs/design/module/ ownership paths and describe its registration lifecycle. tests/npu/ is also the right location for the device check: if it remains a standalone check_*.py script, document its invocation in tests/npu/README.md; if it is intended as a maintained regression, prefer a discoverable test_routing_ids.py with the appropriate dependency/device guards.

The device evidence still references integrated commit e181f22; I have not seen results for the corrected profile-to-live reuse check on this head. This review was static; I did not execute the fork's code locally.

[jiangkuaixue123]

LGTM