[Bugfix][MoE] Restore TP-sharded routed-expert capture for DP=1

[Ronald1995]

Overview

Restore complete routed-expert rows when model-level sequence parallelism shards tokens across TP ranks and DP=1 has no DP metadata. Related platform-specific work: vllm-project/vllm-ascend#18502.

Claims

• Identify TP shards using the existing BatchDescriptor.num_tokens (or DPMetadata for multi-DP), gather in rank order, and trim SP padding.
• Reuse the descriptor in eager, graph, microbatch and decoder-replay contexts; no extra ForwardContext.num_tokens field or count-only caller arguments.
• Preserve multi-DP layout priority and reject malformed shard grouping explicitly, including under optimized Python.

Validation

.venv/bin/python -m pytest -q tests/model_executor/test_routed_experts_capture.py tests/models/test_deepseek_v41_replay_batch.py
pre-commit run --files tests/model_executor/test_routed_experts_capture.py tests/models/test_deepseek_v41_replay_batch.py vllm/forward_context.py vllm/model_executor/layers/fused_moe/routed_experts_capturer.py vllm/models/deepseek_v41/nvidia/model_state.py vllm/v1/worker/gpu/ubatch_utils.py vllm/v1/worker/gpu_ubatch_wrapper.py
git diff --check

• 62 pytest cases passed on GB200, using isolated test-root collection (--confcutdir) against a nightly-image overlay. Includes real sp_padding_mask mutation and compact V4.1 replay contexts. A mask-length implementation fails both new TP-rank regressions.
• Real NCCL/CUDA-graph layout checks passed for TP2/DP1, TP4/DP1 and TP2/DP2: 36 cases per rank, including odd/empty counts, unequal DP counts and buffer reuse.
• Applicable pre-commit hooks (including mypy and test tethering) and diff checks passed.
• Final helper cleanup: 51 capture tests passed locally, with 30 targeted cases rerun during review. 89,469 CPU layout cases compared the previous and simplified production helper in both normal and optimized Python: identical rows, strides, exception types and collective calls. Collective outputs were simulated in this differential check; the GPU/model results below precede this cleanup.
• Historical CPU checks: 25 isolated capture cases (also 25 under optimize=1) and 96,203 multi-DP differential cases passed. Earlier model reproducer and commands: discussion.

Real-weight DeepSeek-V4.1-Flash A/B

Control: fe18b045fd; candidate: this descriptor simplification. GB200, TP2/DP1/EP, deep_gemm_mega_moe, default padding skipping, async scheduling, prefix caching, chunked prefill, SHM/value delivery. Both eager and FULL_AND_PIECEWISE were exercised.

• Five model runs, including a control reload: 41,315 returned rows / 9,915,600 expert IDs, with no missing or mismatched IDs. Compare raw per-rank routes to capture snapshots and HTTP outputs; prefix hits use the existing first-published immutable Store values. Multi-chunk prefill, decode, a measured cache hit and actual graph replays were verified.
• Fixed-token scoring: three sequences (33/257/1025 tokens), three uncached repeats each. Short-sequence logprobs were exactly equal. The long sequence varies even across control runs: same-GPU control/candidate mean absolute logprob difference 0.05360, versus 0.05179 between control initializations. Mean NLL shift was 0.00174, versus control-reload shift 0.00459. No clear additional degradation in these bounded checks; not a claim of bitwise equivalence or broad benchmark accuracy. R3-disabled serving smoke checks also passed.
• DBO is not fixed: real V2 UBatchRunner/TP collectives/graph checks reproduce the same existing microbatch overwrite in both variants. V4.1 Engram itself rejects DBO, so these are separate checks, not a claimed V4.1+DBO model pass. Offset repair remains outside this PR ([Prototype][MRV2] Stage DeepSeek decode tokens for imbalanced DBO graph replay #57184).

Environment: official nightly vllm/vllm-openai@sha256:a7deff2c216def98c193352aca50fb03a2ee29231355b3033dc3789009f6d23e, V4.1 snapshot 2cba9e42aa026125f3ed06c6d98c1db82f7ca027. Only the affected production Python files were overlaid per variant; not a full checkout rebuild. Model runs used --max-model-len 8192 --max-num-batched-tokens 2048 --max-num-seqs 16 --gpu-memory-utilization 0.92, concurrent prompt lengths 33/1024/2051/4096 and 8 output tokens, then repeated the 1024-token prefix. VLLM_DEEP_GEMM_WARMUP=skip bypassed an existing image warmup API mismatch; normal JIT and graph capture still ran. Correctness instrumentation makes these runs unsuitable for throughput claims.

Details

The previous DP=1 path wrote only the received TP shard, leaving later routing rows stale. Restore the local token axis before buffer assignment. Layer bounds are checked before collectives and gather inputs are contiguous.

The pre-SP count must survive model-level sharding: is_padding can itself become a TP shard and is not a reliable count source. The existing batch descriptor retains the original size, including in compact decoder-replay contexts. Creating that descriptor for eager execution removes duplicate state and three count-only caller edits (production diff for this refinement: net -11 lines).

This updates the existing TP fix rather than opening a duplicate. Microbatch write offsets (#57184) and decoder-replay row retention (#61643) remain separate concerns, unchanged here. The restoration helper shares one layout error, reads the declared DPMetadata.local_sizes field directly, and retains the existing layout priority. Tests use actual DPMetadata objects.

AI assistance was used for implementation, review and validation.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill.
• AI assistance was used during the creation of this PR.
• Design Fit: Reuses existing metadata; no count-only changes to model runners or model state remain relative to base.
• Testing and Validation: Focused regressions, distributed checks and model A/B completed; numerical and DBO limits stated above.
• Code Quality and Style: Applicable pre-commit checks, including mypy, and diff checks passed.
• Pull Request Contents: Root cause, duplicate-work scope, evidence and AI assistance are disclosed.

[Ronald1995]

Model end-to-end reproducer for the DP=1 sharded capture (1/2): the test case

This is the model-level reproducer for the "Required model end-to-end reproducer" table. It is deliberately not collected by CI (the filename does not match pytest's python_files), it needs two GPUs, and it is run by path.

What it asserts. For every request it rebuilds the whole batch from each TP rank's raw pre-capture shard (recorded by the probe in the next comment), concatenates the shards in TP-rank order, trims to the pre-SP forward token count, and compares every token/layer/expert id against the routed_experts array returned by /v1/completions. It also asserts that the publishing rank's capture buffer holds that same batch, which separates a capture bug from an auxiliary-output bug. "All ids in range" / "rows are nonzero" are not used as criteria.

Cases covered (PR table + the remaining DP=1 branch of the capturer):

test	branch under test
test_six_token_prefill_captures_every_tp_shard	n == ceil(T/tp) (the bug)
test_three_token_prefill_drops_only_sp_padding	n == ceil(T/tp) with ceil-div padding
test_second_uncached_batch_has_no_stale_rows	n == ceil(T/tp), buffer reuse
test_single_token_steps_capture_without_gather	T == 1, no capture-only gather
test_cuda_graph_matches_eager_reference	CUDA graph (opt-in)
test_ordinary_moe_path_captures_full_batch_without_gather (+ 6b)	n == T control, TP=2+EP vs TP=1 oracle
test_complete_padded_batch_is_trimmed_not_regathered	n == ceil(T/tp)*tp (already-gathered padded batch)
test_unexpected_capture_batch_fails_closed	fail-closed on an unrecognized layout

How to run. With a real checkpoint on a Blackwell host (all three MEGA_MOE_BACKENDS members require SM100/SM103):

export VLLM_ROUTED_EXPERTS_E2E_MODEL=/path/to/deepseek-v4-checkpoint
export VLLM_ROUTED_EXPERTS_E2E_BACKEND=deep_gemm_mega_moe
uv run --no-sync python -m pytest -q -s tests/models/deepseek_v41_routed_experts_capture_e2e.py

Without a checkpoint (weights-free tiny model, any two GPUs) the sharded layout is supplied by the probe's shard injection, because no MegaMoE backend can bind below SM100/SM103:

export VLLM_ROUTED_EXPERTS_E2E_DUMMY=1
export VLLM_ROUTED_EXPERTS_E2E_EMULATE_SP=1
uv run --no-sync python -m pytest -q -s tests/models/deepseek_v41_routed_experts_capture_e2e.py

Base-vs-fix: run the same file by absolute path while the base revision supplies vllm (cd <base worktree>; PYTHONPATH=<base> ... python -m pytest --import-mode=importlib <abs path to this file>); cases 1-3 must fail there and pass on the fix.

The module needs tests/models/routed_experts_capture_probe.py (posted in the next comment).

# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end model reproducer for TP-sharded routed-expert capture at DP=1.

This is the model-level reproducer that vLLM PR #61388 ("[Bugfix][MoE]
Restore TP-sharded routed-expert capture for DP=1") lists as pending, outside
CI: it needs a DeepSeek V4/V4.1 checkpoint and two GPUs, so it is deliberately
not collected by any Buildkite job (the filename does not match pytest's
``python_files``) and must be run by path.

The bug
-------
With ``data_parallel_size == 1``, ``ParallelConfig.use_sequence_parallel_moe``
is false, so the forward context carries no ``dp_metadata``. DeepSeek V4/V4.1
(and Kimi K3) still shard tokens across TP *at the model level* for their
MegaMoE backends - ``_use_sequence_parallel`` only needs EP, TP > 1 and a
MegaMoE backend. The capture callback therefore receives ``ceil(N / TP)`` rows,
and the pre-fix capturer wrote exactly those rows into the DP-1 buffer, leaving
rows ``[ceil(N / TP), N)`` holding whatever an earlier forward left there.
Only TP rank 0 publishes the auxiliary output, so the stale rows reached the
client as if they were the second half of the batch.

What this checks
----------------
For each request it compares the ``routed_experts`` array returned by
``/v1/completions`` against a reference rebuilt from every TP rank's *raw*
pre-capture shard (see ``tests/models/routed_experts_capture_probe.py``):
shards are concatenated in TP-rank order and trimmed to the pre-SP forward
token count, which is exactly the batch the captured rows must describe. It
also asserts that the publishing rank's capture buffer holds that same batch,
which separates a capture bug from an auxiliary-output plumbing bug. Every
token/layer/expert id is compared; "all ids are in range" or "rows are
nonzero" are not used as criteria.

Cases (the rows of the PR's required-model-validation table):

1. six-token prefill, TP=2, EP, DP=1 - all six rows equal the shard reference;
2. three-token prefill - the fourth, SP-padded row is trimmed away;
3. two different six-token batches on one worker - the second batch keeps no
   stale ids from the first;
4. single-token prefill and decodes - no capture-only gather, TP0's rows are
   what is returned;
5. CUDA graph execution (env-gated by ``VLLM_ROUTED_EXPERTS_E2E_GRAPH=1``) -
   returned real-token rows match the eager reference, no hang;
6. complete-batch control on an ordinary MoE path (env-gated by
   ``VLLM_ROUTED_EXPERTS_E2E_CONTROL=1``, default on) - TP=2 + EP + DP=1 with
   no model-level token SP captures the whole batch without a capture gather,
   and its routes equal a TP=1 run of the same model.

How to run
----------
With a checkpoint whose MoE backend supports capture (Blackwell only, see
below)::

    export VLLM_ROUTED_EXPERTS_E2E_MODEL=/path/to/deepseek-v4-checkpoint
    export VLLM_ROUTED_EXPERTS_E2E_BACKEND=deep_gemm_mega_moe   # or
        # flashinfer_moe_ep_mega_deep_gemm / flashinfer_moe_ep_cutedsl
    export VLLM_ROUTED_EXPERTS_E2E_EXTRA_ARGS='--trust-remote-code ...'
    uv run --no-sync python -m pytest -q -s \
        tests/models/deepseek_v41_routed_experts_capture_e2e.py

Without any checkpoint (weights-free tiny model, two GPUs of any arch)::

    export VLLM_ROUTED_EXPERTS_E2E_DUMMY=1            # tiny Mixtral, built in
    export VLLM_ROUTED_EXPERTS_E2E_EMULATE_SP=1       # see "Where it can run"
    uv run --no-sync python -m pytest -q -s \
        tests/models/deepseek_v41_routed_experts_capture_e2e.py

Run the same file against the base commit to reproduce the bug. The probe and
this file are revision-independent, so run *this* copy of the test by absolute
path while the base revision supplies ``vllm`` - from the base worktree, so its
``vllm`` wins over the running directory::

    fixed=<worktree with this file>
    base=<base revision worktree>
    cd "$base"
    PYTHONPATH=$base uv run --no-sync python -m pytest -q -s \
        --import-mode=importlib \
        "$fixed"/tests/models/deepseek_v41_routed_experts_capture_e2e.py

With a real MegaMoE model, cases 1-3 must fail on the base (partial batch
written) and pass on the fix; case 4 is a control that passes on both.

Where it can run, and what ``EMULATE_SP`` is for
------------------------------------------------
All three members of ``MEGA_MOE_BACKENDS`` refuse to bind below SM100/SM103:
``get_deep_gemm_mega_moe_backend`` raises for anything but SM100
(``vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py``) and
``FlashInferMoeEpBackendSpec.supports_current_device`` accepts only (10, 0) and
(10, 3) (``vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py``). So a
Hopper/Ampere host cannot start the model-level sharding the DP=1 bug lives in,
whatever the weights are. ``VLLM_ROUTED_EXPERTS_E2E_EMULATE_SP=1`` covers that
gap: the probe slices the capture input into ``ceil(T / tp)``-row TP shards,
padded the way ``sp_reduce_scatter`` pads, so the real capture code, the real
TP all-gather and the whole engine path run on the layout the bug is about.
Everything else stays real; only "which rows the capture is handed" is
emulated, and the run says so in every summary it prints.

Environment
-----------
``VLLM_ROUTED_EXPERTS_E2E_MODEL``
    Checkpoint to serve. Unset with ``..._DUMMY`` set builds a tiny local one.
``VLLM_ROUTED_EXPERTS_E2E_DUMMY`` / ``..._DUMMY_ARCH``
    Build a weights-free model (``mixtral``, default, any GPU; ``deepseek_v4``,
    Blackwell only) with ``--load-format dummy``.
``VLLM_ROUTED_EXPERTS_E2E_EMULATE_SP``
    Slice the capture input into TP shards (see above).
``VLLM_ROUTED_EXPERTS_E2E_BACKEND`` (default ``deep_gemm_mega_moe``)
``VLLM_ROUTED_EXPERTS_E2E_MEGA_MOE`` (``0`` disables the ``--moe-backend``)
``VLLM_ROUTED_EXPERTS_E2E_TP_SIZE`` (default ``2``)
``VLLM_ROUTED_EXPERTS_E2E_EXTRA_ARGS``
    Extra ``vllm serve`` flags, shell-split (e.g. ``--trust-remote-code``).
``VLLM_ROUTED_EXPERTS_E2E_PROBE_MODE``
    ``worker-extension`` (default) injects the probe through vLLM's
    ``--worker-extension-cls``; ``sitecustomize`` uses a ``sitecustomize.py``
    shim instead, if the extension route does not produce dumps.
``VLLM_ROUTED_EXPERTS_E2E_GRAPH`` / ``_CONTROL`` / ``_SYNC_CHECK`` / ``_FAIL_CLOSED``
    Enable cases 5 / 6 / the ``VLLM_GPU_SYNC_CHECK=error`` run / the
    fail-closed run (6 and fail-closed are on by default). The sync-check run
    has no probe, because dumping synchronizes.
``VLLM_ROUTED_EXPERTS_E2E_CONTROL_MODEL``
    Model for case 6; unset uses the dummy Mixtral.
``VLLM_ROUTED_EXPERTS_E2E_ARTIFACT_DIR``
    Where dumps and probe notes go; defaults to a fresh temp directory that is
    printed at the start of the run.
``VLLM_ROUTED_EXPERTS_E2E_STARTUP_TIMEOUT`` (default ``1800`` seconds)
``VLLM_ROUTED_EXPERTS_E2E_REQUEST_TIMEOUT`` (default ``300`` seconds)

Not covered here: NCCL/CUDA-graph correctness of the collective itself beyond
the row comparison above, MoE numerics, and throughput. Prefix caching stays
enabled (AuxOutput requires it) and prompts are distinct and uncached, as the
PR requires.
"""

from __future__ import annotations

import contextlib
import io
import os
import shlex
import struct
import tempfile
import uuid
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np
import pybase64
import pytest

from tests.utils import RemoteOpenAIServer

MODEL = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_MODEL", "")
MOE_BACKEND = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_BACKEND", "deep_gemm_mega_moe")
TP_SIZE = int(os.environ.get("VLLM_ROUTED_EXPERTS_E2E_TP_SIZE", "2"))
CONTROL_MODEL = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_CONTROL_MODEL", "")
EXTRA_ARGS = shlex.split(os.environ.get("VLLM_ROUTED_EXPERTS_E2E_EXTRA_ARGS", ""))
PROBE_MODE = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_PROBE_MODE", "worker-extension")
RUN_GRAPH = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_GRAPH", "0") == "1"
RUN_CONTROL = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_CONTROL", "1") == "1"
RUN_SYNC_CHECK = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_SYNC_CHECK", "0") == "1"
RUN_FAIL_CLOSED = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_FAIL_CLOSED", "1") == "1"
RUN_TP1_ORACLE = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_TP1_ORACLE", "1") == "1"
# Weights-free local model: no checkpoint, `--load-format dummy` supplies
# finite random weights (packed FP4 bytes and E8M0 scales included).
DUMMY_ARCH = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_DUMMY_ARCH", "")
DUMMY = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_DUMMY", "0") == "1" or bool(DUMMY_ARCH)
# Emulate the model-level token sharding the DP=1 branch is about by slicing the
# capture input into TP shards in the probe. Required on any GPU where no
# MegaMoE backend can bind: `deep_gemm_mega_moe` refuses anything but SM100
# (`deep_gemm_mega_moe.py` get_deep_gemm_mega_moe_backend) and both FlashInfer
# MoE-EP backends require SM100/SM103 (`flashinfer_moe_ep.py`
# FlashInferMoeEpBackendSpec.supports_current_device). The capture code, the TP
# gather and the rest of the serving path stay real.
EMULATE_SP = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_EMULATE_SP", "0") == "1"
REQUEST_TIMEOUT = float(
    os.environ.get("VLLM_ROUTED_EXPERTS_E2E_REQUEST_TIMEOUT", "300")
)
STARTUP_TIMEOUT = float(
    os.environ.get("VLLM_ROUTED_EXPERTS_E2E_STARTUP_TIMEOUT", "1800")
)
ARTIFACT_DIR = Path(
    os.environ.get("VLLM_ROUTED_EXPERTS_E2E_ARTIFACT_DIR")
    or tempfile.mkdtemp(prefix="routed-experts-capture-e2e-")
)

_PROBE_MODULE = Path(__file__).with_name("routed_experts_capture_probe.py")
_MAGIC = b"RXD1"
_HEADER = struct.Struct("<4sqiiiiiiiii")
# Slots keep prompts disjoint, so prefix caching never serves one request's
# routing to another and every request is a genuine capture.
_PROMPT_BASE = 1000
_WARMUP_SLOT = 99

pytestmark = [
    pytest.mark.distributed,
    pytest.mark.skipif(
        not MODEL and not DUMMY,
        reason=(
            "set VLLM_ROUTED_EXPERTS_E2E_MODEL to a DeepSeek V4/V4.1 "
            "checkpoint, or VLLM_ROUTED_EXPERTS_E2E_DUMMY=1 to generate a "
            "weights-free tiny model"
        ),
    ),
]


class _HarnessError(AssertionError):
    """The run cannot produce a reference; reported as a failure, not a skip."""


@dataclass(frozen=True)
class _DummyArch:
    """A local weights-free model: config.json + a generated tokenizer only."""

    config: dict[str, Any]
    serve_args: tuple[str, ...] = ("--load-format", "dummy")
    tokenizer_vocab: int = 8192
    # Model-level token sharding needs a MegaMoE backend on this architecture.
    needs_mega_moe: bool = False


# The published DeepSeek-V4-Flash config, shrunk to two layers / eight experts.
# Kept close to the release geometry (MLA head_dim 512, sliding window 128) so
# the sparse-attention kernels stay on their supported shapes. Every MegaMoE
# backend is SM100/SM103-only, so this entry only runs on Blackwell.
_DEEPSEEK_V4_DUMMY = {
    "architectures": ["DeepseekV4ForCausalLM"],
    "model_type": "deepseek_v4",
    "attention_bias": False,
    "attention_dropout": 0.0,
    "bos_token_id": 0,
    "eos_token_id": 1,
    "expert_dtype": "fp4",
    "hc_eps": 1e-05,
    "hc_mult": 2,
    "hc_sinkhorn_iters": 3,
    "head_dim": 512,
    "hidden_act": "silu",
    "hidden_size": 256,
    "index_head_dim": 128,
    "index_n_heads": 8,
    "index_topk": 64,
    "max_position_embeddings": 512,
    "moe_intermediate_size": 256,
    "n_routed_experts": 8,
    "n_shared_experts": 0,
    "norm_topk_prob": True,
    "num_attention_heads": 8,
    "num_experts_per_tok": 2,
    "num_hidden_layers": 2,
    "num_hash_layers": 0,
    "num_key_value_heads": 1,
    "num_nextn_predict_layers": 0,
    "o_groups": 8,
    "o_lora_rank": 128,
    "q_lora_rank": 128,
    "qk_rope_head_dim": 64,
    "rms_norm_eps": 1e-06,
    "rope_scaling": {"rope_type": "deepseek_yarn", "factor": 1.0},
    "rope_theta": 10000,
    "routed_scaling_factor": 1.0,
    "scoring_func": "sqrtsoftplus",
    "sliding_window": 128,
    "swiglu_limit": None,
    "tie_word_embeddings": False,
    "topk_method": "noaux_tc",
    "torch_dtype": "bfloat16",
    "use_cache": True,
    "vocab_size": 129280,
    "compress_ratios": [0, 0],
    "kv_source_layer_ids": [0],
    "index_source_layer_ids": [],
    "vision_config": {"num_hidden_layers": 0},
}

# A small Mixtral: the ordinary MoE path (TP+EP, no model-level token SP), so it
# serves on any 2 GPUs and doubles as the complete-batch control.
_MIXTRAL_DUMMY = {
    "architectures": ["MixtralForCausalLM"],
    "model_type": "mixtral",
    "hidden_act": "silu",
    "hidden_size": 256,
    "intermediate_size": 512,
    "max_position_embeddings": 512,
    "num_attention_heads": 4,
    "num_experts_per_tok": 2,
    "num_hidden_layers": 2,
    "num_key_value_heads": 2,
    "num_local_experts": 8,
    "rms_norm_eps": 1e-06,
    "rope_theta": 1000000.0,
    "sliding_window": None,
    "tie_word_embeddings": False,
    "torch_dtype": "bfloat16",
    "use_cache": True,
    "vocab_size": 8192,
    "router_jitter_noise": 0.0,
}

_DUMMY_ARCHES = {
    "deepseek_v4": _DummyArch(
        config=_DEEPSEEK_V4_DUMMY,
        serve_args=(
            "--load-format",
            "dummy",
            "--dtype",
            "bfloat16",
            "--limit-mm-per-prompt",
            '{"image": 0}',
            "--block-size",
            "64",
        ),
        tokenizer_vocab=8192,
        needs_mega_moe=True,
    ),
    "mixtral": _DummyArch(config=_MIXTRAL_DUMMY, tokenizer_vocab=8192),
}


def _prompt(slot: int, count: int) -> tuple[int, ...]:
    """Distinct, uncached, tokenizer-valid ids for one request."""
    base = _PROMPT_BASE + slot * 32
    return tuple(base + index for index in range(count))


@dataclass(frozen=True)
class _Record:
    seq: int
    tp_rank: int
    dp_rank: int
    layer_id: int
    num_tokens: int
    num_padding: int
    gathers: int
    ids: np.ndarray
    buffer_rows: np.ndarray


@dataclass
class _Step:
    index: int
    num_tokens: int
    num_padding: int
    layer_ids: tuple[int, ...]
    shards: dict[int, dict[int, np.ndarray]] = field(default_factory=dict)
    buffers: dict[int, dict[int, np.ndarray]] = field(default_factory=dict)
    gathers: dict[tuple[int, int], int] = field(default_factory=dict)

    @property
    def num_real(self) -> int:
        return self.num_tokens - self.num_padding

    @property
    def first_layer(self) -> int:
        return self.layer_ids[0]

    def shard_sizes(self, layer_id: int) -> dict[int, int]:
        return {rank: rows.shape[0] for rank, rows in self.shards[layer_id].items()}


@dataclass
class _Trace:
    steps: list[_Step]
    bounds: list[tuple[int, int]]
    sharded: list[bool]
    trimmed: list[bool]
    gathers: list[int]
    reference: dict[int, np.ndarray]
    published: dict[int, np.ndarray]
    publish_rank: int

    @property
    def total_rows(self) -> int:
        return sum(end - start for start, end in self.bounds)

    @property
    def max_gathers(self) -> int:
        return max(self.gathers, default=0)

    @property
    def layers(self) -> list[int]:
        return sorted(self.reference)


@dataclass(frozen=True)
class _ServerSpec:
    name: str
    model: str
    args: tuple[str, ...]
    env: dict[str, str]
    probe: bool


@dataclass
class _Result:
    tokens: tuple[int, ...]
    prompt_tokens: int
    generated: int
    routed: np.ndarray
    records: list[_Record]
    case: str
    tag: str | None = None

    @property
    def expected_rows(self) -> int:
        return self.prompt_tokens + self.generated - 1


def _write_tiny_tokenizer(model_dir: Path, vocab_size: int) -> None:
    """A byte-level-free WordLevel tokenizer, enough for token-id prompts."""
    import json

    from tokenizers import Tokenizer, models, pre_tokenizers

    vocab = {"<unk>": 0, "<s>": 1, "</s>": 2, "<pad>": 3}
    vocab.update({f"tok{index}": index for index in range(4, vocab_size)})
    tokenizer = Tokenizer(models.WordLevel(vocab=vocab, unk_token="<unk>"))
    tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
    tokenizer.save(str(model_dir / "tokenizer.json"))
    (model_dir / "tokenizer_config.json").write_text(
        json.dumps(
            {
                "tokenizer_class": "PreTrainedTokenizerFast",
                "unk_token": "<unk>",
                "bos_token": "<s>",
                "eos_token": "</s>",
                "pad_token": "<pad>",
                "model_max_length": 512,
                "clean_up_tokenization_spaces": False,
            }
        )
    )


def _dummy_model_dir(arch: str) -> Path:
    """Materialize a weights-free model dir (config + tokenizer, no weights)."""
    import json

    model_dir = ARTIFACT_DIR / f"dummy-{arch}"
    if (model_dir / "config.json").exists():
        return model_dir
    model_dir.mkdir(parents=True, exist_ok=True)
    spec = _DUMMY_ARCHES[arch]
    (model_dir / "config.json").write_text(json.dumps(spec.config, indent=2))
    _write_tiny_tokenizer(model_dir, spec.tokenizer_vocab)
    (model_dir / "generation_config.json").write_text(
        json.dumps(
            {
                "bos_token_id": 1,
                "eos_token_id": 2,
                "pad_token_id": 3,
                "do_sample": False,
            }
        )
    )
    print(f"[routed-experts-e2e] wrote weights-free {arch} model to {model_dir}")
    return model_dir


def _target_model() -> tuple[str, tuple[str, ...], bool]:
    """(model path, extra serve args, whether it uses a MegaMoE backend)."""
    override = os.environ.get("VLLM_ROUTED_EXPERTS_E2E_MEGA_MOE", "")
    if MODEL:
        mega = override != "0"
        return MODEL, (), mega
    arch = DUMMY_ARCH or "mixtral"
    spec = _DUMMY_ARCHES[arch]
    # Emulation replaces the model-level sharding, so no backend is needed.
    mega = spec.needs_mega_moe and not EMULATE_SP and override != "0"
    return str(_dummy_model_dir(arch)), spec.serve_args, mega


def _control_model() -> tuple[str, tuple[str, ...]]:
    if CONTROL_MODEL:
        return CONTROL_MODEL, ()
    if DUMMY:
        spec = _DUMMY_ARCHES["mixtral"]
        return str(_dummy_model_dir("mixtral")), spec.serve_args
    return "", ()


def _moe_serve_args(
    *,
    tp_size: int,
    expert_parallel: bool,
    backend: str | None,
    enforce_eager: bool,
) -> tuple[str, ...]:
    args = [
        "--tensor-parallel-size",
        str(tp_size),
        "--data-parallel-size",
        "1",
        "--max-num-seqs",
        "1",
        "--max-model-len",
        "256",
        "--max-num-batched-tokens",
        "256",
        "--enable-return-routed-experts",
    ]
    if expert_parallel:
        args += ["--enable-expert-parallel"]
    if backend is not None:
        args += ["--moe-backend", backend]
    if enforce_eager:
        args += ["--enforce-eager"]
    return tuple(args) + tuple(EXTRA_ARGS)


def _main_spec(
    *, enforce_eager: bool, probe: bool = True, inject: str | None = None
) -> _ServerSpec:
    env = {"VLLM_USE_V2_MODEL_RUNNER": "1"}
    if inject is None and EMULATE_SP:
        # Force the DP=1 sharded layout the fix is about; see EMULATE_SP.
        inject = "shard"
    if inject:
        env["VLLM_ROUTED_EXPERTS_E2E_INJECT"] = inject
    model, model_args, mega = _target_model()
    name = "eager" if enforce_eager else "graph"
    return _ServerSpec(
        name=f"main-{name}" + (f"-{inject}" if inject else ""),
        model=model,
        args=_moe_serve_args(
            tp_size=TP_SIZE,
            expert_parallel=True,
            backend=MOE_BACKEND if mega else None,
            enforce_eager=enforce_eager,
        )
        + model_args,
        env=env,
        probe=probe,
    )


def _sync_check_spec() -> _ServerSpec:
    spec = _main_spec(enforce_eager=True, probe=False)
    env = dict(spec.env)
    env["VLLM_GPU_SYNC_CHECK"] = os.environ.get("VLLM_GPU_SYNC_CHECK", "error")
    return _ServerSpec(
        name="main-sync-check",
        model=spec.model,
        args=spec.args,
        env=env,
        probe=False,
    )


def _injected_spec(mode: str) -> _ServerSpec:
    """Server whose capture input is rewritten to force one layout branch."""
    return _main_spec(enforce_eager=True, inject=mode)


def _control_spec(*, tp_size: int, probe: bool) -> _ServerSpec:
    """Ordinary MoE path: TP+EP but no model-level token sharding."""
    model, model_args = _control_model()
    args = _moe_serve_args(
        tp_size=tp_size,
        expert_parallel=tp_size > 1,
        backend=None,
        enforce_eager=True,
    )
    # tiny-mixtral publishes a sliding window that AuxOutput retention rejects.
    if "TitanML/tiny-mixtral" in model:
        args += ("--hf-overrides", '{"sliding_window": null}')
    return _ServerSpec(
        name=f"control-tp{tp_size}",
        model=model,
        args=args + model_args,
        env={"VLLM_USE_V2_MODEL_RUNNER": "1"},
        probe=probe,
    )


class _Harness:
    """Runs one server at a time (each spec needs every GPU) and reads dumps."""

    def __init__(self, artifact_dir: Path) -> None:
        assert PROBE_MODE in ("worker-extension", "sitecustomize"), (
            f"VLLM_ROUTED_EXPERTS_E2E_PROBE_MODE={PROBE_MODE!r} is not supported"
        )
        self.root = artifact_dir
        self.dump_dir = artifact_dir / "dumps"
        self.dump_dir.mkdir(parents=True, exist_ok=True)
        self.tag_path = self.dump_dir / "current_tag"
        self.shim_dir = artifact_dir / "probe_shim"
        self._write_probe_shim()
        self._stack = contextlib.ExitStack()
        self._client = None
        self._active: str | None = None
        self._requests = 0
        print(
            f"[routed-experts-e2e] artifacts: {self.root}\n"
            f"[routed-experts-e2e] probe mode: {PROBE_MODE}"
        )

    def close(self) -> None:
        self.tag_path.unlink(missing_ok=True)
        self._stack.close()
        self._client = None
        self._active = None

    def _write_probe_shim(self) -> Path:
        self.shim_dir.mkdir(parents=True, exist_ok=True)
        (self.shim_dir / "vllm_re_capture_probe.py").write_text(
            "import importlib.util as _ilu\n\n"
            f"_PATH = {str(_PROBE_MODULE)!r}\n"
            "_NAME = 'vllm_re_capture_probe_impl'\n"
            "_spec = _ilu.spec_from_file_location(_NAME, _PATH)\n"
            "_impl = _ilu.module_from_spec(_spec)\n"
            "_spec.loader.exec_module(_impl)\n"
            "_impl.install()\n\n"
            "ProbeWorkerExtension = _impl.ProbeWorkerExtension\n"
        )
        if PROBE_MODE == "sitecustomize":
            # Fallback injection: every interpreter on PYTHONPATH imports this.
            (self.shim_dir / "sitecustomize.py").write_text(
                "import os\n\n"
                'if os.environ.get("VLLM_ROUTED_EXPERTS_DUMP_DIR"):\n'
                "    import vllm_re_capture_probe  # noqa: F401\n"
            )
        return self.shim_dir

    def _server_env(self, spec: _ServerSpec) -> dict[str, str]:
        env = dict(spec.env)
        env["VLLM_ROUTED_EXPERTS_DUMP_DIR"] = str(self.dump_dir)
        if spec.probe:
            entries = [str(self.shim_dir)]
            entries += [
                entry
                for entry in os.environ.get("PYTHONPATH", "").split(os.pathsep)
                if entry
            ]
            env["PYTHONPATH"] = os.pathsep.join(entries)
        return env

    def _ensure(self, spec: _ServerSpec) -> None:
        if self._active == spec.name:
            return
        self.close()
        args = list(spec.args)
        if spec.probe and PROBE_MODE == "worker-extension":
            args += [
                "--worker-extension-cls",
                "vllm_re_capture_probe.ProbeWorkerExtension",
            ]
        print(f"[routed-experts-e2e] starting {spec.name}: {' '.join(args)}")
        server = self._stack.enter_context(
            RemoteOpenAIServer(
                spec.model,
                args,
                env_dict=self._server_env(spec),
                max_wait_seconds=STARTUP_TIMEOUT,
            )
        )
        self._client = server.get_client(timeout=REQUEST_TIMEOUT)
        self._active = spec.name
        if spec.probe:
            self._warm_up(spec)

    def _warm_up(self, spec: _ServerSpec) -> None:
        """Burn one request with no dump tag: profiling, JIT and graph capture
        (lazy or not) happen here, never inside a measured request."""
        self._client.completions.create(
            model=spec.model,
            prompt=list(_prompt(_WARMUP_SLOT, 1)),
            max_tokens=1,
            temperature=0.0,
            extra_body={"return_token_ids": True, "ignore_eos": True},
        )

    def generate(
        self,
        spec: _ServerSpec,
        tokens: tuple[int, ...],
        *,
        max_tokens: int,
        case: str,
        dump: bool = True,
        request_timeout: float | None = None,
    ) -> _Result:
        self._ensure(spec)
        self._requests += 1
        tag = f"{spec.name}-{self._requests}-{uuid.uuid4().hex[:8]}" if dump else None
        if tag is not None:
            self.tag_path.write_text(tag)
        extra: dict[str, Any] = {}
        if request_timeout is not None:
            extra["timeout"] = request_timeout
        try:
            response = self._client.completions.create(
                model=spec.model,
                prompt=list(tokens),
                max_tokens=max_tokens,
                temperature=0.0,
                extra_body={"return_token_ids": True, "ignore_eos": True},
                **extra,
            )
        finally:
            self.tag_path.unlink(missing_ok=True)
        choice = response.model_dump()["choices"][0]
        token_ids = choice.get("token_ids") or []
        records = _read_records(self.dump_dir, tag) if tag else []
        return _Result(
            tokens=tokens,
            prompt_tokens=int(response.usage.prompt_tokens),
            generated=len(token_ids),
            routed=_decode_routed_experts(choice.get("routed_experts"), case),
            records=records,
            case=case,
            tag=tag,
        )


@pytest.fixture(scope="module")
def harness() -> _Harness:
    model, _args, mega = _target_model()
    if mega and not _mega_moe_device_supported():
        pytest.skip(
            f"{model} would need a MEGA_MOE_BACKENDS backend for model-level "
            "token sharding, but all of them require SM100/SM103 "
            "(deep_gemm_mega_moe: get_deep_gemm_mega_moe_backend; FlashInfer "
            "MoE-EP: FlashInferMoeEpBackendSpec.supports_current_device) and "
            f"this device reports {_device_name()}. Re-run with "
            "VLLM_ROUTED_EXPERTS_E2E_EMULATE_SP=1 and "
            "VLLM_ROUTED_EXPERTS_E2E_DUMMY_ARCH=mixtral to exercise the DP=1 "
            "sharded layout through the probe's shard injection instead."
        )
    instance = _Harness(ARTIFACT_DIR)
    try:
        yield instance
    finally:
        instance.close()


def _device_name() -> str:
    try:
        import torch

        if not torch.cuda.is_available():
            return "no CUDA device"
        major, minor = torch.cuda.get_device_capability()
        return f"{torch.cuda.get_device_name()} (SM{major}{minor})"
    except Exception as exc:  # pragma: no cover - defensive
        return f"unknown device ({exc!r})"


def _decode_routed_experts(encoded: str | None, case: str) -> np.ndarray:
    assert encoded is not None, (
        f"{case}: the response carried no routed_experts. Check that "
        "--enable-return-routed-experts is active, that the model runner is V2 "
        'and that prefix caching is enabled ("AuxOutput Connector requires '
        'prefix caching").'
    )
    array = np.load(io.BytesIO(pybase64.b64decode(encoded)))
    assert array.ndim == 3, f"{case}: routed_experts must be [tokens, layers, topk]"
    return array.astype(np.int32)


def _read_records(dump_dir: Path, tag: str) -> list[_Record]:
    records: list[_Record] = []
    for path in sorted(dump_dir.glob(f"{tag}.*.bin")):
        data = path.read_bytes()
        offset = 0
        while offset < len(data):
            if data[offset : offset + 4] != _MAGIC:
                raise _HarnessError(f"{path.name}: corrupt probe dump at {offset}")
            (
                _magic,
                seq,
                tp_rank,
                dp_rank,
                layer_id,
                num_tokens,
                num_padding,
                n_rows,
                topk,
                gathers,
                buffer_rows,
            ) = _HEADER.unpack_from(data, offset)
            offset += _HEADER.size
            count = n_rows * topk
            ids = np.frombuffer(data, dtype=np.int32, count=count, offset=offset)
            offset += count * 4
            count = buffer_rows * topk
            rows = np.frombuffer(data, dtype=np.int32, count=count, offset=offset)
            offset += count * 4
            records.append(
                _Record(
                    seq=seq,
                    tp_rank=tp_rank,
                    dp_rank=dp_rank,
                    layer_id=layer_id,
                    num_tokens=num_tokens,
                    num_padding=num_padding,
                    gathers=gathers,
                    ids=ids.reshape(n_rows, topk).copy(),
                    buffer_rows=rows.reshape(buffer_rows, topk).copy(),
                )
            )
    return records


def _split_steps(records: list[_Record]) -> list[list[_Record]]:
    """Split one rank's records into forwards: the layer id restarts each time."""
    steps: list[list[_Record]] = []
    current: list[_Record] = []
    previous: int | None = None
    for record in records:
        if previous is not None and record.layer_id <= previous:
            steps.append(current)
            current = []
        current.append(record)
        previous = record.layer_id
    if current:
        steps.append(current)
    return steps


def _merge_step(index: int, group: dict[int, list[_Record]]) -> _Step:
    _, records = next(iter(group.items()))
    layer_ids = tuple(record.layer_id for record in records)
    step = _Step(
        index=index,
        num_tokens=records[0].num_tokens,
        num_padding=records[0].num_padding,
        layer_ids=layer_ids,
    )
    for rank, records in group.items():
        if tuple(record.layer_id for record in records) != layer_ids:
            raise _HarnessError(
                f"forward {index}: TP rank {rank} captured layers "
                f"{tuple(r.layer_id for r in records)}, another rank captured "
                f"{layer_ids}; the ranks did not run the same forward"
            )
        for record in records:
            if (record.num_tokens, record.num_padding) != (
                step.num_tokens,
                step.num_padding,
            ):
                raise _HarnessError(
                    f"forward {index}: TP rank {rank} reports token count "
                    f"({record.num_tokens}, padding {record.num_padding}), "
                    f"another rank reports ({step.num_tokens}, padding "
                    f"{step.num_padding})"
                )
            step.shards.setdefault(record.layer_id, {})[rank] = record.ids
            step.buffers.setdefault(record.layer_id, {})[rank] = record.buffer_rows
            step.gathers[(record.layer_id, rank)] = record.gathers
    return step


def _group_steps(records: list[_Record], tp_size: int) -> list[_Step]:
    if not records:
        raise _HarnessError(
            "no capture records were dumped. The probe did not run in the "
            "worker processes; see probe_notes.log in the artifact directory. "
            "Try VLLM_ROUTED_EXPERTS_E2E_PROBE_MODE=sitecustomize if "
            "--worker-extension-cls is not honoured on this stack."
        )
    by_rank: dict[int, list[_Record]] = {}
    for record in records:
        by_rank.setdefault(record.tp_rank, []).append(record)
    if any(rank < 0 for rank in by_rank):
        raise _HarnessError(
            "the probe could not resolve a TP rank for some captures, so the "
            "shards cannot be ordered; run with vLLM's distributed state "
            "initialized as usual and report this"
        )
    expected_ranks = set(range(tp_size))
    if set(by_rank) != expected_ranks:
        raise _HarnessError(
            f"capture records came from TP ranks {sorted(by_rank)}, expected "
            f"{sorted(expected_ranks)}"
        )
    per_rank = {
        rank: _split_steps(sorted(records, key=lambda record: record.seq))
        for rank, records in by_rank.items()
    }
    counts = {rank: len(steps) for rank, steps in per_rank.items()}
    if len(set(counts.values())) != 1:
        raise _HarnessError(
            f"TP ranks captured a different number of forwards: {counts}; a "
            "rank skipped or repeated a forward"
        )
    return [
        _merge_step(index, {rank: steps[index] for rank, steps in per_rank.items()})
        for index in range(next(iter(counts.values()), 0))
    ]


def _layer_reference(
    step: _Step, layer_id: int, tp_size: int
) -> tuple[np.ndarray, bool, bool]:
    """Rebuild one forward's routing rows for one layer from the raw shards.

    Returns:
        The unpadded batch's rows, whether the shards were TP-split, and
        whether any row had to be trimmed from the gathered batch.

    """
    sizes = step.shard_sizes(layer_id)
    ranks = sorted(sizes)
    padded = step.num_tokens
    real = step.num_real
    if padded <= 0 or real <= 0:
        raise _HarnessError(
            f"forward {step.index}: the probe saw no forward token count, so "
            "the shards cannot be trimmed to the batch. The capturer's context "
            "must expose num_tokens (or a batch descriptor with num_tokens)."
        )
    if all(size == padded for size in sizes.values()):
        # No token sharding: every rank holds the whole batch.
        return step.shards[layer_id][ranks[0]][:real], False, padded > real
    shard_n = (padded + tp_size - 1) // tp_size
    unexpected = {rank: size for rank, size in sizes.items() if size != shard_n}
    if unexpected:
        raise _HarnessError(
            f"forward {step.index} layer {layer_id}: capture was fed "
            f"{sizes} rows per TP rank for a {padded}-token forward at "
            f"tp_size={tp_size}; expected {padded} on every rank (no sharding) "
            f"or {shard_n} (TP shard)"
        )
    gathered = np.concatenate([step.shards[layer_id][rank] for rank in ranks], axis=0)
    return gathered[:real], True, gathered.shape[0] > real


def _build_trace(records: list[_Record], tp_size: int) -> _Trace:
    steps = _group_steps(records, tp_size)
    publish_rank = min(steps[0].shards[steps[0].layer_ids[0]])
    reference: dict[int, list[np.ndarray]] = {}
    published: dict[int, list[np.ndarray]] = {}
    bounds: list[tuple[int, int]] = []
    sharded: list[bool] = []
    trimmed: list[bool] = []
    gathers: list[int] = []
    cursor = 0
    for step in steps:
        step_sharded = False
        step_trimmed = False
        for layer_id in step.layer_ids:
            rows, was_sharded, was_trimmed = _layer_reference(step, layer_id, tp_size)
            step_sharded = step_sharded or was_sharded
            step_trimmed = step_trimmed or was_trimmed
            reference.setdefault(layer_id, []).append(rows)
            buffer = step.buffers.get(layer_id, {}).get(publish_rank)
            if buffer is not None:
                published.setdefault(layer_id, []).append(buffer)
        bounds.append((cursor, cursor + step.num_real))
        cursor += step.num_real
        sharded.append(step_sharded)
        trimmed.append(step_trimmed)
        gathers.append(max(step.gathers.values(), default=0))
    return _Trace(
        steps=steps,
        bounds=bounds,
        sharded=sharded,
        trimmed=trimmed,
        gathers=gathers,
        reference={
            layer: np.concatenate(chunks, axis=0) for layer, chunks in reference.items()
        },
        published={
            layer: np.concatenate(chunks, axis=0)
            for layer, chunks in published.items()
            if len(chunks) == len(steps)
        },
        publish_rank=publish_rank,
    )


def _summary(label: str, result: _Result, trace: _Trace) -> str:
    per_step = ", ".join(
        f"#{step.index}: {step.num_real} tokens / "
        f"{step.shard_sizes(step.first_layer)} rows per rank"
        for step in trace.steps
    )
    return (
        f"[{label}] returned rows={result.routed.shape[0]} "
        f"(prompt {result.prompt_tokens} + generated {result.generated} - 1), "
        f"reference rows={trace.total_rows}, captured layers={len(trace.layers)}, "
        f"publishing rank={trace.publish_rank}, "
        f"sharded={trace.sharded}, sp-padding trimmed={trace.trimmed}, "
        f"capture gathers per forward={trace.gathers}; steps: {per_step}"
    )


def _mismatch_examples(
    got: np.ndarray, expected: np.ndarray, start: int, layer_id: int, limit: int = 5
) -> list[str]:
    rows = np.flatnonzero((got != expected).any(axis=1))
    examples = []
    for offset in rows[:limit]:
        examples.append(
            f"  token {start + int(offset)} layer {layer_id}: got "
            f"{np.array2string(got[offset], max_line_width=200)} expected "
            f"{np.array2string(expected[offset], max_line_width=200)}"
        )
    if rows.size > limit:
        examples.append(f"  ... and {rows.size - limit} more rows")
    return examples


def _check_capture(label: str, result: _Result, tp_size: int) -> _Trace:
    """Assert the returned routes equal the raw-shard reference, end to end."""
    trace = _build_trace(result.records, tp_size)
    print(_summary(label, result, trace))
    _warn_if_routing_is_degenerate(label, trace)
    problems: list[str] = []
    if result.routed.shape[0] != result.expected_rows:
        problems.append(
            f"{label}: response returned {result.routed.shape[0]} routed-expert "
            f"rows, expected prompt + generated - 1 = {result.expected_rows}"
        )
    if trace.total_rows != result.expected_rows:
        problems.append(
            f"{label}: the capture covered {trace.total_rows} rows but the "
            f"response returned {result.expected_rows}. Either a forward was "
            "recorded that is not part of this request (replay or warmup step) "
            "or a forward is missing"
        )
    if not problems:
        for layer_id in trace.layers:
            expected_layer = trace.reference[layer_id]
            for step, (start, end) in zip(trace.steps, trace.bounds):
                got = result.routed[start:end, layer_id]
                expected = expected_layer[start:end]
                if got.shape != expected.shape:
                    problems.append(
                        f"{label}: forward {step.index} layer {layer_id}: "
                        f"response rows {got.shape} != reference {expected.shape}"
                    )
                    continue
                if not np.array_equal(got, expected):
                    problems.extend(
                        [
                            f"{label}: forward {step.index} layer {layer_id}: "
                            f"{int((got != expected).any(axis=1).sum())} of "
                            f"{got.shape[0]} rows differ from the raw-shard "
                            f"reference",
                            *_mismatch_examples(got, expected, start, layer_id),
                        ]
                    )
            published = trace.published.get(layer_id)
            if published is None:
                continue
            expected_layer = trace.reference[layer_id]
            if published.shape != expected_layer.shape:
                problems.append(
                    f"{label}: layer {layer_id}: TP rank {trace.publish_rank}'s "
                    f"capture buffer holds {published.shape[0]} rows for a "
                    f"{expected_layer.shape[0]}-row batch; the capture wrote a "
                    "partial batch (capture, not the auxiliary-output path, is "
                    "wrong)"
                )
            elif not np.array_equal(published, expected_layer):
                problems.extend(
                    [
                        f"{label}: layer {layer_id}: TP rank "
                        f"{trace.publish_rank}'s capture buffer does not hold "
                        "the whole batch (capture, not the auxiliary-output "
                        "path, is wrong)",
                        *_mismatch_examples(
                            published, expected_layer, 0, layer_id, limit=3
                        ),
                    ]
                )
    if problems:
        detail = "\n".join(problems)
        pytest.fail(f"{detail}\nprobe notes:\n{_probe_notes_text()}")
    return trace


def _probe_notes_text() -> str:
    notes = ARTIFACT_DIR / "dumps" / "probe_notes.log"
    return notes.read_text().strip() if notes.exists() else "(no probe notes)"


def _assert_injection(tag: str | None, needle: str, label: str) -> None:
    """Assert the probe rewrote this request's capture input as intended."""
    assert tag, f"{label}: no dump tag for this request"
    matches = "\n".join(
        line for line in _probe_notes_text().splitlines() if tag in line
    )
    assert needle in matches, (
        f"{label}: the probe did not record {needle!r} for tag {tag}; probe "
        f"notes for this request:\n{matches or '(none)'}"
    )


def _assert_dp1_sharded(label: str, result: _Result, trace: _Trace) -> None:
    """Assert the DP=1 sharded capture branch was really the one exercised."""
    if EMULATE_SP:
        _assert_injection(result.tag, "injected shard layout", label)
    else:
        step = trace.steps[0]
        assert trace.sharded[0], (
            f"{label}: rank rows {step.shard_sizes(step.first_layer)} show that "
            "no rank captured a TP shard, so this configuration does not "
            "exercise the DP=1 sharded-capture path. Check "
            "--enable-expert-parallel, the --moe-backend, and the tensor "
            "parallel size"
        )
    assert trace.max_gathers == 1, (
        f"{label}: the capture must reconstruct the batch with one TP gather "
        f"per forward, saw {trace.gathers}"
    )


def _mega_moe_device_supported() -> bool:
    """Whether any MEGA_MOE_BACKENDS member can bind on this device."""
    try:
        import torch

        if not torch.cuda.is_available():
            return False
        return torch.cuda.get_device_capability() in ((10, 0), (10, 3))
    except Exception:
        return False


def _warn_if_routing_is_degenerate(label: str, trace: _Trace) -> None:
    """Flag weights whose routing cannot show a stale row.

    Dummy ("dummy" load format) or degenerate weights can route every token to
    the same experts; then a stale row is indistinguishable from a fresh one by
    value. The capture-buffer shape check still catches a partial write, but the
    value comparison stops being evidence, so the run has to say so.
    """
    flat = [
        layer_id
        for layer_id in trace.layers
        if len(np.unique(trace.reference[layer_id], axis=0)) <= 1
    ]
    if flat:
        shown = ", ".join(str(layer_id) for layer_id in flat[:8])
        more = f" (+{len(flat) - 8} more)" if len(flat) > 8 else ""
        print(
            f"[{label}] WARNING: every token routed to the same experts in "
            f"{len(flat)} of {len(trace.layers)} captured layers [{shown}{more}]; "
            "stale rows cannot be told apart by value with these weights, so "
            "only the capture-buffer shape check discriminates here"
        )


def _compare_rows(
    label: str, got: np.ndarray, expected: np.ndarray, note: str = ""
) -> None:
    assert got.shape == expected.shape, (
        f"{label}: routed_experts shape {got.shape} != reference {expected.shape}{note}"
    )
    if np.array_equal(got, expected):
        return
    differing = (got != expected).any(axis=2)
    examples = []
    for row, layer in zip(*np.nonzero(differing)):  # noqa: B905 - numpy arrays
        examples.append(
            f"  token {int(row)} layer {int(layer)}: got "
            f"{np.array2string(got[row, layer], max_line_width=200)} expected "
            f"{np.array2string(expected[row, layer], max_line_width=200)}"
        )
        if len(examples) >= 5:
            break
    pytest.fail(
        f"{label}: {int(differing.sum())} (token, layer) pairs differ from the "
        f"reference{note}\n" + "\n".join(examples)
    )


def test_six_token_prefill_captures_every_tp_shard(harness: _Harness):
    """Case 1: a 6-token prefill at TP=2 must return all six rows."""
    result = harness.generate(
        _main_spec(enforce_eager=True), _prompt(0, 6), max_tokens=1, case="six-token"
    )
    trace = _check_capture("six-token prefill", result, TP_SIZE)
    _assert_dp1_sharded("six-token prefill", result, trace)


def test_three_token_prefill_drops_only_sp_padding(harness: _Harness):
    """Case 2: ceil(3 / 2) = 2 rows per rank, and the 4th row is padding."""
    result = harness.generate(
        _main_spec(enforce_eager=True), _prompt(1, 3), max_tokens=1, case="three-token"
    )
    trace = _check_capture("three-token prefill", result, TP_SIZE)
    _assert_dp1_sharded("three-token prefill", result, trace)
    step = trace.steps[0]
    if EMULATE_SP:
        # The probe slices the batch into ceil(3 / 2) = 2 rows per rank and pads
        # rank 1 to that width, exactly like sp_reduce_scatter's ceil-div pad;
        # the capture has to trim back to 3 rows, so the returned data ends at
        # token 2 and never includes the padded row.
        assert step.num_real == 3 and result.routed.shape[0] == 3
    else:
        assert trace.trimmed[0], (
            "three-token prefill: the SP-padded row must be trimmed, saw shard "
            f"rows {step.shard_sizes(step.first_layer)} for {step.num_real} tokens"
        )
        assert result.routed.shape[0] == 3


def test_second_uncached_batch_has_no_stale_rows(harness: _Harness):
    """Case 3: a second, different 6-token batch may not inherit old ids."""
    first = harness.generate(
        _main_spec(enforce_eager=True), _prompt(2, 6), max_tokens=1, case="batch-1"
    )
    _assert_dp1_sharded(
        "first six-token batch",
        first,
        _check_capture("first six-token batch", first, TP_SIZE),
    )
    second = harness.generate(
        _main_spec(enforce_eager=True), _prompt(3, 6), max_tokens=1, case="batch-2"
    )
    _assert_dp1_sharded(
        "second six-token batch",
        second,
        _check_capture("second six-token batch", second, TP_SIZE),
    )
    assert not np.array_equal(first.routed, second.routed), (
        "the two batches are different prompts but returned identical routes; "
        "the dump tag or the prompt slots are not distinct"
    )


def test_single_token_steps_capture_without_gather(harness: _Harness):
    """Case 4: 1-token prefill and decodes need no capture-only gather."""
    result = harness.generate(
        _main_spec(enforce_eager=True), _prompt(4, 1), max_tokens=3, case="single-token"
    )
    trace = _check_capture("single-token steps", result, TP_SIZE)
    if EMULATE_SP:
        _assert_injection(result.tag, "injected shard layout", "single-token steps")
    assert len(trace.steps) == 3, (
        f"expected one prefill and two decode forwards, saw {len(trace.steps)}"
    )
    assert all(step.num_real == 1 for step in trace.steps)
    assert trace.max_gathers == 0, (
        "a 1-token forward must not run a capture-only gather: TP0 owns the "
        f"real row, saw gathers {trace.gathers}"
    )


@pytest.mark.skipif(
    not RUN_GRAPH, reason="set VLLM_ROUTED_EXPERTS_E2E_GRAPH=1 to run case 5"
)
def test_cuda_graph_matches_eager_reference(harness: _Harness):
    """Case 5: graph execution returns the eager routes and does not hang."""
    references: dict[tuple[int, ...], np.ndarray] = {}
    for slot, count in ((5, 6), (6, 3)):
        tokens = _prompt(slot, count)
        eager = harness.generate(
            _main_spec(enforce_eager=True), tokens, max_tokens=1, case="graph-ref"
        )
        _check_capture(f"eager reference ({count} tokens)", eager, TP_SIZE)
        references[tokens] = eager.routed
    for tokens, expected in references.items():
        graph = harness.generate(
            _main_spec(enforce_eager=False),
            tokens,
            max_tokens=1,
            case="cuda-graph",
        )
        print(
            f"[cuda-graph] prompt {len(tokens)} tokens: dumped "
            f"{len(graph.records)} capture records during graph execution"
        )
        if graph.records:
            # Graph replay need not execute the Python capture callback, so the
            # records may describe capture-time dummy batches only. When they do
            # form a usable reference, hold the replay output to it as well.
            try:
                _build_trace(graph.records, TP_SIZE)
            except _HarnessError as exc:
                print(f"[cuda-graph] capture records are not a usable reference: {exc}")
            else:
                _check_capture(
                    f"cuda-graph ({len(tokens)} tokens), reference from its own "
                    "capture records",
                    graph,
                    TP_SIZE,
                )
        _compare_rows(
            f"cuda-graph ({len(tokens)} tokens) vs eager",
            graph.routed,
            expected,
            note=(
                "; re-run the eager reference at the same padded forward shape "
                "if the padded batch perturbs the router's floating point "
                "reduction (see the run's printed shard sizes)"
            ),
        )


@pytest.mark.skipif(
    not RUN_CONTROL or not (CONTROL_MODEL or DUMMY),
    reason=(
        "set VLLM_ROUTED_EXPERTS_E2E_CONTROL_MODEL, or enable the generated "
        "dummy model, to run case 6"
    ),
)
def test_ordinary_moe_path_captures_full_batch_without_gather(harness: _Harness):
    """Case 6: TP+EP without model-level token SP stays a plain full batch."""
    tokens = _prompt(7, 6)
    sharded = harness.generate(
        _control_spec(tp_size=2, probe=True), tokens, max_tokens=1, case="control-tp2"
    )
    trace = _check_capture("ordinary MoE path at TP=2", sharded, 2)
    step = trace.steps[0]
    assert not any(trace.sharded), (
        f"{CONTROL_MODEL} must not shard tokens at the model level, saw shard "
        f"rows {step.shard_sizes(step.first_layer)} for {step.num_real} tokens"
    )
    assert trace.max_gathers == 0, (
        "the complete-batch path must not perform a capture-only gather, saw "
        f"{trace.gathers}"
    )


@pytest.mark.skipif(
    not (RUN_CONTROL and RUN_TP1_ORACLE) or not (CONTROL_MODEL or DUMMY),
    reason=(
        "set VLLM_ROUTED_EXPERTS_E2E_TP1_ORACLE=1 (and enable case 6) to run "
        "the single-GPU route comparison"
    ),
)
def test_ordinary_moe_path_matches_single_gpu_routes(harness: _Harness):
    """Case 6b: independent oracle - TP=2+EP routing equals a TP=1 run.

    Kept separate from case 6 because this oracle is only exact while no
    expert pair is tied in the router scores: a TP=1 and a TP=2 run reduce the
    attention output in a different order, so a tie can flip an id.
    """
    tokens = _prompt(7, 6)
    sharded = harness.generate(
        _control_spec(tp_size=2, probe=True), tokens, max_tokens=1, case="control-tp2"
    )
    single = harness.generate(
        _control_spec(tp_size=1, probe=False), tokens, max_tokens=1, case="control-tp1"
    )
    _compare_rows(
        "ordinary MoE path, TP=2 vs TP=1",
        sharded.routed,
        single.routed,
        note="; the ordinary MoE path must not change routing with TP size",
    )


@pytest.mark.skipif(
    not RUN_SYNC_CHECK,
    reason="set VLLM_ROUTED_EXPERTS_E2E_SYNC_CHECK=1 to run the sync check",
)
def test_capture_gather_adds_no_host_sync(harness: _Harness):
    """The added TP gather must not add a GPU-to-CPU read (no probe here)."""
    result = harness.generate(
        _sync_check_spec(), _prompt(8, 6), max_tokens=2, case="sync-check"
    )
    assert result.routed.shape[0] == result.expected_rows, (
        "the sync-check run returned "
        f"{result.routed.shape[0]} rows, expected {result.expected_rows}"
    )
    print(
        "[sync-check] a sharded 6-token prefill plus one decode completed under "
        f"VLLM_GPU_SYNC_CHECK with {result.routed.shape[0]} routed-expert rows"
    )


def test_complete_padded_batch_is_trimmed_not_regathered(harness: _Harness):
    """DP=1 branch: n == ceil(T / tp) * tp, an already gathered padded batch.

    The probe gathers each rank's shard before calling the capture, so the
    capture is fed the complete, SP-padded batch instead of a shard. The fix
    must trim it to the forward token count and must not gather a second time;
    the pre-fix code would write those padded rows into the buffer unchanged
    (and, with the gather-count check, a re-gather is visible too).
    """
    spec = _injected_spec("already-gathered")
    for slot, count in ((9, 3), (10, 5)):
        tokens = _prompt(slot, count)
        result = harness.generate(
            spec, tokens, max_tokens=1, case=f"already-gathered-{count}"
        )
        trace = _check_capture(
            f"already-gathered batch ({count} tokens)", result, TP_SIZE
        )
        step = trace.steps[0]
        assert step.num_real == count and result.routed.shape[0] == count
        if EMULATE_SP:
            _assert_injection(
                result.tag, "injected already-gathered batch", f"{count}-token batch"
            )
        else:
            assert (
                step.shard_sizes(step.first_layer)[0]
                == (count + TP_SIZE - 1) // TP_SIZE
            )
            assert trace.trimmed[0], (
                f"{count}-token forward: the padded batch must be trimmed, shard "
                f"rows {step.shard_sizes(step.first_layer)} for {step.num_real} "
                "tokens"
            )
        assert trace.max_gathers == 0, (
            "an already gathered batch must be trimmed without a second gather, "
            f"saw {trace.gathers} gather calls per forward"
        )


@pytest.mark.skipif(
    not RUN_FAIL_CLOSED,
    reason="set VLLM_ROUTED_EXPERTS_E2E_FAIL_CLOSED=0 to skip the fail-closed run",
)
def test_unexpected_capture_batch_fails_closed(harness: _Harness):
    """DP=1 branch: an unrecognized row count must raise, not write stale rows.

    Runs last on purpose: the worker's assertion kills the engine core, so the
    server is not reused afterwards.
    """
    spec = _injected_spec("unexpected")
    with pytest.raises(Exception) as caught:
        harness.generate(
            spec,
            _prompt(11, 6),
            max_tokens=1,
            case="unexpected",
            request_timeout=REQUEST_TIMEOUT,
        )
    print(f"[fail-closed] the request failed as required: {caught.value!r}")
    notes = _probe_notes_text()
    assert "unexpected topk_ids batch" in notes, (
        f"the capture must fail closed with its shape assertion; probe notes: {notes}"
    )

Vs base / validation results are posted in a separate comment.

[Ronald1995]

Model end-to-end reproducer for the DP=1 sharded capture (2/2): the validation probe

RoutedExpertsCapturer.capture receives only this rank's rows, so no single rank holds the whole batch: the reference has to be the raw per-rank shards. This probe records, for every capture call and every TP rank, the tensor the capture path is fed (read before the capture logic can gather it), the rows the call leaves in that rank's buffer, whether the call performed a TP collective, and the forward token count before model-level SP (ForwardContext.num_tokens, falling back to batch_descriptor.num_tokens so it also works on the base revision).

It changes nothing on disk: the test puts a generated shim on PYTHONPATH and starts the server with --worker-extension-cls vllm_re_capture_probe.ProbeWorkerExtension, so vLLM resolves the class - importing this module - inside every worker before the model is built. Dumping is gated on a current_tag file that the test rewrites per request, so warmup, profiling and graph-capture forwards are never recorded.

It can also force a capture-input layout, which is how the DP=1 branches can be exercised on a host where no MEGA_MOE_BACKENDS backend binds (get_deep_gemm_mega_moe_backend requires SM100; FlashInferMoeEpBackendSpec.supports_current_device accepts only (10,0)/(10,3)):

• shard - slice the batch into ceil(T/tp)-row TP shards, padded the way sp_reduce_scatter pads, so the real capture code and a real TP all-gather run on the layout the bug is about;
• already-gathered - hand over the complete SP-padded batch;
• unexpected - hand over a row count that matches no branch, to check that the capture fails closed.

# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Validation-only instrumentation for the routed-experts capture reproducer.

The probe records, for every ``RoutedExpertsCapturer.capture`` call and every
TP rank: the tensor the capture path is fed (this rank's shard, read *before*
the capture logic can gather it), the rows the call leaves in this rank's
capture buffer, whether the call performed a TP collective, and the forward
token count before model-level sequence parallel (SP).

That is the only way to get a per-rank reference for a capture that is
supposed to write a whole batch from TP shards: a single batch's routing is
split across ranks, so no single rank holds the answer, and comparing against
a differently sharded run (TP=1) is not exact because the router's floating
point reduction order changes with the shard shape.

This module changes nothing on disk. The reproducer puts a generated shim
directory on ``PYTHONPATH`` and starts the server with
``--worker-extension-cls vllm_re_capture_probe.ProbeWorkerExtension``, so vLLM
resolves the extension class - importing this module - inside every worker
process before the model is built. See
``tests/models/deepseek_v41_routed_experts_capture_e2e.py``.

Dumping is off unless ``VLLM_ROUTED_EXPERTS_DUMP_DIR`` is set *and* that
directory holds a ``current_tag`` file. The reproducer rewrites the tag before
each request, so warmup, profiling and CUDA-graph capture steps are never
recorded.

Record layout, appended as
``<dump dir>/<tag>.rank<tp_rank>.<pid>.bin``::

    <4s magic><q seq><i tp_rank><i dp_rank><i layer_id><i num_tokens>
    <i num_padding><i n_rows><i topk><i gathers><i buffer_rows>
    <int32 ids[n_rows, topk]><int32 buffer[buffer_rows, topk]>

``num_tokens`` is the forward token count before model-level SP, including
batch padding, or ``-1`` when the producer does not report one; ``num_padding``
counts trailing padding rows within it. ``ids`` is what the capture path was
fed on this rank and ``buffer`` is ``device_buffer[:buffer_rows, layer_id]``
read right after the call. Files are appended and flushed per record, so a
crashed run still leaves usable dumps.
"""

from __future__ import annotations

import os
import struct
import time
from pathlib import Path
from typing import Any

ENV_DUMP_DIR = "VLLM_ROUTED_EXPERTS_DUMP_DIR"
ENV_INJECT = "VLLM_ROUTED_EXPERTS_E2E_INJECT"
TAG_FILENAME = "current_tag"
NOTES_FILENAME = "probe_notes.log"

# Layouts the capture path can be fed, used to force a branch E2E:
#   "shard": slice the batch into TP shards (with the ceil-div padding
#     sp_reduce_scatter applies), so a model that would hand over the whole
#     batch exercises the DP=1 sharded branch. This is how the sharded layout
#     can be reached on GPUs where no MegaMoE backend can bind (all three are
#     SM100/SM103 gated), with the real capture code and a real TP gather.
#   "already-gathered": gather the shard first, so the capture sees the
#     complete SP-padded batch and must trim it without gathering again;
#   "unexpected": feed a row count that matches no branch, so the capture
#     must fail closed instead of writing a partial batch.
INJECT_MODES = ("shard", "already-gathered", "unexpected")

_MAGIC = b"RXD1"
_HEADER = struct.Struct("<4sqiiiiiiiii")
_INSTALLED_ATTR = "_routed_experts_capture_probe_installed"


class ProbeWorkerExtension:
    """Worker extension whose import installs the probe.

    ``--worker-extension-cls`` makes vLLM resolve this class inside every
    worker process before the model is built, which imports this module. The
    body is intentionally empty: only the import has an effect.
    """


def install() -> bool:
    """Wrap ``RoutedExpertsCapturer.capture`` with the dumping recorder.

    Returns:
        True if the wrapper is installed, False if the probe is disabled or
        vLLM's capture module could not be patched.

    """
    dump_dir = os.environ.get(ENV_DUMP_DIR)
    if not dump_dir:
        return False
    directory = Path(dump_dir)
    try:
        from vllm.model_executor.layers.fused_moe import routed_experts_capturer
    except Exception as exc:  # pragma: no cover - defensive
        _note(directory, f"import failed: {exc!r}")
        return False
    capturer_cls = getattr(routed_experts_capturer, "RoutedExpertsCapturer", None)
    if capturer_cls is None:  # pragma: no cover - defensive
        _note(directory, "RoutedExpertsCapturer not found")
        return False
    if getattr(capturer_cls, _INSTALLED_ATTR, False):
        return True
    setattr(capturer_cls, _INSTALLED_ATTR, True)

    recorder = _Recorder(directory, routed_experts_capturer)
    original = capturer_cls.capture

    def capture(self: Any, layer_id: int, topk_ids: Any) -> None:
        recorder.capture(original, self, layer_id, topk_ids)

    capturer_cls.capture = capture
    _note(directory, f"installed pid={os.getpid()}")
    return True


class _Recorder:
    """Per-process recorder; one append-only file per (tag, TP rank)."""

    def __init__(self, dump_dir: Path, module: Any) -> None:
        self._dir = dump_dir
        self._module = module
        self._seq = 0
        self._tp_rank: int | None = None
        # Error notes are capped so a broken dump cannot flood the log; layout
        # notes are emitted once per (tag, TP rank), so the reproducer can
        # always confirm which layout a request was fed.
        self._error_notes = 0
        self._layout_noted: set[tuple[str, int]] = set()

    def capture(
        self, original: Any, capturer: Any, layer_id: int, topk_ids: Any
    ) -> None:
        tag = _read_tag(self._dir)
        if tag is None:
            original(capturer, layer_id, topk_ids)
            return

        module = self._module
        ids = context = None
        try:
            ids = _to_int32(topk_ids)
            ctx = module.get_forward_context()
            num_tokens = _forward_token_count(ctx)
            num_padding = _padding_rows(ctx, num_tokens)
            context = (num_tokens, num_padding, self._rank(capturer))
        except Exception as exc:
            self._note_error(f"capture metadata failed: {exc!r}")

        if context is not None:
            try:
                topk_ids = self._inject(tag, capturer, topk_ids, context[0])
            except Exception as exc:
                self._note_error(f"injection failed: {exc!r}")

        gathers = 0
        real_get_tp_group = module.get_tp_group

        def counting_get_tp_group() -> Any:
            nonlocal gathers
            gathers += 1
            return real_get_tp_group()

        module.get_tp_group = counting_get_tp_group
        try:
            original(capturer, layer_id, topk_ids)
        except BaseException as exc:
            self._note_error(f"capture raised {type(exc).__name__}: {exc}")
            raise
        finally:
            module.get_tp_group = real_get_tp_group

        if ids is None or context is None:
            return
        try:
            num_tokens, num_padding, tp_rank = context
            rows = _buffer_rows(capturer, layer_id, num_tokens, num_padding)
            self._write(
                tag,
                tp_rank,
                capturer,
                layer_id,
                num_tokens,
                num_padding,
                gathers,
                ids,
                rows,
            )
        except Exception as exc:
            self._note_error(f"dump failed: {exc!r}")

    def _inject(self, tag: str, capturer: Any, topk_ids: Any, num_tokens: int) -> Any:
        """Rewrite the capture input to force a layout branch (test only).

        The dump always keeps the rank's original tensor, so the reference the
        reproducer rebuilds is unaffected by the injection.
        """
        mode = os.environ.get(ENV_INJECT)
        if not mode or num_tokens <= 0:
            return topk_ids
        tp_size = max(int(getattr(capturer, "tp_size", 1)), 1)
        shard_n = (num_tokens + tp_size - 1) // tp_size
        rows = int(topk_ids.shape[0])
        rank = max(self._rank(capturer), 0)
        if mode == "shard":
            if rows != num_tokens:
                # Already a shard (the real model-level SP path): leave it.
                self._note_layout(
                    tag,
                    rank,
                    f"{tag}: shard injection skipped, capture already has "
                    f"{rows} rows for a {num_tokens}-token forward",
                )
                return topk_ids
            start = min(rank * shard_n, max(rows - 1, 0))
            local = topk_ids[start : start + shard_n]
            if local.shape[0] < shard_n:
                # sp_reduce_scatter pads with ceil-div, so every rank holds
                # exactly shard_n rows; the extra ones are trimmed downstream.
                pad = topk_ids[: shard_n - local.shape[0]]
                local = _concat(local, pad)
            self._note_layout(
                tag,
                rank,
                f"{tag}: injected shard layout: {rows} rows -> {local.shape[0]} "
                f"rows on tp_rank {rank} for a {num_tokens}-token forward",
            )
            return local.contiguous()
        if mode == "already-gathered":
            if rows == num_tokens and tp_size > 1:
                # Unsharded input (emulated layout): shard it the way the model
                # would, then gather back into one padded batch.
                start = min(rank * shard_n, max(rows - 1, 0))
                local = topk_ids[start : start + shard_n]
                if local.shape[0] < shard_n:
                    pad = topk_ids[: shard_n - local.shape[0]]
                    local = _concat(local, pad)
            else:
                local = topk_ids
            gathered = self._module.get_tp_group().all_gather(local.contiguous(), dim=0)
            self._note_layout(
                tag,
                rank,
                f"{tag}: injected already-gathered batch: {rows} rows -> "
                f"{gathered.shape[0]} rows for a {num_tokens}-token forward",
            )
            return gathered
        if mode == "unexpected":
            # A row count that matches none of the capturer's layouts, equal on
            # every rank so the failure is not a rank disagreement.
            forbidden = {num_tokens, shard_n, shard_n * tp_size}
            count = 1
            while count in forbidden:
                count += 1
            if rows >= count:
                injected = topk_ids[:count]
            else:
                repeats = (count + rows - 1) // rows
                injected = topk_ids.repeat(repeats, 1)[:count]
            self._note_layout(
                tag,
                rank,
                f"{tag}: injected unexpected batch: {rows} rows -> {count} rows "
                f"for a {num_tokens}-token forward at tp_size={tp_size}",
            )
            return injected.contiguous()
        raise ValueError(f"unknown {ENV_INJECT}={mode!r}")

    def _rank(self, capturer: Any) -> int:
        if self._tp_rank is None:
            try:
                from vllm.distributed.parallel_state import (
                    get_tensor_model_parallel_rank,
                )

                self._tp_rank = int(get_tensor_model_parallel_rank())
            except Exception:
                self._tp_rank = None
        if self._tp_rank is not None and self._tp_rank >= 0:
            return self._tp_rank
        # Encode the DP rank so a rank-less run still writes distinct files.
        return -1 - int(getattr(capturer, "dp_rank", 0))

    def _write(
        self,
        tag: str,
        tp_rank: int,
        capturer: Any,
        layer_id: int,
        num_tokens: int,
        num_padding: int,
        gathers: int,
        ids: Any,
        rows: Any,
    ) -> None:
        self._seq += 1
        path = self._dir / f"{tag}.rank{tp_rank}.{os.getpid()}.bin"
        with open(path, "ab") as handle:
            handle.write(
                _HEADER.pack(
                    _MAGIC,
                    self._seq,
                    tp_rank,
                    int(getattr(capturer, "dp_rank", 0)),
                    layer_id,
                    num_tokens,
                    num_padding,
                    int(ids.shape[0]),
                    int(ids.shape[1]),
                    gathers,
                    int(rows.shape[0]),
                )
            )
            handle.write(ids.tobytes())
            handle.write(rows.tobytes())

    def _note_layout(self, tag: str, tp_rank: int, message: str) -> None:
        """One layout note per (tag, TP rank): enough for the reproducer to
        confirm which layout a request was fed, without one line per layer."""
        key = (tag, tp_rank)
        if key in self._layout_noted:
            return
        self._layout_noted.add(key)
        _note(self._dir, message)

    def _note_error(self, message: str) -> None:
        # Keep a broken dump from flooding the log with one line per layer.
        if self._error_notes < 8:
            self._error_notes += 1
            _note(self._dir, message)


def _read_tag(dump_dir: Path) -> str | None:
    try:
        tag = (dump_dir / TAG_FILENAME).read_text().strip()
    except OSError:
        return None
    return tag or None


def _forward_token_count(ctx: Any) -> int:
    """Forward token count before model-level SP, including batch padding."""
    num_tokens = getattr(ctx, "num_tokens", None)
    if num_tokens is None:
        descriptor = getattr(ctx, "batch_descriptor", None)
        num_tokens = getattr(descriptor, "num_tokens", None)
    return -1 if num_tokens is None else int(num_tokens)


def _padding_rows(ctx: Any, num_tokens: int) -> int:
    is_padding = getattr(ctx, "is_padding", None)
    if is_padding is None or num_tokens < 0:
        return 0
    try:
        return int(is_padding[:num_tokens].sum().item())
    except Exception:  # pragma: no cover - defensive
        return 0


def _buffer_rows(
    capturer: Any, layer_id: int, num_tokens: int, num_padding: int
) -> Any:
    import numpy as np
    import torch

    buffer = capturer.device_buffer
    if num_tokens < 0 or layer_id >= buffer.shape[1]:
        return np.zeros((0, buffer.shape[2]), dtype=np.int32)
    wanted = max(num_tokens - max(num_padding, 0), 0)
    rows = min(wanted, buffer.shape[0])
    return buffer[:rows, layer_id].detach().to("cpu", torch.int32).numpy()


def _to_int32(tensor: Any) -> Any:
    import numpy as np
    import torch

    return np.ascontiguousarray(tensor.detach().to("cpu", torch.int32).numpy())


def _concat(first: Any, second: Any) -> Any:
    import torch

    return torch.cat((first, second), dim=0)


def _note(dump_dir: Path, message: str) -> None:
    try:
        dump_dir.mkdir(parents=True, exist_ok=True)
        stamp = time.strftime("%Y-%m-%dT%H:%M:%S")
        with open(dump_dir / NOTES_FILENAME, "a") as handle:
            handle.write(f"{stamp} {message}\n")
    except OSError:  # pragma: no cover - defensive
        pass

[Ronald1995]

Model end-to-end reproducer (3/3): base vs fix results, and a hardware gate you should know about

Headline

• Base revision 206e45d470: 4 failed, 4 passed, 2 skipped -?the three required sharded-capture cases fail, and so does the fail-closed guard that does not exist yet.
• PR head 446dc9aa9b: 8 passed, 2 skipped -?every returned token/layer/expert id equals the raw-shard reference, and the fail-closed guard fires.
• But a real DL/DSV4/V4.1/Kimi-K3 MegaMoE forward cannot run on this cluster at all: all three MEGA_MOE_BACKENDS members are hard SM100/SM103 gates, and this cluster is H20 (SM90). Detail and evidence below.

The hardware gate (affects what "required model E2E" can mean)

backend	gate
deep_gemm_mega_moe	get_deep_gemm_mega_moe_backend raises NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.") unless is_device_capability_family(100) (vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py)
flashinfer_moe_ep_cutedsl, flashinfer_moe_ep_mega_deep_gemm	FlashInferMoeEpBackendSpec.supports_current_device() accepts only (10, 0) / (10, 3) (vllm/model_executor/layers/fused_moe/flashinfer_moe_ep.py)

_use_sequence_parallel is a pure config predicate, so it is True on SM90 with --moe-backend deep_gemm_mega_moe, but startup then dies in the post-load finalize chain. Net effect: the model-level token sharding this PR fixes can only be produced by a real MegaMoE model on Blackwell (B200/B300). (--tensor-parallel-size >1 + --data-parallel-size >1 also turns SP on for other archs, but then dp_metadata exists and the DP branch -?not this fix -?handles it.)

So on H20 I ran the reproducer with the shard layout supplied by the validation probe (VLLM_ROUTED_EXPERTS_E2E_EMULATE_SP=1 + a weights-free tiny Mixtral): the capture code, the TP all-gather (real NCCL across 2 GPUs) and the whole engine/AuxOutput path are real; only "which rows the capture is handed" is emulated, and every summary line says so. On Blackwell, run the same file with VLLM_ROUTED_EXPERTS_E2E_MODEL=<checkpoint> --moe-backend deep_gemm_mega_moe and no emulation flag to get the real path.

Environment

pyxis: imported docker image: nvcr.io#nvidia/cuda:12.4.1-devel-ubuntu22.04
Filesystem      Size  Used Avail Use% Mounted on
tmpfs          1008G  984G   25G  98% /dev/shm
index, name, compute_cap, memory.total [MiB]
0, NVIDIA H20-3e, 9.0, 143771 MiB
1, NVIDIA H20-3e, 9.0, 143771 MiB
torch 2.13.0+cu130 cuda True 2
vllm 0.31.1rc1.dev295+gd4e92ccc3

vllm serve flags (identical for both revisions, from the test's own log):

--tensor-parallel-size 2 --data-parallel-size 1 --max-num-seqs 1 --max-model-len 256
--max-num-batched-tokens 256 --enable-return-routed-experts --enable-expert-parallel
--enforce-eager --gpu-memory-utilization 0.3 --load-format dummy
--worker-extension-cls vllm_re_capture_probe.ProbeWorkerExtension

The two revisions are compared by overlaying the PR's own files on one installed vLLM build, so the diff under test is exactly this PR's diff -?verified by blob hash: the wheel's routed_experts_capturer.py is 5d111fa635-- and forward_context.pyis9ddc8e9255--, i.e. byte-identical to the base revision; the fixed overlay holds e7acb23c04-- / 7b57428dbf--, the PR head versions. [RemoteOpenAIServer] was driven with the probe injected into both workers:

(VllmWorker-1 pid=3307839) INFO [worker_base.py:317] Injected <class 'vllm_re_capture_probe_impl.ProbeWorkerExtension'> into <class 'vllm.v1.worker.gpu_worker.Worker'>
(VllmWorker-0 pid=3307837) INFO [worker_base.py:317] Injected <class 'vllm_re_capture_probe_impl.ProbeWorkerExtension'> into <class 'vllm.v1.worker.gpu_worker.Worker'>

The probe confirms every measured request really was fed a TP shard (6 rows -> 3 rows on tp_rank 0/1, `3 rows -> 2 rows --), so the emulated layout is not assumed, it is recorded.

Base revision -?4 failed, 4 passed, 2 skipped in 439s

FAILED -?:test_six_token_prefill_captures_every_tp_shard
FAILED -?:test_three_token_prefill_drops_only_sp_padding
FAILED -?:test_second_uncached_batch_has_no_stale_rows
FAILED -?:test_unexpected_capture_batch_fails_closed

Failure detail (the bug signature):

E  Failed: six-token prefill: forward 0 layer 0: 3 of 6 rows differ from the raw-shard reference
E  six-token prefill: layer 0: TP rank 0's capture buffer does not hold the whole batch (capture, not the auxiliary-output path, is wrong)
E  Failed: three-token prefill: forward 0 layer 0: 1 of 3 rows differ from the raw-shard reference
E  three-token prefill: layer 0: TP rank 0's capture buffer does not hold the whole batch (-?
E  Failed: first six-token batch: forward 0 layer 0: 3 of 6 rows differ from the raw-shard reference
E  first six-token batch: layer 0: TP rank 0's capture buffer does not hold the whole batch (-?
E  Failed: DID NOT RAISE Exception            # no shape guard on the base revision

Note capture gathers per forward=[0] in the base summary lines: the pre-fix capturer writes only this rank's ceil(T/tp) rows and never gathers, so rows [ceil(T/tp), T) are left over from an earlier forward. Token/layer examples for the 6-token case: token 4 layer 1: got [1 4] expected [4 7], token 5 layer 1: got [1 0] expected [1 7].

PR head -?8 passed, 2 skipped in 445s

[six-token prefill]   returned rows=6, reference rows=6, capture gathers per forward=[1]
[three-token prefill] returned rows=3, reference rows=3, capture gathers per forward=[1]
[first six-token batch]  returned rows=6, reference rows=6, capture gathers per forward=[1]
[second six-token batch] returned rows=6, reference rows=6, capture gathers per forward=[1]
[single-token steps]  returned rows=3, reference rows=3, capture gathers per forward=[0, 0, 0]
[ordinary MoE path at TP=2] returned rows=6, reference rows=6, capture gathers per forward=[0]
[already-gathered batch (3 tokens)] returned rows=3, reference rows=3, gathers=[0]
[already-gathered batch (5 tokens)] returned rows=5, reference rows=5, gathers=[0]
[fail-closed] the request failed as required: InternalServerError("-?EngineCore encountered an issue -?)

• Cases 1-3 return a full batch that matches every raw-shard id, with exactly one gather per forward.
• Case 4 (1-token prefill + 2 decodes) gathers nothing -?TP0 owns the real row.
• The complete-batch control (TP=2+EP, no model-level SP) keeps n == T, gathers nothing, and matches the TP=1 run of the same model.
• Branch n == ceil(T/tp)*tp (already gathered, SP-padded) is trimmed without a second gather; this branch only trims, so it is a guard, not a client-visible difference -?it passes on both revisions.
• Fail-closed evidence, from the probe notes and the worker assertion:

main-eager-unexpected-11-36f659d5: injected unexpected batch: 6 rows -> 1 rows for a 6-token forward at tp_size=2
capture raised AssertionError: RoutedExpertsCapturer: unexpected topk_ids batch dim 1 (expected 6, 3, or 6 for single DP, tp_size=2)

How we know the model shards at DP=1 without running MegaMoE

The shard is not produced by the capture path; it comes from the model's own sequence-parallel flag, which is a pure config predicate. Evaluated on this exact build with the flags the reproducer uses (tp=2, dp=1, ep on, pp=1):

moe_backend=auto                             -> dsv4=False dsv41=False
moe_backend=deep_gemm_mega_moe               -> dsv4=True  dsv41=True
moe_backend=flashinfer_moe_ep_cutedsl        -> dsv4=True  dsv41=True
moe_backend=flashinfer_moe_ep_mega_deep_gemm -> dsv4=True  dsv41=True
framework use_sequence_parallel_moe:         dp=1 -> False (so dp_metadata is None), dp=2 -> True

The repo's own tests/models/test_deepseek_v4_fi_moe_ep.py -k token_sharding pins the same predicate for all three backends on this revision: 3 passed. And dp=1 is precisely why dp_metadata is absent while the model still shards -?the two facts the fix has to reconcile.

That flag is consumed at the SP call sites (vllm/models/deepseek_v41/nvidia/model.py:588,594, vllm/models/deepseek_v4/nvidia/model.py:1253,1257), which call sp_reduce_scatter / sp_shard:

def sp_reduce_scatter(x):                     # vllm/models/common/ops/sequence_parallel.py
    tp_size = get_tensor_model_parallel_world_size()
    sp_pad = (-x.shape[0]) % tp_size            # ceil-div padding, exactly the tail the fix trims
    if sp_pad > 0:
        x = torch.nn.functional.pad(x, (0, 0, 0, sp_pad))
    return tensor_model_parallel_reduce_scatter(x, 0)   # rank r gets contiguous chunk r

so a rank holds exactly rows [r*ceil(T/tp), (r+1)*ceil(T/tp)) -?i.e. n == ceil(T/tp), the branch the fix adds (routed_experts_capturer.py:125-145: n == shard_n -> all_gather(topk_ids)[:num_tokens]), with a zero-padded tail for T % tp != 0. Note auto (or EP off, or pp>1) gives False, which is why the case-6 control never shards.

What this does not claim: that a live DSV4 forward reaches those call sites and what the MegaMoE kernels then do with the shard -?that needs the model to build and load, which on SM90 stops at the backend gate above. The layout the capturer is handed is not invented, though: it is the layout these collectives define, the probe records the row count either way, and the fix's own shape guard fails closed if it ever sees something else.

What is not verified here

• The real MegaMoE kernels and the model-level sp_reduce_scatter that creates the shards: impossible on SM90 (see the gate above). The shard layout is emulated; the capture code, the NCCL gather and the delivery path are not.
• CUDA-graph execution and the VLLM_GPU_SYNC_CHECK=error run (both opt-in in the test, not enabled for this run).
• Multi-DP capture logic, which this PR leaves unchanged.
• MoE numerics/throughput: the weights are random (--load-format dummy), so only routing ids are meaningful. The test prints a warning if a run's routing turns out to be token-independent, because then stale rows could not be told apart by value.

Two environment issues found and fixed while setting this up, for anyone re-running: the Pyxis container's /dev/shm is ~33 MiB by default and vLLM's MessageQueue needs ~240 MiB (bind-mount the host /dev/shm, as done above), and the probe originally capped its note file per process, which hid the injection record for later requests (fixed: one layout note per request/rank).

Logs: ~/vllm-tasks/dsv41-routing-e2e/logs/{base,fix}.log, dumps and probe notes under logs/artifacts-{base,fix}/dumps/.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

✅ @Ronald1995, CI is now available for this PR.

• /ci run starts upstream CI; /amd-ci run starts AMD CI only.
• Your branch must contain every commit currently on its upstream target branch. Merge or rebase onto the latest target branch, then rerun the command. Append --allow-stale to a run command to test an outdated branch at your own risk.
• /ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
• /amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
• /ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

[Ronald1995]

CUDA-graph model E2E + the new regression cases (2xH20-3e, SM90)

Tested revision: fe18b045 ("Propagate TP token counts into replay contexts"), capturer blob 4eadfdd34a--; control = base **206e45d**, blob 5d111fa635--. Both run against one installed vLLM (0.31.1rc1.dev295+gd4e92ccc3, torch 2.13.0+cu130) with only the PR's own files overlaid, so the diff under test is the PR's.

1. The new capturer cases: PR head vs base

pytest tests/model_executor/test_routed_experts_capture.py

overlay	result
PR head (4eadfdd34a)	41 passed (12.8 s)
base (5d111fa635) with the same test file	16 failed, 25 passed

The base failures split cleanly:

• Behavioural -?this PR's point (3): test_single_dp_ep_capture_restores_tp_shards[2-6], [2-3], [4-2] ->?assert False at the shard comparison (test file line 370). The base capturer writes only ceil(T/tp) rows.
• Contract change (1): test_routed_experts_capturer_rejects_incomplete_dp_shard_groups ->?base raises AssertionError where this revision raises ValueError.
• Scope, not behaviour (12): TypeError: ForwardContext.__init__() got an unexpected keyword argument 'num_tokens'. These cases build the PR's ForwardContext.num_tokens, which the base revision does not have at all, so base cannot even construct the context (single_dp_no_metadata, narrows_snapshot[2], sp_modular_gathers_tp_shards[2], does_not_regather_full_batch[3], rejects_unexpected_batch_without_writing[2], rejects_layer_before_restoring_rows[2]).

Independent check of the layout helper: I extract the real _restore_local_topk_ids body from the head file and drive it through 23 constructed DP=1 / DP>1 / edge cases (extracted source lines 97-185, sha1 1177b8559854): 23/23. The head version also closes the review points raised earlier (layer index validated before the collective, ValueError instead of assert, TP gather input always contiguous, layout contracts documented, PCP declared unsupported).

2. CUDA-graph model E2E (case 5 of the reproducer, 2 GPUs)

--tensor-parallel-size 2 --data-parallel-size 1 --max-num-seqs 1 --max-model-len 256 --max-num-batched-tokens 256 --enable-return-routed-experts --enable-expert-parallel --load-format dummy --gpu-memory-utilization 0.3 plus a worker-extension probe. Case 5 builds eager references for a 6-token and a 3-token prompt (validating that reference first), then re-runs both prompts without --enforce-eager and compares the returned routed experts.

base -?fails before it can compare. The eager reference itself is wrong, i.e. the same pre-fix defect:

eager reference (6 tokens): layer 0: TP rank 0's capture buffer does not hold the whole batch (capture, not the auxiliary-output path, is wrong)
eager reference (6 tokens): forward 0 layer 0: 3 of 6 rows differ from the raw-shard reference
1 failed, 9 deselected

PR head -?1 passed (203.9 s). CUDA graphs really were captured in the graph-mode server (Capturing CUDA graphs (PIECEWISE): 2/2), and the probe recorded the shard layout for both graph-mode forwards, so the TP-shard branch is what ran:

main-graph-shard-3-f4b31dd2: injected shard layout: 6 rows -> 3 rows on tp_rank 0 / tp_rank 1 for a 6-token forward
main-graph-shard-4-e1dfcf25: injected shard layout: 3 rows -> 2 rows on tp_rank 0 / tp_rank 1 for a 3-token forward
[eager reference (6 tokens)] returned rows=6, reference rows=6, capture gathers per forward=[1]
[eager reference (3 tokens)] returned rows=3, reference rows=3, capture gathers per forward=[1]
[cuda-graph] prompt 6 tokens: dumped 4 capture records during graph execution      (2 layers x 2 ranks)
[cuda-graph (6 tokens), reference from its own capture records] returned rows=6, reference rows=6, gathers=[1]
[cuda-graph] prompt 3 tokens: dumped 4 capture records during graph execution
[cuda-graph (3 tokens), reference from its own capture records] returned rows=3, reference rows=3, gathers=[1]
1 passed, 9 deselected

So in graph mode the capture path ran (4 records per prompt), each forward took exactly one all-gather -?the n == ceil(T/tp) branch with padding trimmed -?the returned rows match the eager reference for both prompt shapes, and the replay did not hang.

Environment note (not a PR issue): the graph-mode server first died at startup building flashinfer's sampling JIT -?this image's nvcc is CUDA 12.4 and rejects the --compress-mode=size flag flashinfer 0.7.0 passes (nvcc fatal: Unknown option '--compress-mode=size'). Eager servers never need that module. The run above uses VLLM_USE_FLASHINFER_SAMPLER=0; the sampler backend is orthogonal to routed-expert capture.

3. Full reproducer suite at this revision

pytest tests/models/deepseek_v41_routed_experts_capture_e2e.py on the PR-head overlay (2 GPUs, eager, shard layout probe-injected):

• 9 passed, 1 skipped in 522 s -?the skip is the opt-in VLLM_GPU_SYNC_CHECK case.
• That case on its own: 1 passed in 86 s -?[sync-check] a sharded 6-token prefill plus one decode completed under VLLM_GPU_SYNC_CHECK with 7 routed-expert rows (PyTorch sync-debug mode active).

Every case therefore passes at fe18b045: 6-token and 3-token sharded prefill, a second uncached batch, 1-token steps with no gather, CUDA-graph replay, the TP+EP control and its TP=1 oracle, both already-gathered padded batches, the fail-closed guard and the sync check.

[six-token prefill]        returned rows=6, reference rows=6, capture gathers per forward=[1]
[three-token prefill]      returned rows=3, reference rows=3, capture gathers per forward=[1]
[first/second six-token batch] returned rows=6, reference rows=6, gathers=[1]
[single-token steps]       returned rows=3, reference rows=3, gathers=[0, 0, 0]
[cuda-graph (6 / 3 tokens)] returned rows=6/3, reference rows=6/3, gathers=[1]
[ordinary MoE path at TP=2] returned rows=6, reference rows=6, gathers=[0]
[already-gathered batch (3 / 5 tokens)] returned 3 == 3 / 5 == 5, gathers=[0]
[fail-closed] the request failed as required: InternalServerError("... 500 ...")

4. What this still does not cover

• Real DL/DSV4/V4.1/Kimi-K3 MegaMoE kernels and the model-level sp_reduce_scatter: all three MEGA_MOE_BACKENDS members are hard SM100/SM103 gates, so on H20 the shard layout is supplied by the probe. The capture code, the NCCL all-gather, the AuxOutput plumbing, the engine and the CUDA graphs are real. On a B200/B300 host, run the same file with a real checkpoint and without the emulation flag to cover the remainder.
• PIECEWISE CUDA graphs only (vLLM default); --cudagraph-mode full was not run. The VLLM_GPU_SYNC_CHECK=error case is covered in section 3.
• The model is weights-free (--load-format dummy), so only routing ids are meaningful; the run warns if a prompt's routing degenerates to a token-independent expert set.

Logs and dumps: ~/vllm-tasks/dsv41-routing-e2e/logs/{graph-base,graph-fix,graph-fix2}.log, logs/graph-fix2-artifacts/dumps/. Results are for fe18b045; if the head moves again I can re-run the same two commands.

[Ronald1995]

/ci run

[github-actions]

✅ Triggered Buildkite CI #94984 for commit 92a36b2e26df.

[Ronald1995]

/ci retry

[github-actions]

✅ No failed, timed-out, or expired jobs need retrying: https://buildkite.com/vllm/ci/builds/94984

[Ronald1995]

/ci retry

[github-actions]

✅ Queued 1 failed job(s) for retry in Buildkite CI #94984.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95067 for commit 3269afcf6983.

[aoshen02]

/ci run

[github-actions]

✅ Triggered Buildkite CI #95145 for commit 3f13cb853b58.