[Perf] Avoid torch Stream construction across grouped MoE launches

aiter/ops/flydsl/grouped_gemm_mxfp4.py

@@ -10,7 +10,7 @@
 import torch
 
 from .kernels.mega_moe_gfx1250.types import Stage2ScatterContext
-from .kernels.tensor_shim import _run_compiled, ptr_arg
+from .kernels.tensor_shim import _run_compiled, current_stream_arg, ptr_arg
 
 _SUPPORTED_CLUSTER_M = (1, 2, 4)
 _SUPPORTED_CLUSTER_N = (1, 2, 3, 4)
@@ -208,7 +208,7 @@ def flydsl_grouped_gemm_a8w4_masked(
     from .kernels.mxfp4_preshuffle_gfx1250_tdm import launch_gemm_a8w4_tdm
 
     if stream is None:
-        stream = torch.cuda.current_stream()
+        stream = current_stream_arg()
     if stage1_act == 3:
         if float(situ_beta) <= 0.0:
             raise ValueError(f"situ_beta must be > 0, got {situ_beta!r}")

aiter/ops/flydsl/grouped_moe_gfx1250.py

@@ -20,7 +20,11 @@
 from aiter import ActivationType, QuantType, dtypes, logger
 from aiter.jit.utils.chip_info import get_gfx
 from aiter.ops.flydsl.kernels.mega_moe_gfx1250.types import Stage2ScatterContext
-from aiter.ops.flydsl.kernels.tensor_shim import _run_compiled, ptr_arg
+from aiter.ops.flydsl.kernels.tensor_shim import (
+    _run_compiled,
+    current_stream_arg,
+    ptr_arg,
+)
 
 # Opt-in switch for the gfx1250 FlyDSL grouped-GEMM path.
 _TRUTHY_ENV = ("1", "true", "True", "yes", "YES")
@@ -1987,7 +1991,7 @@ def fused_route_psum_remap(
         experts,
         int(max_m),
         int(tile_m),
-        torch.cuda.current_stream(),
+        current_stream_arg(),
     )
     return masked_m, topids_to_rows.view(token_num, topk), psum
 
@@ -2157,7 +2161,7 @@ def contiguous_psum_remap(
             int(ep_scatter_params["max_tok"]),
             int(ep_scatter_params["slot_stride"]),
             ptr_arg(row_to_token_i32),
-            torch.cuda.current_stream(),
+            current_stream_arg(),
         )
         return starts, psum, contiguous_m_t
     from aiter.ops.flydsl.kernels.moe_contiguous_psum import MAX_REMAP_EXPERTS
@@ -2181,7 +2185,7 @@ def contiguous_psum_remap(
         ptr_arg(num_valid_routes_i32),
         ptr_arg(row_to_token_i32),
         int(topk),
-        torch.cuda.current_stream(),
+        current_stream_arg(),
     )
     return starts, psum, contiguous_m_t
 
@@ -2285,7 +2289,7 @@ def flydsl_moe_gather_reduce(
         token_num,
         slice_stride_dw,
         ptr_arg(num_valid_tokens_i32),
-        torch.cuda.current_stream(),
+        current_stream_arg(),
     )
     return out
 
@@ -2412,7 +2416,7 @@ def flydsl_moe_scatter_preshuffle_scale(
         max_m,
         E,
         tiles_per_expert,
-        torch.cuda.current_stream(),
+        current_stream_arg(),
     )
     return grouped_a1_scale
 

aiter/ops/flydsl/kernels/tensor_shim.py

@@ -277,6 +277,18 @@ def check_e8m0(
         raise RuntimeError(f"[{owner}] {name} must be contiguous")
 
 
+# Resolve the optional API once, but query its device/stream on every call.
+# This is not a stream cache: graph capture and thread-local stream switches
+# must still be observed. Older torch builds keep the public API fallback.
+_get_current_raw_stream = getattr(torch._C, "_cuda_getCurrentRawStream", None)
+
+
+def current_stream_arg():
+    if _get_current_raw_stream is None:
+        return torch.cuda.current_stream()
+    return fx.Stream(_get_current_raw_stream(torch.cuda.current_device()))
+
+
 def ptr_arg(t: torch.Tensor, dtype=None):
     """Wrap a torch.Tensor as an fx.Pointer (PointerJitArg) for kernel launch."""
     if dtype is None:

aiter/ops/flydsl/moe_kernels.py

@@ -13,6 +13,7 @@
     _run_compiled as _dispatch_compiled,
 )
 from aiter.ops.flydsl.kernels.tensor_shim import (
+    current_stream_arg,
     ptr_arg,
 )
 
@@ -2605,7 +2606,7 @@ def flydsl_moe_topids_to_rows(
             numel,
             int(max_m),
             int(E),
-            torch.cuda.current_stream(),
+            current_stream_arg(),
         )
         return counter, topids_to_rows.view(token_num, topk)
 
@@ -2671,7 +2672,7 @@ def flydsl_moe_topids_to_rows(
             topids_to_rows_kernel,
             *route_args,
             *_ep_tail,
-            torch.cuda.current_stream(),
+            current_stream_arg(),
         )
     else:
         from aiter.jit.utils.chip_info import get_gfx_runtime
@@ -2700,7 +2701,7 @@ def flydsl_moe_topids_to_rows(
             numel,
             int(max_m),
             route_grid,
-            torch.cuda.current_stream(),
+            current_stream_arg(),
         )
     return counter, topids_to_rows.view(token_num, topk)
 
@@ -3530,7 +3531,7 @@ def flydsl_moe_fused_quant_preshuffle(
                 ptr_arg(row_to_token.reshape(-1)),
                 token_num,
                 (token_num + warps_per_block - 1) // warps_per_block,
-                torch.cuda.current_stream(),
+                current_stream_arg(),
             )
             return out_payload, out_scale
         if use_token_multidest:
@@ -3561,7 +3562,7 @@ def flydsl_moe_fused_quant_preshuffle(
                 ptr_arg(topids_to_rows_i32),
                 token_num,
                 token_grid,
-                torch.cuda.current_stream(),
+                current_stream_arg(),
             )
             return out_payload, out_scale
         # The M16 scatter consumes a finalized inverse map and K1024 tiles.
@@ -3584,7 +3585,7 @@ def flydsl_moe_fused_quant_preshuffle(
             from aiter.jit.utils.chip_info import get_cu_num
 
             apre_workers = max(1, 4 * (get_cu_num() or 1))
-            stream = torch.cuda.current_stream()
+            stream = current_stream_arg()
             if prequantized:
                 # MoRI already produced packed FP4 and padded E8M0 rows.
                 compact_payload = grouped_in.contiguous()
@@ -3669,7 +3670,7 @@ def flydsl_moe_fused_quant_preshuffle(
             ),
             *_routeks_ep_psum_args(device, ep_psum_params),
             grid_blocks,
-            torch.cuda.current_stream(),
+            current_stream_arg(),
         )
         return out_payload, out_scale
 
@@ -3691,6 +3692,6 @@ def flydsl_moe_fused_quant_preshuffle(
         n_rows,
         max_m,
         grid_blocks,
-        torch.cuda.current_stream(),
+        current_stream_arg(),
     )
     return out_payload, out_scale

op_tests/test_grouped_gemm_stream.py

@@ -0,0 +1,195 @@
+# SPDX-License-Identifier: MIT
+# Copyright (C) 2026, Advanced Micro Devices, Inc. All rights reserved.
+"""Default stream handle selection must preserve explicit streams and ordering."""
+
+import ctypes
+import importlib
+import importlib.util
+from concurrent.futures import ThreadPoolExecutor
+from pathlib import Path
+
+import pytest
+import torch
+
+gm = importlib.import_module("aiter.ops.flydsl.grouped_gemm_mxfp4")
+shim = importlib.import_module("aiter.ops.flydsl.kernels.tensor_shim")
+moe = importlib.import_module("aiter.ops.flydsl.moe_kernels")
+grouped = importlib.import_module("aiter.ops.flydsl.grouped_moe_gfx1250")
+
+
+def _handle(value):
+    from flydsl.expr import Stream
+
+    stream = value if isinstance(value, Stream) else Stream(value)
+    storage = ctypes.c_void_p()
+    _, fill = stream.__c_abi_spec__()[0]
+    fill(stream, storage)
+    return storage.value or 0
+
+
+def test_raw_handle_is_queried_each_time(monkeypatch):
+    from flydsl.expr import Stream
+
+    state = {"device": 0, "handles": [101, 202]}
+    seen = []
+
+    def raw(device):
+        seen.append(device)
+        return state["handles"][device]
+
+    monkeypatch.setattr(shim, "_get_current_raw_stream", raw)
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: state["device"])
+    monkeypatch.setattr(
+        torch.cuda,
+        "current_stream",
+        lambda *a, **kw: pytest.fail("constructed a torch Stream on the fast path"),
+    )
+    first = shim.current_stream_arg()
+    assert isinstance(first, Stream) and _handle(first) == 101
+    state["handles"][0] = 303
+    assert _handle(shim.current_stream_arg()) == 303
+    state["device"] = 1
+    assert _handle(shim.current_stream_arg()) == 202
+    assert seen == [0, 0, 1]
+    assert _handle(first) == 101  # No mutable Stream object reused across calls.
+
+
+def test_missing_private_api_uses_public_fallback(monkeypatch):
+    sentinel = object()
+    monkeypatch.setattr(shim, "_get_current_raw_stream", None)
+    monkeypatch.setattr(torch.cuda, "current_stream", lambda: sentinel)
+    monkeypatch.setattr(
+        torch.cuda,
+        "current_device",
+        lambda: pytest.fail("fallback must delegate device selection to torch"),
+    )
+    assert shim.current_stream_arg() is sentinel
+
+
+def _gpu():
+    if not torch.cuda.is_available():
+        pytest.skip("requires ROCm GPU")
+    from aiter.jit.utils.chip_info import get_gfx_runtime
+
+    if get_gfx_runtime() != "gfx1250":
+        pytest.skip("grouped GEMM requires gfx1250")
+
+
+def test_device_and_thread_local_handles():
+    _gpu()
+    # Handle selection itself is device independent. Do not reuse a compiled
+    # FlyDSL executable on another device; that is a separate runtime concern.
+    for device in range(min(torch.cuda.device_count(), 2)):
+        with torch.cuda.device(device):
+            streams = [torch.cuda.default_stream(), torch.cuda.Stream()]
+            for stream in streams:
+                with torch.cuda.stream(stream):
+                    assert _handle(shim.current_stream_arg()) == stream.cuda_stream
+
+    device = torch.cuda.current_device()
+    streams = [torch.cuda.Stream(device=device), torch.cuda.Stream(device=device)]
+
+    def worker(index):
+        with torch.cuda.device(device), torch.cuda.stream(streams[index]):
+            return all(
+                _handle(shim.current_stream_arg()) == streams[index].cuda_stream
+                for _ in range(100)
+            )
+
+    with ThreadPoolExecutor(max_workers=2) as pool:
+        assert all(pool.map(worker, range(2)))
+
+
+@pytest.mark.parametrize("data_format", ["a4w4", "a8w4"])
+@pytest.mark.parametrize("tokens", [1, 16, 257])
+def test_real_moe_stream_order_and_graph(data_format, tokens, monkeypatch):
+    _gpu()
+    fm = importlib.import_module("aiter.fused_moe")
+    path = Path(__file__).parent / "flydsl_tests" / "test_flydsl_grouped_gemm.py"
+    spec = importlib.util.spec_from_file_location("_stream_grouped_reference", path)
+    reference = importlib.util.module_from_spec(spec)
+    default_device = torch.get_default_device()
+    try:
+        spec.loader.exec_module(reference)
+        inputs = []
+        original_public = reference.fused_moe
+
+        def capture(*args, **kwargs):
+            kwargs["quant_dtype_a"] = (
+                fm.dtypes.fp8 if data_format == "a8w4" else fm.dtypes.fp4x2
+            )
+            inputs.append((args, kwargs))
+            return original_public(*args, **kwargs)
+
+        reference.fused_moe = capture
+        actual, expected, _, _ = reference._run_grouped_via_fused_moe(
+            experts=8,
+            tokens=tokens,
+            topk=2,
+            model_dim=256,
+            inter_dim=256,
+            data_format=data_format,
+            activation=fm.ActivationType.Swiglu,
+            use_bias=True,
+            seed=42,
+            bench=False,
+        )
+        assert reference._logits_diff(actual, expected) < reference.LOGITS_DIFF_TOL
+        args, kwargs = inputs[-1]
+        new = gm.flydsl_grouped_gemm_a8w4_masked
+
+        def old(*a, **kw):
+            if kw.get("stream") is None:
+                kw["stream"] = torch.cuda.current_stream()
+            return new(*a, **kw)
+
+        with monkeypatch.context() as patch:
+            patch.setattr(gm, "flydsl_grouped_gemm_a8w4_masked", old)
+            patch.setattr(moe, "current_stream_arg", torch.cuda.current_stream)
+            patch.setattr(grouped, "current_stream_arg", torch.cuda.current_stream)
+            before = fm.fused_moe(*args, **kwargs)
+        assert torch.equal(before, actual)
+
+        launches = []
+
+        def record(*a, **kw):
+            launches.append((a, kw))
+            return new(*a, **kw)
+
+        with monkeypatch.context() as patch:
+            patch.setattr(gm, "flydsl_grouped_gemm_a8w4_masked", record)
+            fm.fused_moe(*args, **kwargs)
+        a, kw = launches[0]
+        stream = torch.cuda.Stream()
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            # Explicit stream must be forwarded untouched, without querying the
+            # default-stream helper. Compare against the original launch.
+            old(*a, **{**kw, "stream": stream})
+            explicit_expected = a[0].clone()
+            with monkeypatch.context() as patch:
+                patch.setattr(
+                    gm,
+                    "current_stream_arg",
+                    lambda: pytest.fail("explicit stream replaced"),
+                )
+                new(*a, **{**kw, "stream": stream})
+            assert torch.equal(a[0], explicit_expected)
+            fm.fused_moe(*args, **kwargs)  # Warm before capture.
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph, stream=stream):
+                replayed = fm.fused_moe(*args, **kwargs)
+            for _ in range(2):
+                args[0].neg_()
+                graph.replay()
+                with monkeypatch.context() as patch:
+                    patch.setattr(gm, "flydsl_grouped_gemm_a8w4_masked", old)
+                    patch.setattr(moe, "current_stream_arg", torch.cuda.current_stream)
+                    patch.setattr(
+                        grouped, "current_stream_arg", torch.cuda.current_stream
+                    )
+                    eager = fm.fused_moe(*args, **kwargs)
+                assert torch.equal(replayed, eager)
+        torch.cuda.current_stream().wait_stream(stream)
+    finally:
+        torch.set_default_device(default_device)