fix(plugin): gate dual-stream MoE by runtime graph mode

[XiaobingSuper]

Problem

In vLLM plugin mode, CUDA/HIP graph capture is owned by vLLM and the concrete runtime mode is selected per forward (FULL, PIECEWISE, or NONE). ATOM's MoE dispatcher instead checks its static compilation config.

For FULL_AND_PIECEWISE, the plugin maps vLLM's compile mode to ATOM level 3, which initializes the ATOM-side static graph mode as PIECEWISE. As a result, decode forwards that vLLM actually dispatches as FULL are incorrectly forced onto the single-stream MoE path. This was observed while validating Kimi-K3 dual-stream shared/routed expert overlap in #1752.

Solution

• Resolve the concrete graph mode from the active frontend's forward context.
• Normalize frontend enums by name instead of relying on matching numeric values.
• Fall back to ATOM's native forward context when vLLM context is unavailable.
• Disable dual-stream MoE only for an actual PIECEWISE forward; preserve it for eager/NONE and whole-model FULL capture.

Validation

• Black formatting check
• Ruff lint check
• Kimi-K3 TP8 with vLLM FULL_DECODE_ONLY and FULL_AND_PIECEWISE
• Profiler confirmed decode dual-stream graph execution on all TP ranks; with the native-matched online quant configuration, measured stream overlap was 99.1%–100%

[github-actions]

🏷️ CI Guide

Runs automatically on every eligible PR before approval:

• ✅ Pre Checkin: Black, Ruff, catalog schema validation, non-GPU unit tests

Heavy model tests:

• ✅ Run after the PR is approved and Pre Checkin passes
• ✅ Run immediately when an approval review is submitted
• ✅ Can be requested before approval with labels

Label	Tests
ci:full	Run all heavy PR model tests: native ATOM, vLLM, and SGLang
ci:atom	Run native ATOM model accuracy tests
ci:vllm	Run ATOM vLLM OOT model accuracy tests
ci:sglang	Run ATOM SGLang model accuracy tests

Heavy jobs are skipped when the PR is not approved and no matching ci:* label is present.
Add labels via the sidebar or gh pr edit 1791 --add-label <label>

[Copilot]

Pull request overview

This PR fixes dual-stream MoE gating in vLLM plugin mode by switching from ATOM’s static compilation config to the per-forward runtime cudagraph mode selected by the active frontend (vLLM or native ATOM). This ensures decode forwards that vLLM runs under FULL capture are not incorrectly forced onto the single-stream MoE path when ATOM is configured with composite cudagraph modes (e.g., FULL_AND_PIECEWISE).

Changes:

• Add a frontend-aware helper (get_current_cudagraph_runtime_mode) that resolves the concrete runtime graph mode by normalizing enum values by name (NONE/PIECEWISE/FULL), with safe fallback to native ATOM context and defaulting unknown/unavailable to NONE.
• Update dual-stream MoE dispatch gating to disable dual-stream only for actual PIECEWISE runtime forwards (allowing eager/NONE and whole-model FULL capture).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
atom/utils/forward_context.py	Adds normalization + runtime-mode resolution across vLLM/native forward contexts, defaulting to NONE when unavailable/unknown.
atom/model_ops/module_dispatch_ops.py	Switches dual-stream MoE safety gating from static compilation cudagraph config to concrete per-forward runtime mode (disables only for PIECEWISE).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated no new comments.

Suppressed comments (1)

atom/utils/forward_context.py:651

• New runtime-mode normalization + vLLM fallback logic is now gating dual-stream MoE behavior, but there’s no regression test covering the key cases (vLLM context present with NONE/PIECEWISE/FULL by name, context unavailable, and composite modes being rejected). This is easy to regress and would silently flip MoE dispatch decisions in plugin mode.

def get_current_cudagraph_runtime_mode() -> CUDAGraphMode:
    """Return the concrete graph mode for the active model forward.

    In vLLM plugin mode graph capture/replay is owned by vLLM, so its forward
    context is authoritative.  Native ATOM records the same decision on its
    own ForwardContext.  An unavailable/unknown context is treated as NONE:
    eager dual-stream execution is valid, and some vLLM runners expose NONE
    while a whole-model FULL graph is being captured.  Replay does not execute
    this Python dispatcher.
    """
    from atom.plugin import is_vllm

    if is_vllm():
        try:
            from vllm.forward_context import (
                get_forward_context as get_vllm_forward_context,
            )
            from vllm.forward_context import (
                is_forward_context_available,
            )

            if is_forward_context_available():
                mode = _normalize_cudagraph_runtime_mode(
                    get_vllm_forward_context().cudagraph_runtime_mode
                )
                if mode is not None:
                    return mode
        except (ImportError, AttributeError, AssertionError):
            pass

    mode = _normalize_cudagraph_runtime_mode(
        getattr(get_forward_context(), "cudagraph_runtime_mode", None)
    )
    return mode if mode is not None else CUDAGraphMode.NONE