Skip to content

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

Closed
wuhuikx wants to merge 1 commit into
mainfrom
fix/vllm-runtime-cudagraph-dual-stream
Closed

fix(plugin): gate dual-stream MoE by runtime graph mode#1790
wuhuikx wants to merge 1 commit into
mainfrom
fix/vllm-runtime-cudagraph-dual-stream

Conversation

@wuhuikx

@wuhuikx wuhuikx commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator

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 checked 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.
  • Add focused tests for enum translation, frontend fallback, runtime dispatch, token threshold, TBO, and the existing _use_dual_stream gate.

Validation

  • python -m pytest -q tests/test_module_dispatch_ops.py tests/plugin/test_plugin_config_translation.py tests/plugin/test_vllm_kimi_k3.py — 22 passed
  • 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%

Made with Cursor

Use the frontend's per-forward graph decision so FULL decode can retain overlap while PIECEWISE capture remains safely single-stream.

Co-authored-by: Cursor <cursoragent@cursor.com>
@github-actions

github-actions Bot commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

🏷️ 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 1790 --add-label <label>

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ [ruff] <I001> reported by reviewdog 🐶
Import block is un-sorted or un-formatted

import logging
import threading
from contextlib import contextmanager
from dataclasses import dataclass, field, fields
from enum import Enum
from typing import Any, Union
import numpy as np
import torch
from atom.config import CUDAGraphMode, Config, KVCacheTensor, ParallelConfig

@wuhuikx

wuhuikx commented Aug 4, 2026

Copy link
Copy Markdown
Collaborator Author

Closing to resubmit the change from @XiaobingSuper.

@wuhuikx wuhuikx closed this Aug 4, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants