[eager MOE] Selected expert IDs should probably not be treated as 0-d GPU tensors

- `transformers` version: 5.15.1
- Platform: Linux-5.15.0-70-generic-x86_64-with-glibc2.35
- Python version: 3.12.13
- Huggingface_hub version: 1.27.0
- Safetensors version: 0.8.0
- Accelerate version: 1.14.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (accelerator?): 2.12.0+git6bbd260 (CUDA)
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: yes
- GPU type: AMD MI350

...
        for expert_idx in expert_hit[:, 0].tolist():
...

device=cuda, dtype=torch.bfloat16, tokens=4096, top_k=8, experts=256, hidden=2048, intermediate=512
Each forward selects every expert; baseline runs first, then tolist().
upstream CUDA-scalar loop      median=0.055140s mean=0.055138s min=0.054309s
one-time expert_hit.tolist()   median=0.038659s mean=0.038688s min=0.038511s
outputs: identical; median speedup: 1.43x
/usr/local/lib/python3.12/dist-packages/torch/profiler/profiler.py:272: UserWarning: Warning: Profiler clears events at the end of each cycle.Only events from the current cycle will be reported.To keep events across cycles, set acc_events=True.
  _warn_once(
USDT:2026-09-29 15:20:40 3162222:3162222 ActivityProfilerController.cpp:415] profiler_start
USDT:2026-09-29 15:20:40 3162222:3162222 ActivityProfilerController.cpp:455] profiler_stop
[W929 15:20:40.278328056 collection.cpp:1182] Warning: ROCTracer produced duplicate flow start: 1 (function operator())

Torch Profiler summary: upstream CUDA-scalar loop
  all aten:: CPU ops: count=19,227, inclusive_cpu_time=0.233808s
  aten::item: count=1,024, inclusive_cpu_time=0.013411s
  aten::_local_scalar_dense: count=1,024, inclusive_cpu_time=0.012870s
  Chrome trace: .debug/qwen35_moe_traces/qwen35_moe_scalar_indexing_trace.json
USDT:2026-09-29 15:20:42 3162222:3162222 ActivityProfilerController.cpp:415] profiler_start
USDT:2026-09-29 15:20:42 3162222:3162222 ActivityProfilerController.cpp:455] profiler_stop

Torch Profiler summary: one-time expert_hit.tolist()
  all aten:: CPU ops: count=15,650, inclusive_cpu_time=0.197983s
  aten::item: count=0, inclusive_cpu_time=0.000000s
  aten::_local_scalar_dense: count=0, inclusive_cpu_time=0.000000s
  Chrome trace: .debug/qwen35_moe_traces/qwen35_moe_tolist_trace.json

ATen CPU-op count differences (baseline -> tolist()):
  operation                                      baseline    tolist      delta   baseline CPU   tolist CPU
  aten::item                                        1,024         0    -1,024     0.013411s     0.000000s
  aten::_local_scalar_dense                         1,024         0    -1,024     0.012870s     0.000000s
  aten::select                                      1,792     1,281      -511     0.001646s     0.001237s
  aten::as_strided                                  3,845     3,334      -511     0.001128s     0.001020s
  aten::is_nonzero                                    256         0      -256     0.004266s     0.000000s
  aten::eq                                            256         0      -256     0.003840s     0.000000s
  aten::unbind                                        257       256        -1     0.001033s     0.000790s
  aten::to                                            256       257        +1     0.000043s     0.000116s
  aten::resolve_neg                                     0         1        +1     0.000000s     0.000001s
  aten::resolve_conj                                    0         1        +1     0.000000s     0.000001s
  aten::empty_strided                                   1         2        +1     0.000039s     0.000041s
  aten::copy_                                         256       257        +1     0.003247s     0.003434s
  aten::_to_copy                                        0         1        +1     0.000000s     0.000076s

#!/usr/bin/env python3
"""Standalone reproducer for Qwen3.5 MoE CUDA-scalar indexing overhead.

This intentionally has no AMD Quark or Transformers dependency.  It contains
the relevant structure of Transformers' ``Qwen3_5MoeExperts.forward`` and
compares its original CUDA-scalar loop against a semantically equivalent loop
that materializes all selected expert IDs once with ``tolist()``.

Example on CUDA or ROCm::

    HIP_VISIBLE_DEVICES=7 python .debug/benchmark_qwen35_moe_expert_indexing.py

For a closer approximation of one Qwen calibration forward, increase tokens::

    HIP_VISIBLE_DEVICES=7 python .debug/benchmark_qwen35_moe_expert_indexing.py \
        --tokens 16384 --warmup 2 --iterations 5

The default dimensions deliberately fit on a typical accelerator while still
selecting all 256 experts.  The two models use identical weights and inputs;
the script validates their outputs before reporting synchronized timings.
"""

from __future__ import annotations

import argparse
import statistics
import time
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path

import torch
from torch import nn


@dataclass(frozen=True)
class Qwen3_5MoeConfig:
    """Minimal subset of the Transformers Qwen3.5 MoE config used below."""

    num_experts: int
    hidden_size: int
    moe_intermediate_size: int


class Qwen3_5MoeExperts(nn.Module):
    """Minimal copy of Transformers' 3-D-parameter Qwen3.5 experts module.

    ``materialize_expert_ids`` is the only intentional difference between the
    two benchmark variants.  The ``False`` path preserves the upstream loop:
    ``expert_idx`` remains a CUDA scalar tensor and is used for the sentinel
    check and three tensor indexes.  The ``True`` path transfers the short
    list of selected ids once and then uses ordinary Python integers.
    """

    def __init__(self, config: Qwen3_5MoeConfig, *, materialize_expert_ids: bool) -> None:
        super().__init__()
        self.num_experts = config.num_experts
        self.hidden_dim = config.hidden_size
        self.intermediate_dim = config.moe_intermediate_size
        self.gate_up_proj = nn.Parameter(
            torch.empty(self.num_experts, 2 * self.intermediate_dim, self.hidden_dim)
        )
        self.down_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_dim, self.intermediate_dim))
        self.materialize_expert_ids = materialize_expert_ids

    def forward(
        self,
        hidden_states: torch.Tensor,
        top_k_index: torch.Tensor,
        top_k_weights: torch.Tensor,
    ) -> torch.Tensor:
        final_hidden_states = torch.zeros_like(hidden_states)
        with torch.no_grad():
            expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
            expert_mask = expert_mask.permute(2, 1, 0)
            expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()

        if self.materialize_expert_ids:
            # Proposed Transformers change: one D2H transfer for all selected
            # IDs, rather than converting the same CUDA scalar for every use.
            expert_ids: list[int] | torch.Tensor = expert_hit[:, 0].tolist()
        else:
            # Upstream behavior: iteration produces one-device scalar tensors.
            expert_ids = expert_hit

        for expert_idx in expert_ids:
            if not self.materialize_expert_ids:
                expert_idx = expert_idx[0]
            if expert_idx == self.num_experts:
                continue
            top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
            current_state = hidden_states[token_idx]
            gate, up = torch.nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
            current_hidden_states = torch.nn.functional.silu(gate) * up
            current_hidden_states = torch.nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
            current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
            final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))

        return final_hidden_states


def parse_dtype(name: str) -> torch.dtype:
    try:
        return {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[name]
    except KeyError as error:
        raise ValueError(f"Unsupported dtype {name!r}.") from error


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--dtype", choices=("bfloat16", "float16", "float32"), default="bfloat16")
    parser.add_argument("--num-experts", type=int, default=256)
    parser.add_argument("--hidden-size", type=int, default=2048)
    parser.add_argument("--intermediate-size", type=int, default=512)
    parser.add_argument("--tokens", type=int, default=4096)
    parser.add_argument("--top-k", type=int, default=8)
    parser.add_argument("--warmup", type=int, default=2)
    parser.add_argument("--iterations", type=int, default=8)
    parser.add_argument(
        "--torch-profiler",
        action="store_true",
        help="Run one additional torch.profiler forward for each variant and print CPU-op summaries.",
    )
    parser.add_argument(
        "--profile-trace-dir",
        type=Path,
        help="Optional directory in which --torch-profiler writes one Chrome trace per variant.",
    )
    parser.add_argument(
        "--profile-diff-limit",
        type=int,
        default=30,
        help="Maximum number of differing ATen operation counts to display (default: 30).",
    )
    return parser.parse_args()


def synchronize(device: torch.device) -> None:
    if device.type == "cuda":
        torch.cuda.synchronize(device)


@torch.no_grad()
def benchmark(
    name: str,
    model: nn.Module,
    hidden_states: torch.Tensor,
    top_k_index: torch.Tensor,
    top_k_weights: torch.Tensor,
    warmup: int,
    iterations: int,
) -> tuple[torch.Tensor, list[float]]:
    for _ in range(warmup):
        model(hidden_states, top_k_index, top_k_weights)

    synchronize(hidden_states.device)
    samples: list[float] = []
    output: torch.Tensor | None = None
    for _ in range(iterations):
        start = time.perf_counter()
        output = model(hidden_states, top_k_index, top_k_weights)
        synchronize(hidden_states.device)
        samples.append(time.perf_counter() - start)

    assert output is not None
    print(
        f"{name:30} median={statistics.median(samples):.6f}s "
        f"mean={statistics.mean(samples):.6f}s min={min(samples):.6f}s"
    )
    return output, samples


@torch.no_grad()
def profile_one_forward(
    name: str,
    model: nn.Module,
    hidden_states: torch.Tensor,
    top_k_index: torch.Tensor,
    top_k_weights: torch.Tensor,
    trace_path: Path | None,
) -> dict[str, tuple[int, float]]:
    """Profile one synchronized forward and print the requested CPU-op totals.

    CPU-op durations are inclusive, matching a Chrome-trace category summary:
    nested ATen calls therefore must not be summed as wall-clock time.
    """
    activities = [torch.profiler.ProfilerActivity.CPU]
    if hidden_states.device.type == "cuda":
        activities.append(torch.profiler.ProfilerActivity.CUDA)

    synchronize(hidden_states.device)
    with torch.profiler.profile(activities=activities) as profiler:
        model(hidden_states, top_k_index, top_k_weights)
        synchronize(hidden_states.device)

    if trace_path is not None:
        trace_path.parent.mkdir(parents=True, exist_ok=True)
        profiler.export_chrome_trace(str(trace_path))

    events = [event for event in profiler.events() if event.name.startswith("aten::")]
    per_op: dict[str, list[float]] = defaultdict(lambda: [0.0, 0.0])
    for event in events:
        per_op[event.name][0] += 1
        per_op[event.name][1] += event.cpu_time_total
    metrics = {name: (int(count), cpu_us) for name, (count, cpu_us) in per_op.items()}

    def summarize(event_name: str) -> tuple[int, float]:
        matches = [event for event in events if event.name == event_name]
        return len(matches), sum(event.cpu_time_total for event in matches)

    total_cpu_us = sum(event.cpu_time_total for event in events)
    print(f"\nTorch Profiler summary: {name}")
    print(f"  all aten:: CPU ops: count={len(events):,}, inclusive_cpu_time={total_cpu_us / 1e6:.6f}s")
    for event_name in ("aten::item", "aten::_local_scalar_dense"):
        count, cpu_us = summarize(event_name)
        print(f"  {event_name}: count={count:,}, inclusive_cpu_time={cpu_us / 1e6:.6f}s")
    if trace_path is not None:
        print(f"  Chrome trace: {trace_path}")
    return metrics


def print_profile_count_differences(
    baseline: dict[str, tuple[int, float]],
    patched: dict[str, tuple[int, float]],
    limit: int,
) -> None:
    """Display the ATen operations whose event counts differ between variants."""
    differences = []
    for name in baseline.keys() | patched.keys():
        baseline_count, baseline_cpu_us = baseline.get(name, (0, 0.0))
        patched_count, patched_cpu_us = patched.get(name, (0, 0.0))
        if baseline_count != patched_count:
            differences.append((abs(patched_count - baseline_count), name, baseline_count, patched_count, baseline_cpu_us, patched_cpu_us))

    print("\nATen CPU-op count differences (baseline -> tolist()):")
    print("  operation                                      baseline    tolist      delta   baseline CPU   tolist CPU")
    for _, name, baseline_count, patched_count, baseline_cpu_us, patched_cpu_us in sorted(
        differences, reverse=True
    )[:limit]:
        print(
            f"  {name:45} {baseline_count:9,} {patched_count:9,} "
            f"{patched_count - baseline_count:+9,} {baseline_cpu_us / 1e6:12.6f}s {patched_cpu_us / 1e6:12.6f}s"
        )


def main() -> None:
    args = parse_args()
    device = torch.device(args.device)
    dtype = parse_dtype(args.dtype)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("This reproducer requires a CUDA or ROCm accelerator.")
    if args.tokens < args.num_experts:
        raise ValueError("--tokens must be at least --num-experts so every expert can be selected.")

    config = Qwen3_5MoeConfig(args.num_experts, args.hidden_size, args.intermediate_size)
    torch.manual_seed(0)
    baseline = Qwen3_5MoeExperts(config, materialize_expert_ids=False).to(device=device, dtype=dtype).eval()
    patched = Qwen3_5MoeExperts(config, materialize_expert_ids=True).to(device=device, dtype=dtype).eval()
    patched.load_state_dict(baseline.state_dict())

    # Cycling ensures every expert is selected even with small token counts.
    expert_ids = torch.arange(args.tokens * args.top_k, device=device).reshape(args.tokens, args.top_k)
    top_k_index = expert_ids.remainder(args.num_experts)
    top_k_weights = torch.full((args.tokens, args.top_k), 1.0 / args.top_k, device=device, dtype=dtype)
    hidden_states = torch.randn(args.tokens, args.hidden_size, device=device, dtype=dtype)

    print(
        f"device={device}, dtype={dtype}, tokens={args.tokens}, top_k={args.top_k}, "
        f"experts={args.num_experts}, hidden={args.hidden_size}, intermediate={args.intermediate_size}"
    )
    print("Each forward selects every expert; baseline runs first, then tolist().")

    baseline_output, baseline_samples = benchmark(
        "upstream CUDA-scalar loop",
        baseline,
        hidden_states,
        top_k_index,
        top_k_weights,
        args.warmup,
        args.iterations,
    )
    patched_output, patched_samples = benchmark(
        "one-time expert_hit.tolist()",
        patched,
        hidden_states,
        top_k_index,
        top_k_weights,
        args.warmup,
        args.iterations,
    )

    torch.testing.assert_close(baseline_output, patched_output, rtol=0, atol=0)
    speedup = statistics.median(baseline_samples) / statistics.median(patched_samples)
    print(f"outputs: identical; median speedup: {speedup:.2f}x")

    if args.torch_profiler:
        profile_dir = args.profile_trace_dir
        baseline_metrics = profile_one_forward(
            "upstream CUDA-scalar loop",
            baseline,
            hidden_states,
            top_k_index,
            top_k_weights,
            None if profile_dir is None else profile_dir / "qwen35_moe_scalar_indexing_trace.json",
        )
        patched_metrics = profile_one_forward(
            "one-time expert_hit.tolist()",
            patched,
            hidden_states,
            top_k_index,
            top_k_weights,
            None if profile_dir is None else profile_dir / "qwen35_moe_tolist_trace.json",
        )
        print_profile_count_differences(baseline_metrics, patched_metrics, args.profile_diff_limit)


if __name__ == "__main__":
    main()