manitest

benchmarks/ops/bench_moe_expand_to_fused.py

@@ -0,0 +1,504 @@
+"""Benchmark for MoE expand-to-fused scatter ops.
+
+Benchmarks:
+  - MoeExpandToFusedFwdOp: bf16/fp16 activation scatter into fused layout.
+  - MoeExpandToFusedWithSFFwdOp: fp8 activation scatter with scale factors.
+
+Baselines:
+  - PyTorch reference: vectorized scatter into the same expanded layout.
+
+Real model configurations cover Kimi K2 and Qwen3-30B shapes.
+"""
+
+from typing import Any
+
+import pytest
+import torch
+
+from benchmarks.benchmark_base import BenchmarkReport, ManifestBenchmark
+from tileops.manifest import load_workloads
+from tileops.ops.moe import MoeExpandToFusedFwdOp, MoeExpandToFusedWithSFFwdOp
+from workloads.workload_base import WorkloadBase
+
+_FP8_E4M3 = getattr(torch, "float8_e4m3fn", None)
+_EXPAND_OP_NAME = "MoeExpandToFusedFwdOp"
+_EXPAND_WITH_SF_OP_NAME = "MoeExpandToFusedWithSFFwdOp"
+_UE8M0_PER_INT32 = 4
+
+_BASE_FALLBACK_WORKLOADS = [
+    {
+        "x_shape": [1, 7168],
+        "token_topk_to_pos_shape": [1, 8],
+        "pos_to_expert_shape": [16],
+        "dtypes": ["bfloat16"],
+        "label": "kimi-k2-decode",
+    },
+    {
+        "x_shape": [32, 7168],
+        "token_topk_to_pos_shape": [32, 8],
+        "pos_to_expert_shape": [256],
+        "dtypes": ["bfloat16"],
+        "label": "kimi-k2-small",
+    },
+    {
+        "x_shape": [512, 7168],
+        "token_topk_to_pos_shape": [512, 8],
+        "pos_to_expert_shape": [4096],
+        "dtypes": ["bfloat16"],
+        "label": "kimi-k2-medium",
+    },
+    {
+        "x_shape": [4096, 7168],
+        "token_topk_to_pos_shape": [4096, 8],
+        "pos_to_expert_shape": [32768],
+        "dtypes": ["bfloat16"],
+        "label": "kimi-k2-prefill",
+    },
+    {
+        "x_shape": [1, 3072],
+        "token_topk_to_pos_shape": [1, 8],
+        "pos_to_expert_shape": [16],
+        "dtypes": ["bfloat16"],
+        "label": "qwen3-30b-decode",
+    },
+    {
+        "x_shape": [32, 3072],
+        "token_topk_to_pos_shape": [32, 8],
+        "pos_to_expert_shape": [256],
+        "dtypes": ["bfloat16"],
+        "label": "qwen3-30b-small",
+    },
+    {
+        "x_shape": [512, 3072],
+        "token_topk_to_pos_shape": [512, 8],
+        "pos_to_expert_shape": [4096],
+        "dtypes": ["bfloat16"],
+        "label": "qwen3-30b-medium",
+    },
+    {
+        "x_shape": [4096, 3072],
+        "token_topk_to_pos_shape": [4096, 8],
+        "pos_to_expert_shape": [32768],
+        "dtypes": ["bfloat16"],
+        "label": "qwen3-30b-prefill",
+    },
+]
+
+_SF_FALLBACK_WORKLOADS = [
+    {
+        "x_shape": [1, 7168],
+        "x_sf_shape": [1, 56],
+        "token_topk_to_pos_shape": [1, 8],
+        "pos_to_expert_shape": [16],
+        "num_per_channels": 128,
+        "use_tma_aligned_col_major_sf": False,
+        "dtypes": ["float8_e4m3fn"],
+        "label": "kimi-k2-decode-sf-row-major",
+    },
+    {
+        "x_shape": [512, 7168],
+        "x_sf_shape": [512, 56],
+        "token_topk_to_pos_shape": [512, 8],
+        "pos_to_expert_shape": [4096],
+        "num_per_channels": 128,
+        "use_tma_aligned_col_major_sf": False,
+        "dtypes": ["float8_e4m3fn"],
+        "label": "kimi-k2-medium-sf-row-major",
+    },
+    {
+        "x_shape": [1, 7168],
+        "x_sf_shape": [1, 14],
+        "token_topk_to_pos_shape": [1, 8],
+        "pos_to_expert_shape": [16],
+        "num_per_channels": 128,
+        "use_tma_aligned_col_major_sf": True,
+        "dtypes": ["float8_e4m3fn"],
+        "label": "kimi-k2-decode-sf-packed",
+    },
+    {
+        "x_shape": [512, 7168],
+        "x_sf_shape": [512, 14],
+        "token_topk_to_pos_shape": [512, 8],
+        "pos_to_expert_shape": [4096],
+        "num_per_channels": 128,
+        "use_tma_aligned_col_major_sf": True,
+        "dtypes": ["float8_e4m3fn"],
+        "label": "kimi-k2-medium-sf-packed",
+    },
+    {
+        "x_shape": [32, 3072],
+        "x_sf_shape": [32, 24],
+        "token_topk_to_pos_shape": [32, 8],
+        "pos_to_expert_shape": [256],
+        "num_per_channels": 128,
+        "use_tma_aligned_col_major_sf": False,
+        "dtypes": ["float8_e4m3fn"],
+        "label": "qwen3-30b-small-sf-row-major",
+    },
+    {
+        "x_shape": [32, 3072],
+        "x_sf_shape": [32, 6],
+        "token_topk_to_pos_shape": [32, 8],
+        "pos_to_expert_shape": [256],
+        "num_per_channels": 128,
+        "use_tma_aligned_col_major_sf": True,
+        "dtypes": ["float8_e4m3fn"],
+        "label": "qwen3-30b-small-sf-packed",
+    },
+]
+
+
+def _ceil_div(x: int, y: int) -> int:
+    return -(-x // y)
+
+
+def _load_workloads_or_fallback(
+    op_name: str,
+    fallback: list[dict[str, Any]],
+) -> list[dict[str, Any]]:
+    try:
+        workloads = load_workloads(op_name)
+    except KeyError:
+        return fallback
+    return workloads or fallback
+
+
+def _make_routing(
+    total_tokens: int,
+    top_k: int,
+    num_expanded_tokens: int,
+    *,
+    device: str = "cuda",
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build deterministic routing with extra capacity marked unassigned."""
+    token_topk_to_pos = torch.full(
+        (total_tokens, top_k), -1, dtype=torch.int32, device=device
+    )
+    pos_to_expert = torch.full(
+        (num_expanded_tokens,), -1, dtype=torch.int32, device=device
+    )
+
+    slot = 0
+    for t in range(total_tokens):
+        for k in range(top_k):
+            if slot >= num_expanded_tokens:
+                return token_topk_to_pos, pos_to_expert
+            token_topk_to_pos[t, k] = slot
+            pos_to_expert[slot] = k
+            slot += 1
+    return token_topk_to_pos, pos_to_expert
+
+
+def _infer_sf_dtype(
+    hidden_size: int,
+    hidden_sf: int,
+    num_per_channels: int,
+) -> torch.dtype:
+    """Infer float32 vs packed UE8M0 int32 from the manifest SF width."""
+    sf_blocks = _ceil_div(hidden_size, num_per_channels)
+    if hidden_sf == _ceil_div(sf_blocks, _UE8M0_PER_INT32):
+        return torch.int32
+    if hidden_sf == sf_blocks:
+        return torch.float32
+    raise ValueError(
+        f"Cannot infer scale-factor dtype for hidden_size={hidden_size}, "
+        f"hidden_sf={hidden_sf}, num_per_channels={num_per_channels}"
+    )
+
+
+def _torch_expand_to_fused(
+    x: torch.Tensor,
+    token_topk_to_pos: torch.Tensor,
+    pos_to_expert: torch.Tensor,
+) -> torch.Tensor:
+    expanded_x = torch.zeros(
+        (pos_to_expert.shape[0], x.shape[1]), dtype=x.dtype, device=x.device
+    )
+
+    flat_pos = token_topk_to_pos.reshape(-1).to(torch.int64)
+    valid = flat_pos >= 0
+    token_ids = torch.arange(x.shape[0], device=x.device).repeat_interleave(
+        token_topk_to_pos.shape[1]
+    )
+    expanded_x[flat_pos[valid]] = x[token_ids[valid]]
+    return expanded_x
+
+
+def _torch_expand_to_fused_with_sf(
+    x: torch.Tensor,
+    x_sf: torch.Tensor,
+    token_topk_to_pos: torch.Tensor,
+    pos_to_expert: torch.Tensor,
+    *,
+    use_tma_aligned_col_major_sf: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    expanded_x = _torch_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+    if use_tma_aligned_col_major_sf:
+        num_expanded_sf_tokens = _ceil_div(pos_to_expert.shape[0], 4) * 4
+        expanded_x_sf = torch.zeros(
+            (x_sf.shape[1], num_expanded_sf_tokens),
+            dtype=x_sf.dtype,
+            device=x_sf.device,
+        )[:, : pos_to_expert.shape[0]].T
+    else:
+        expanded_x_sf = torch.zeros(
+            (pos_to_expert.shape[0], x_sf.shape[1]),
+            dtype=x_sf.dtype,
+            device=x_sf.device,
+        )
+
+    flat_pos = token_topk_to_pos.reshape(-1).to(torch.int64)
+    valid = flat_pos >= 0
+    token_ids = torch.arange(x.shape[0], device=x.device).repeat_interleave(
+        token_topk_to_pos.shape[1]
+    )
+    expanded_x_sf[flat_pos[valid]] = x_sf[token_ids[valid]]
+    return expanded_x, expanded_x_sf
+
+
+class MoeExpandToFusedBench(WorkloadBase):
+    def __init__(
+        self,
+        total_tokens: int,
+        top_k: int,
+        hidden_size: int,
+        num_expanded_tokens: int,
+        dtype: torch.dtype,
+    ) -> None:
+        self.total_tokens = total_tokens
+        self.top_k = top_k
+        self.hidden_size = hidden_size
+        self.num_expanded_tokens = num_expanded_tokens
+        self.dtype = dtype
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        return (self.total_tokens, self.hidden_size)
+
+    def gen_inputs(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        torch.manual_seed(42)
+        x = torch.randn(
+            self.total_tokens,
+            self.hidden_size,
+            dtype=self.dtype,
+            device="cuda",
+        )
+        token_topk_to_pos, pos_to_expert = _make_routing(
+            self.total_tokens,
+            self.top_k,
+            self.num_expanded_tokens,
+        )
+        return x, token_topk_to_pos, pos_to_expert
+
+
+class MoeExpandToFusedWithSFBench(WorkloadBase):
+    def __init__(
+        self,
+        total_tokens: int,
+        top_k: int,
+        hidden_size: int,
+        hidden_sf: int,
+        num_expanded_tokens: int,
+        num_per_channels: int,
+        dtype: torch.dtype,
+        sf_dtype: torch.dtype,
+    ) -> None:
+        self.total_tokens = total_tokens
+        self.top_k = top_k
+        self.hidden_size = hidden_size
+        self.hidden_sf = hidden_sf
+        self.num_expanded_tokens = num_expanded_tokens
+        self.num_per_channels = num_per_channels
+        self.dtype = dtype
+        self.sf_dtype = sf_dtype
+
+    @property
+    def shape(self) -> tuple[int, int]:
+        return (self.total_tokens, self.hidden_size)
+
+    def gen_inputs(
+        self,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        torch.manual_seed(42)
+        x = torch.randn(
+            self.total_tokens,
+            self.hidden_size,
+            dtype=torch.float32,
+            device="cuda",
+        ).to(self.dtype)
+        if self.sf_dtype == torch.float32:
+            x_sf = torch.randn(
+                self.total_tokens,
+                self.hidden_sf,
+                dtype=self.sf_dtype,
+                device="cuda",
+            )
+        else:
+            x_sf = torch.randint(
+                0,
+                0x7F7F7F7F,
+                (self.total_tokens, self.hidden_sf),
+                dtype=self.sf_dtype,
+                device="cuda",
+            )
+        token_topk_to_pos, pos_to_expert = _make_routing(
+            self.total_tokens,
+            self.top_k,
+            self.num_expanded_tokens,
+        )
+        return x, x_sf, token_topk_to_pos, pos_to_expert
+
+
+def _expand_manifest_params() -> list[Any]:
+    params = []
+    for workload in _load_workloads_or_fallback(
+        _EXPAND_OP_NAME, _BASE_FALLBACK_WORKLOADS
+    ):
+        label = workload.get("label", "unlabeled")
+        total_tokens, hidden_size = workload["x_shape"]
+        topk_tokens, top_k = workload["token_topk_to_pos_shape"]
+        assert topk_tokens == total_tokens
+        (num_expanded_tokens,) = workload["pos_to_expert_shape"]
+        for dtype_str in workload["dtypes"]:
+            dtype = getattr(torch, dtype_str)
+            params.append(
+                pytest.param(
+                    total_tokens,
+                    top_k,
+                    hidden_size,
+                    num_expanded_tokens,
+                    dtype,
+                    id=f"{label}-{dtype_str}",
+                )
+            )
+    return params
+
+
+def _expand_with_sf_manifest_params() -> list[Any]:
+    if _FP8_E4M3 is None:
+        return []
+    params = []
+    for workload in _load_workloads_or_fallback(
+        _EXPAND_WITH_SF_OP_NAME, _SF_FALLBACK_WORKLOADS
+    ):
+        label = workload.get("label", "unlabeled")
+        total_tokens, hidden_size = workload["x_shape"]
+        sf_tokens, hidden_sf = workload["x_sf_shape"]
+        topk_tokens, top_k = workload["token_topk_to_pos_shape"]
+        assert sf_tokens == total_tokens
+        assert topk_tokens == total_tokens
+        (num_expanded_tokens,) = workload["pos_to_expert_shape"]
+        num_per_channels = workload["num_per_channels"]
+        use_tma_aligned_col_major_sf = workload["use_tma_aligned_col_major_sf"]
+        sf_dtype = _infer_sf_dtype(hidden_size, hidden_sf, num_per_channels)
+        for dtype_str in workload["dtypes"]:
+            dtype = getattr(torch, dtype_str)
+            params.append(
+                pytest.param(
+                    total_tokens,
+                    top_k,
+                    hidden_size,
+                    hidden_sf,
+                    num_expanded_tokens,
+                    num_per_channels,
+                    use_tma_aligned_col_major_sf,
+                    dtype,
+                    sf_dtype,
+                    id=f"{label}-{dtype_str}-{str(sf_dtype).removeprefix('torch.')}",
+                )
+            )
+    return params
+
+
+@pytest.mark.parametrize(
+    "total_tokens, top_k, hidden_size, num_expanded_tokens, dtype",
+    _expand_manifest_params(),
+)
+def test_moe_expand_to_fused_bench(
+    total_tokens: int,
+    top_k: int,
+    hidden_size: int,
+    num_expanded_tokens: int,
+    dtype: torch.dtype,
+) -> None:
+    test = MoeExpandToFusedBench(
+        total_tokens,
+        top_k,
+        hidden_size,
+        num_expanded_tokens,
+        dtype,
+    )
+    inputs = test.gen_inputs()
+
+    op = MoeExpandToFusedFwdOp(dtype=dtype)
+    bm = ManifestBenchmark(_EXPAND_OP_NAME, op, test)
+    op(*inputs)  # warmup / JIT compile
+    torch.cuda.synchronize()
+
+    result = bm.profile(op, *inputs)
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
+
+    _torch_expand_to_fused(*inputs)  # warmup
+    torch.cuda.synchronize()
+
+    result_torch = bm.profile(_torch_expand_to_fused, *inputs)
+    BenchmarkReport.record(op, locals(), result_torch, tag="torch-ref")
+
+
+@pytest.mark.skipif(_FP8_E4M3 is None, reason="torch fp8 is unavailable")
+@pytest.mark.parametrize(
+    "total_tokens, top_k, hidden_size, hidden_sf, num_expanded_tokens, "
+    "num_per_channels, use_tma_aligned_col_major_sf, dtype, sf_dtype",
+    _expand_with_sf_manifest_params(),
+)
+def test_moe_expand_to_fused_with_sf_bench(
+    total_tokens: int,
+    top_k: int,
+    hidden_size: int,
+    hidden_sf: int,
+    num_expanded_tokens: int,
+    num_per_channels: int,
+    use_tma_aligned_col_major_sf: bool,
+    dtype: torch.dtype,
+    sf_dtype: torch.dtype,
+) -> None:
+    test = MoeExpandToFusedWithSFBench(
+        total_tokens,
+        top_k,
+        hidden_size,
+        hidden_sf,
+        num_expanded_tokens,
+        num_per_channels,
+        dtype,
+        sf_dtype,
+    )
+    inputs = test.gen_inputs()
+
+    op = MoeExpandToFusedWithSFFwdOp(
+        num_per_channels=num_per_channels,
+        use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf,
+        dtype=dtype,
+        sf_dtype=sf_dtype,
+    )
+    bm = ManifestBenchmark(_EXPAND_WITH_SF_OP_NAME, op, test)
+    op(*inputs)  # warmup / JIT compile
+    torch.cuda.synchronize()
+
+    result = bm.profile(op, *inputs)
+    BenchmarkReport.record(op, locals(), result, tag="tileops")
+
+    def torch_ref_with_sf(*args):
+        return _torch_expand_to_fused_with_sf(
+            *args,
+            use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf,
+        )
+
+    torch_ref_with_sf(*inputs)  # warmup
+    torch.cuda.synchronize()
+
+    result_torch = bm.profile(torch_ref_with_sf, *inputs)
+    BenchmarkReport.record(op, locals(), result_torch, tag="torch-ref")
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-vvs"])

tests/ops/test_moe_expand_to_fused.py

@@ -0,0 +1,999 @@
+"""Reference and optional kernel tests for moe_expand_to_fused.
+
+The operator contract is a scatter-copy into the fused expert layout:
+  - token_topk_to_pos[t, k] >= 0 copies x[t] to output[token_topk_to_pos[t, k]]
+  - token_topk_to_pos[t, k] < 0 is skipped
+  - pos_to_expert[pos] < 0 marks padding; padding output rows stay zero
+  - no weights, sums, or reductions are applied
+
+The cases mirror the TileKernels-Metax source implementation at
+79461d72d6f97f91b89403705a841c9d4c1eddd6:
+  - base expand_to_fused for bf16/fp16 activation rows
+  - expand_to_fused_with_sf for fp8 activations plus row-major fp32 or packed
+    int32 scale factors
+  - TMA-aligned scale factors represented publicly as shape [X, S] with a
+    column-major stride
+
+Most tests exercise the pure-PyTorch reference so they can run without a C500
+runtime. CUDA smoke tests import TileLang-backed modules lazily; they compile
+and execute the new kernel only when CUDA and the MACA TileLang build are
+available.
+"""
+
+import os
+
+import pytest
+import torch
+
+
+os.environ.setdefault("TILELANG_PRINT_ON_COMPILATION", "0")
+
+
+_FP8_E4M3 = getattr(torch, "float8_e4m3fn", None)
+_DOCUMENTED_HIDDEN_SIZES = [576, 2048, 2560, 3072, 4096, 6144, 7168]
+_DOCUMENTED_TOPK = [2, 6, 8, 9]
+_DOCUMENTED_NUM_PER_CHANNELS = [32, 128]
+
+_BASE_MANIFEST_WORKLOADS = [
+    pytest.param(
+        "kimi-k2-decode",
+        1,
+        7168,
+        8,
+        16,
+        marks=pytest.mark.smoke,
+        id="kimi-k2-decode",
+    ),
+    pytest.param(
+        "kimi-k2-small",
+        32,
+        7168,
+        8,
+        256,
+        marks=pytest.mark.full,
+        id="kimi-k2-small",
+    ),
+    pytest.param(
+        "kimi-k2-medium",
+        512,
+        7168,
+        8,
+        4096,
+        marks=pytest.mark.full,
+        id="kimi-k2-medium",
+    ),
+    pytest.param(
+        "kimi-k2-prefill",
+        4096,
+        7168,
+        8,
+        32768,
+        marks=pytest.mark.full,
+        id="kimi-k2-prefill",
+    ),
+    pytest.param(
+        "qwen3-30b-decode",
+        1,
+        3072,
+        8,
+        16,
+        marks=pytest.mark.full,
+        id="qwen3-30b-decode",
+    ),
+    pytest.param(
+        "qwen3-30b-small",
+        32,
+        3072,
+        8,
+        256,
+        marks=pytest.mark.full,
+        id="qwen3-30b-small",
+    ),
+    pytest.param(
+        "qwen3-30b-medium",
+        512,
+        3072,
+        8,
+        4096,
+        marks=pytest.mark.full,
+        id="qwen3-30b-medium",
+    ),
+    pytest.param(
+        "qwen3-30b-prefill",
+        4096,
+        3072,
+        8,
+        32768,
+        marks=pytest.mark.full,
+        id="qwen3-30b-prefill",
+    ),
+]
+
+_SF_MANIFEST_WORKLOADS = [
+    pytest.param(
+        "kimi-k2-decode-sf-row-major",
+        1,
+        7168,
+        8,
+        16,
+        128,
+        False,
+        torch.float32,
+        False,
+        marks=pytest.mark.smoke,
+        id="kimi-k2-decode-row-major",
+    ),
+    pytest.param(
+        "kimi-k2-medium-sf-row-major",
+        512,
+        7168,
+        8,
+        4096,
+        128,
+        False,
+        torch.float32,
+        False,
+        marks=pytest.mark.full,
+        id="kimi-k2-medium-row-major",
+    ),
+    pytest.param(
+        "kimi-k2-decode-sf-packed",
+        1,
+        7168,
+        8,
+        16,
+        128,
+        True,
+        torch.int32,
+        True,
+        marks=pytest.mark.full,
+        id="kimi-k2-decode-packed",
+    ),
+    pytest.param(
+        "kimi-k2-medium-sf-packed",
+        512,
+        7168,
+        8,
+        4096,
+        128,
+        True,
+        torch.int32,
+        True,
+        marks=pytest.mark.full,
+        id="kimi-k2-medium-packed",
+    ),
+    pytest.param(
+        "qwen3-30b-small-sf-row-major",
+        32,
+        3072,
+        8,
+        256,
+        128,
+        False,
+        torch.float32,
+        False,
+        marks=pytest.mark.full,
+        id="qwen3-30b-small-row-major",
+    ),
+    pytest.param(
+        "qwen3-30b-small-sf-packed",
+        32,
+        3072,
+        8,
+        256,
+        128,
+        True,
+        torch.int32,
+        True,
+        marks=pytest.mark.full,
+        id="qwen3-30b-small-packed",
+    ),
+]
+
+def _ceil_div(x: int, y: int) -> int:
+    return (x + y - 1) // y
+
+
+def _align(x: int, alignment: int) -> int:
+    return _ceil_div(x, alignment) * alignment
+
+
+def _is_float8_dtype(dtype: torch.dtype) -> bool:
+    fp8_dtypes = {
+        dt
+        for dt in (
+            getattr(torch, "float8_e4m3fn", None),
+            getattr(torch, "float8_e5m2", None),
+        )
+        if dt is not None
+    }
+    return dtype in fp8_dtypes
+
+
+def _sf_channels(hidden: int, num_per_channels: int, *, packed: bool) -> int:
+    channels = _ceil_div(hidden, num_per_channels)
+    return _ceil_div(channels, 4) if packed else channels
+
+
+def _assert_exact(actual: torch.Tensor, expected: torch.Tensor) -> None:
+    assert actual.shape == expected.shape
+    assert actual.dtype == expected.dtype
+    if _is_float8_dtype(actual.dtype):
+        assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
+    elif actual.is_floating_point():
+        torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    else:
+        assert torch.equal(actual, expected)
+
+
+def _make_x(num_tokens: int, hidden: int, dtype: torch.dtype) -> torch.Tensor:
+    values = torch.arange(num_tokens * hidden, dtype=torch.float32).view(num_tokens, hidden)
+    if _is_float8_dtype(dtype):
+        return ((values.remainder(32) - 16) / 8).to(dtype)
+    return (values / 100).to(dtype)
+
+
+def _make_scale_factors(num_tokens: int, channels: int, dtype: torch.dtype) -> torch.Tensor:
+    values = torch.arange(num_tokens * channels, dtype=torch.float32).view(num_tokens, channels)
+    if dtype == torch.int32:
+        return (values.to(torch.int32) + 1) * 17
+    return (values / 100 + 1).to(dtype)
+
+
+def _source_rows(token_topk_to_pos: torch.Tensor) -> torch.Tensor:
+    return torch.arange(
+        token_topk_to_pos.shape[0],
+        dtype=torch.int64,
+        device=token_topk_to_pos.device,
+    ).view(-1, 1).expand_as(token_topk_to_pos)
+
+
+def _validate_expand_inputs(
+    x: torch.Tensor,
+    token_topk_to_pos: torch.Tensor,
+    pos_to_expert: torch.Tensor,
+) -> None:
+    if x.ndim != 2:
+        raise ValueError(f"x must be 2-D, got shape={tuple(x.shape)}")
+    if token_topk_to_pos.ndim != 2:
+        raise ValueError(
+            f"token_topk_to_pos must be 2-D, got shape={tuple(token_topk_to_pos.shape)}"
+        )
+    if token_topk_to_pos.shape[0] != x.shape[0]:
+        raise ValueError(
+            "token_topk_to_pos.shape[0] must match x.shape[0], "
+            f"got {token_topk_to_pos.shape[0]} and {x.shape[0]}"
+        )
+    if pos_to_expert.ndim != 1:
+        raise ValueError(f"pos_to_expert must be 1-D, got shape={tuple(pos_to_expert.shape)}")
+    if token_topk_to_pos.dtype != torch.int32:
+        raise ValueError(f"token_topk_to_pos must be int32, got {token_topk_to_pos.dtype}")
+    if pos_to_expert.dtype != torch.int32:
+        raise ValueError(f"pos_to_expert must be int32, got {pos_to_expert.dtype}")
+
+    # These malformed-mapping checks keep synthetic test inputs unambiguous.
+    # They are reference sanity checks, not a requirement that the eventual
+    # kernel performs runtime validation for undefined inputs.
+    valid = token_topk_to_pos >= 0
+    if valid.any():
+        valid_pos = token_topk_to_pos[valid]
+        if valid_pos.max().item() >= pos_to_expert.numel():
+            raise ValueError("token_topk_to_pos contains a position outside pos_to_expert")
+        if torch.unique(valid_pos).numel() != valid_pos.numel():
+            raise ValueError("token_topk_to_pos contains duplicate valid positions")
+        if (pos_to_expert[valid_pos.long()] < 0).any():
+            raise ValueError("token_topk_to_pos points to a padding position")
+
+
+def _ref_moe_expand_to_fused(
+    x: torch.Tensor,
+    token_topk_to_pos: torch.Tensor,
+    pos_to_expert: torch.Tensor,
+) -> torch.Tensor:
+    """Pure-PyTorch reference for MoeExpandToFusedFwdOp."""
+    _validate_expand_inputs(x, token_topk_to_pos, pos_to_expert)
+
+    out = torch.zeros(
+        (pos_to_expert.numel(), x.shape[1]),
+        dtype=x.dtype,
+        device=x.device,
+    )
+    valid = token_topk_to_pos >= 0
+    if valid.any():
+        src_rows = _source_rows(token_topk_to_pos)
+        dst_rows = token_topk_to_pos[valid].long()
+        src_rows = src_rows[valid].long()
+        if _is_float8_dtype(x.dtype):
+            # PyTorch's CPU float8 tensors do not implement advanced indexing
+            # (`index_cpu`), so copy the exact encoded bytes for reference tests.
+            out.view(torch.uint8)[dst_rows] = x.view(torch.uint8)[src_rows]
+        else:
+            out[dst_rows] = x[src_rows]
+    return out
+
+
+def _ref_moe_expand_to_fused_with_sf(
+    x: torch.Tensor,
+    x_sf: torch.Tensor,
+    num_per_channels: int,
+    token_topk_to_pos: torch.Tensor,
+    pos_to_expert: torch.Tensor,
+    *,
+    use_tma_aligned_col_major_sf: bool = False,
+    use_packed_ue8m0: bool = False,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pure-PyTorch reference for MoeExpandToFusedWithSFFwdOp."""
+    if num_per_channels not in _DOCUMENTED_NUM_PER_CHANNELS:
+        raise ValueError(f"unsupported num_per_channels={num_per_channels}")
+    if use_packed_ue8m0 and not use_tma_aligned_col_major_sf:
+        raise ValueError("packed UE8M0 scale factors require TMA column-major layout")
+    if x_sf.ndim != 2:
+        raise ValueError(f"x_sf must be 2-D, got shape={tuple(x_sf.shape)}")
+
+    expected_s = _sf_channels(
+        x.shape[1],
+        num_per_channels,
+        packed=use_packed_ue8m0,
+    )
+    if x_sf.shape != (x.shape[0], expected_s):
+        raise ValueError(
+            "x_sf must have shape (T, S), where S follows the manifest scale-factor "
+            f"rule; got {tuple(x_sf.shape)} for x.shape={tuple(x.shape)}"
+        )
+    if use_packed_ue8m0 and x_sf.dtype != torch.int32:
+        raise ValueError("packed UE8M0 scale factors are represented as int32")
+
+    out = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+    if use_tma_aligned_col_major_sf:
+        aligned_tokens = _align(pos_to_expert.numel(), 4)
+        out_sf = torch.zeros(
+            (x_sf.shape[1], aligned_tokens),
+            dtype=x_sf.dtype,
+            device=x_sf.device,
+        )
+        # Intentionally return a transposed view to emulate the TMA-aligned
+        # column-major scale-factor layout while preserving public shape [X, S].
+        out_sf = out_sf[:, : pos_to_expert.numel()].T
+    else:
+        out_sf = torch.zeros(
+            (pos_to_expert.numel(), x_sf.shape[1]),
+            dtype=x_sf.dtype,
+            device=x_sf.device,
+        )
+
+    valid = token_topk_to_pos >= 0
+    if valid.any():
+        src_rows = _source_rows(token_topk_to_pos)
+        out_sf[token_topk_to_pos[valid].long()] = x_sf[src_rows[valid].long()]
+    return out, out_sf
+
+
+def _make_expert_major_mapping(
+    topk_ids: torch.Tensor,
+    *,
+    num_experts: int,
+    block_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build deterministic expert-major padded mapping for test inputs.
+
+    Only the observable mapping contract required by expand_to_fused tests is
+    reproduced here. This is not intended to match every implementation detail
+    of moe_get_fused_mapping.
+    """
+    if topk_ids.ndim != 2:
+        raise ValueError(f"topk_ids must be 2-D, got shape={tuple(topk_ids.shape)}")
+
+    counts = [0] * num_experts
+    for eid in topk_ids.flatten().tolist():
+        if eid < 0:
+            continue
+        if eid >= num_experts:
+            raise ValueError(f"expert id {eid} is outside [0, {num_experts})")
+        counts[eid] += 1
+
+    offsets = [0] * (num_experts + 1)
+    for eid, count in enumerate(counts):
+        offsets[eid + 1] = offsets[eid] + _align(count, block_size)
+
+    token_topk_to_pos = torch.full_like(topk_ids, -1, dtype=torch.int32)
+    pos_to_expert = torch.full((offsets[-1],), -1, dtype=torch.int32, device=topk_ids.device)
+
+    write_ptr = offsets[:-1].copy()
+    for token in range(topk_ids.shape[0]):
+        for kth in range(topk_ids.shape[1]):
+            eid = int(topk_ids[token, kth].item())
+            if eid < 0:
+                continue
+            pos = write_ptr[eid]
+            token_topk_to_pos[token, kth] = pos
+            pos_to_expert[pos] = eid
+            write_ptr[eid] += 1
+
+    return token_topk_to_pos, pos_to_expert
+
+
+def _make_manifest_mapping(
+    num_tokens: int,
+    topk: int,
+    num_expanded_tokens: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Build a shape-driven mapping matching manifest workload dimensions."""
+    token_topk_to_pos = torch.full((num_tokens, topk), -1, dtype=torch.int32)
+    pos_to_expert = torch.full((num_expanded_tokens,), -1, dtype=torch.int32)
+
+    num_valid = min(num_tokens * topk, num_expanded_tokens)
+    if num_valid == 0:
+        return token_topk_to_pos, pos_to_expert
+
+    flat = token_topk_to_pos.view(-1)
+    flat[:num_valid] = torch.arange(num_valid, dtype=torch.int32)
+    pos_to_expert[:num_valid] = torch.arange(num_valid, dtype=torch.int32) % 16
+    return token_topk_to_pos, pos_to_expert
+
+
+def _make_topk_ids(
+    num_tokens: int,
+    topk: int,
+    num_experts: int,
+    *,
+    invalid_every: int | None = None,
+) -> torch.Tensor:
+    token = torch.arange(num_tokens, dtype=torch.int32).view(-1, 1)
+    kth = torch.arange(topk, dtype=torch.int32).view(1, -1)
+    topk_ids = (token * 3 + kth * 5) % num_experts
+    if invalid_every is not None and num_tokens > 0:
+        topk_ids[::invalid_every, -1] = -1
+    return topk_ids
+
+
+def _assert_expand_obeys_mapping(
+    out: torch.Tensor,
+    x: torch.Tensor,
+    token_topk_to_pos: torch.Tensor,
+    pos_to_expert: torch.Tensor,
+) -> None:
+    assert out.shape == (pos_to_expert.numel(), x.shape[1])
+    assert out.dtype == x.dtype
+
+    valid = token_topk_to_pos >= 0
+    targeted = torch.zeros(out.shape[0], dtype=torch.bool, device=out.device)
+    if valid.any():
+        target_pos = token_topk_to_pos[valid].long()
+        src_rows = _source_rows(token_topk_to_pos)[valid].long()
+        targeted[target_pos] = True
+        assert torch.unique(target_pos).numel() == target_pos.numel()
+        if _is_float8_dtype(out.dtype):
+            assert torch.equal(
+                out.view(torch.uint8)[target_pos],
+                x.view(torch.uint8)[src_rows],
+            )
+        else:
+            _assert_exact(out[target_pos], x[src_rows])
+    if _is_float8_dtype(out.dtype):
+        out_bytes = out.view(torch.uint8)
+        assert torch.equal(out_bytes[~targeted], torch.zeros_like(out_bytes[~targeted]))
+    else:
+        assert torch.equal(out[~targeted], torch.zeros_like(out[~targeted]))
+
+
+def _assert_sf_obeys_mapping(
+    out_sf: torch.Tensor,
+    x_sf: torch.Tensor,
+    token_topk_to_pos: torch.Tensor,
+    pos_to_expert: torch.Tensor,
+) -> None:
+    assert out_sf.shape == (pos_to_expert.numel(), x_sf.shape[1])
+    assert out_sf.dtype == x_sf.dtype
+
+    valid = token_topk_to_pos >= 0
+    targeted = torch.zeros(out_sf.shape[0], dtype=torch.bool, device=out_sf.device)
+    if valid.any():
+        target_pos = token_topk_to_pos[valid].long()
+        src_rows = _source_rows(token_topk_to_pos)
+        targeted[target_pos] = True
+        assert torch.unique(target_pos).numel() == target_pos.numel()
+        _assert_exact(out_sf[target_pos], x_sf[src_rows[valid].long()])
+    assert torch.equal(out_sf[~targeted], torch.zeros_like(out_sf[~targeted]))
+
+
+def _run_reference_case(
+    *,
+    num_tokens: int,
+    hidden: int,
+    topk: int,
+    num_experts: int,
+    block_size: int,
+    dtype: torch.dtype = torch.bfloat16,
+    invalid_every: int | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    topk_ids = _make_topk_ids(num_tokens, topk, num_experts, invalid_every=invalid_every)
+    token_topk_to_pos, pos_to_expert = _make_expert_major_mapping(
+        topk_ids,
+        num_experts=num_experts,
+        block_size=block_size,
+    )
+    x = _make_x(num_tokens, hidden, dtype)
+    out = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+    assert out.shape == (pos_to_expert.numel(), hidden)
+    assert out.dtype == dtype
+    _assert_expand_obeys_mapping(out, x, token_topk_to_pos, pos_to_expert)
+    return out, x, token_topk_to_pos, pos_to_expert
+
+
+@pytest.mark.smoke
+def test_moe_expand_to_fused_manual_mapping() -> None:
+    x = torch.tensor(
+        [
+            [1.0, 2.0, 3.0],
+            [10.0, 20.0, 30.0],
+            [100.0, 200.0, 300.0],
+        ],
+        dtype=torch.bfloat16,
+    )
+    token_topk_to_pos = torch.tensor(
+        [
+            [0, 3],
+            [1, 4],
+            [2, -1],
+        ],
+        dtype=torch.int32,
+    )
+    pos_to_expert = torch.tensor([0, 0, 0, 1, 1, -1, -1], dtype=torch.int32)
+
+    out = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+    expected = torch.tensor(
+        [
+            [1.0, 2.0, 3.0],
+            [10.0, 20.0, 30.0],
+            [100.0, 200.0, 300.0],
+            [1.0, 2.0, 3.0],
+            [10.0, 20.0, 30.0],
+            [0.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0],
+        ],
+        dtype=torch.bfloat16,
+    )
+    _assert_exact(out, expected)
+
+
+@pytest.mark.smoke
+def test_moe_expand_to_fused_invalid_topk_and_padding() -> None:
+    _, x, token_topk_to_pos, pos_to_expert = _run_reference_case(
+        num_tokens=8,
+        hidden=16,
+        topk=4,
+        num_experts=4,
+        block_size=8,
+        dtype=torch.float16,
+        invalid_every=2,
+    )
+    out = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+    assert (token_topk_to_pos < 0).any()
+    assert (pos_to_expert < 0).any()
+    assert torch.equal(out[pos_to_expert < 0], torch.zeros_like(out[pos_to_expert < 0]))
+
+
+@pytest.mark.smoke
+def test_moe_expand_to_fused_does_not_weight_or_reduce() -> None:
+    x = torch.tensor([[2.0, 4.0], [8.0, 16.0]], dtype=torch.float16)
+    token_topk_to_pos = torch.tensor([[0, 2], [1, 3]], dtype=torch.int32)
+    pos_to_expert = torch.tensor([0, 0, 1, 1], dtype=torch.int32)
+
+    out = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+    _assert_exact(out[0], x[0])
+    _assert_exact(out[2], x[0])
+    assert not torch.equal(out[0], x[0] * 2)
+
+
+@pytest.mark.smoke
+def test_moe_expand_to_fused_random_mapping() -> None:
+    torch.manual_seed(0)
+    num_tokens, topk, num_experts, hidden = 13, 6, 8, 37
+    topk_ids = torch.randint(0, num_experts, (num_tokens, topk), dtype=torch.int32)
+    topk_ids[::4, -1] = -1
+    token_topk_to_pos, pos_to_expert = _make_expert_major_mapping(
+        topk_ids,
+        num_experts=num_experts,
+        block_size=16,
+    )
+    x = _make_x(num_tokens, hidden, torch.bfloat16)
+
+    out = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+    _assert_expand_obeys_mapping(out, x, token_topk_to_pos, pos_to_expert)
+
+
+@pytest.mark.parametrize(
+    "num_expanded_tokens",
+    [
+        pytest.param(0, marks=pytest.mark.smoke, id="no-expanded-slots"),
+        pytest.param(4, marks=pytest.mark.full, id="padding-only-slots"),
+    ],
+)
+def test_moe_expand_to_fused_empty_tokens_reference_only(num_expanded_tokens: int) -> None:
+    x = torch.empty((0, 576), dtype=torch.bfloat16)
+    token_topk_to_pos = torch.empty((0, 8), dtype=torch.int32)
+    pos_to_expert = torch.full((num_expanded_tokens,), -1, dtype=torch.int32)
+
+    out = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+    assert out.shape == (num_expanded_tokens, 576)
+    assert out.dtype == torch.bfloat16
+    assert torch.equal(out, torch.zeros_like(out))
+
+
+@pytest.mark.smoke
+def test_moe_expand_to_fused_all_invalid_mapping() -> None:
+    x = _make_x(num_tokens=3, hidden=16, dtype=torch.bfloat16)
+    token_topk_to_pos = torch.full((3, 4), -1, dtype=torch.int32)
+    pos_to_expert = torch.full((8,), -1, dtype=torch.int32)
+
+    out = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+    assert not (token_topk_to_pos >= 0).any()
+    assert torch.equal(out, torch.zeros_like(out))
+
+
+@pytest.mark.parametrize(
+    "hidden",
+    [
+        pytest.param(
+            hidden,
+            marks=pytest.mark.smoke
+            if hidden == _DOCUMENTED_HIDDEN_SIZES[0]
+            else pytest.mark.full,
+            id=f"h{hidden}",
+        )
+        for hidden in _DOCUMENTED_HIDDEN_SIZES
+    ],
+)
+def test_moe_expand_to_fused_documented_hidden_sizes(hidden: int) -> None:
+    _run_reference_case(
+        num_tokens=5,
+        hidden=hidden,
+        topk=2,
+        num_experts=4,
+        block_size=4,
+    )
+
+
+@pytest.mark.parametrize(
+    "topk",
+    [
+        pytest.param(
+            topk,
+            marks=pytest.mark.smoke if topk == _DOCUMENTED_TOPK[0] else pytest.mark.full,
+            id=f"topk{topk}",
+        )
+        for topk in _DOCUMENTED_TOPK
+    ],
+)
+def test_moe_expand_to_fused_documented_topk(topk: int) -> None:
+    _run_reference_case(
+        num_tokens=7,
+        hidden=128,
+        topk=topk,
+        num_experts=12,
+        block_size=8,
+    )
+
+
+@pytest.mark.smoke
+def test_moe_expand_to_fused_manifest_decode_padding_shape() -> None:
+    x = _make_x(num_tokens=1, hidden=64, dtype=torch.bfloat16)
+    token_topk_to_pos, pos_to_expert = _make_manifest_mapping(
+        num_tokens=1,
+        topk=8,
+        num_expanded_tokens=16,
+    )
+
+    out = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+    assert token_topk_to_pos.numel() == 8
+    assert pos_to_expert.numel() == 16
+    assert out.shape == (16, 64)
+    _assert_expand_obeys_mapping(out, x, token_topk_to_pos, pos_to_expert)
+
+
+@pytest.mark.smoke
+def test_moe_expand_to_fused_num_send_tokens_4001_boundary() -> None:
+    _run_reference_case(
+        num_tokens=4001,
+        hidden=1,
+        topk=8,
+        num_experts=16,
+        block_size=64,
+        invalid_every=17,
+    )
+
+
+@pytest.mark.parametrize(
+    "dtype",
+    [
+        pytest.param(torch.bfloat16, marks=pytest.mark.smoke, id="bf16"),
+        pytest.param(torch.float16, marks=pytest.mark.smoke, id="fp16"),
+    ],
+)
+def test_moe_expand_to_fused_supported_dtypes(dtype: torch.dtype) -> None:
+    token_topk_to_pos, pos_to_expert = _make_manifest_mapping(
+        num_tokens=5,
+        topk=2,
+        num_expanded_tokens=10,
+    )
+    x = _make_x(5, 7, dtype)
+
+    out = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+    _assert_expand_obeys_mapping(out, x, token_topk_to_pos, pos_to_expert)
+
+
+@pytest.mark.parametrize(
+    ("label", "num_tokens", "hidden", "topk", "num_expanded_tokens"),
+    _BASE_MANIFEST_WORKLOADS,
+)
+def test_moe_expand_to_fused_manifest_workload_shapes(
+    label: str,
+    num_tokens: int,
+    hidden: int,
+    topk: int,
+    num_expanded_tokens: int,
+) -> None:
+    assert label
+    x = _make_x(num_tokens, hidden, torch.bfloat16)
+    token_topk_to_pos, pos_to_expert = _make_manifest_mapping(
+        num_tokens,
+        topk,
+        num_expanded_tokens,
+    )
+
+    out = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+    assert out.shape == (num_expanded_tokens, hidden)
+    _assert_expand_obeys_mapping(out, x, token_topk_to_pos, pos_to_expert)
+
+
+@pytest.mark.skipif(_FP8_E4M3 is None, reason="torch fp8 is unavailable")
+@pytest.mark.parametrize(
+    "num_per_channels",
+    [
+        pytest.param(32, marks=pytest.mark.smoke, id="npc32"),
+        pytest.param(128, marks=pytest.mark.smoke, id="npc128"),
+    ],
+)
+@pytest.mark.parametrize(
+    ("use_tma_aligned_col_major_sf", "sf_dtype", "use_packed_ue8m0"),
+    [
+        pytest.param(False, torch.float32, False, id="row-major-fp32-sf"),
+        pytest.param(True, torch.int32, True, id="tma-col-major-packed-ue8m0"),
+    ],
+)
+def test_moe_expand_to_fused_with_sf(
+    num_per_channels: int,
+    use_tma_aligned_col_major_sf: bool,
+    sf_dtype: torch.dtype,
+    use_packed_ue8m0: bool,
+) -> None:
+    num_tokens, topk, num_expanded_tokens, hidden = 6, 8, 64, 576
+    token_topk_to_pos, pos_to_expert = _make_manifest_mapping(
+        num_tokens,
+        topk,
+        num_expanded_tokens,
+    )
+    x = _make_x(num_tokens, hidden, _FP8_E4M3)
+    channels = _sf_channels(hidden, num_per_channels, packed=use_packed_ue8m0)
+    x_sf = _make_scale_factors(num_tokens, channels, sf_dtype)
+
+    out, out_sf = _ref_moe_expand_to_fused_with_sf(
+        x,
+        x_sf,
+        num_per_channels,
+        token_topk_to_pos,
+        pos_to_expert,
+        use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf,
+        use_packed_ue8m0=use_packed_ue8m0,
+    )
+
+    _assert_expand_obeys_mapping(out, x, token_topk_to_pos, pos_to_expert)
+    _assert_sf_obeys_mapping(out_sf, x_sf, token_topk_to_pos, pos_to_expert)
+    assert torch.equal(out_sf[pos_to_expert < 0], torch.zeros_like(out_sf[pos_to_expert < 0]))
+    if use_tma_aligned_col_major_sf and channels > 1 and pos_to_expert.numel() > 0:
+        assert out_sf.stride() == (1, _align(pos_to_expert.numel(), 4))
+    else:
+        assert out_sf.is_contiguous()
+
+
+@pytest.mark.skipif(_FP8_E4M3 is None, reason="torch fp8 is unavailable")
+@pytest.mark.parametrize(
+    (
+        "label",
+        "num_tokens",
+        "hidden",
+        "topk",
+        "num_expanded_tokens",
+        "num_per_channels",
+        "use_tma_aligned_col_major_sf",
+        "sf_dtype",
+        "use_packed_ue8m0",
+    ),
+    _SF_MANIFEST_WORKLOADS,
+)
+def test_moe_expand_to_fused_with_sf_manifest_workload_shapes(
+    label: str,
+    num_tokens: int,
+    hidden: int,
+    topk: int,
+    num_expanded_tokens: int,
+    num_per_channels: int,
+    use_tma_aligned_col_major_sf: bool,
+    sf_dtype: torch.dtype,
+    use_packed_ue8m0: bool,
+) -> None:
+    assert label
+    x = _make_x(num_tokens, hidden, _FP8_E4M3)
+    channels = _sf_channels(hidden, num_per_channels, packed=use_packed_ue8m0)
+    x_sf = _make_scale_factors(num_tokens, channels, sf_dtype)
+    token_topk_to_pos, pos_to_expert = _make_manifest_mapping(
+        num_tokens,
+        topk,
+        num_expanded_tokens,
+    )
+
+    out, out_sf = _ref_moe_expand_to_fused_with_sf(
+        x,
+        x_sf,
+        num_per_channels,
+        token_topk_to_pos,
+        pos_to_expert,
+        use_tma_aligned_col_major_sf=use_tma_aligned_col_major_sf,
+        use_packed_ue8m0=use_packed_ue8m0,
+    )
+
+    assert out.shape == (num_expanded_tokens, hidden)
+    assert out_sf.shape == (num_expanded_tokens, channels)
+    _assert_expand_obeys_mapping(out, x, token_topk_to_pos, pos_to_expert)
+    _assert_sf_obeys_mapping(out_sf, x_sf, token_topk_to_pos, pos_to_expert)
+
+
+@pytest.mark.skipif(_FP8_E4M3 is None, reason="torch fp8 is unavailable")
+@pytest.mark.smoke
+def test_moe_expand_to_fused_with_sf_reference_rejects_shape_rules() -> None:
+    x = _make_x(num_tokens=2, hidden=576, dtype=_FP8_E4M3)
+    token_topk_to_pos, pos_to_expert = _make_manifest_mapping(
+        num_tokens=2,
+        topk=2,
+        num_expanded_tokens=4,
+    )
+
+    x_sf_row = _make_scale_factors(2, _sf_channels(576, 128, packed=False), torch.float32)
+    with pytest.raises(ValueError, match="unsupported num_per_channels"):
+        _ref_moe_expand_to_fused_with_sf(
+            x,
+            x_sf_row,
+            64,
+            token_topk_to_pos,
+            pos_to_expert,
+        )
+
+    x_sf_packed = _make_scale_factors(2, _sf_channels(576, 128, packed=True), torch.int32)
+    with pytest.raises(ValueError, match="packed UE8M0 scale factors require"):
+        _ref_moe_expand_to_fused_with_sf(
+            x,
+            x_sf_packed,
+            128,
+            token_topk_to_pos,
+            pos_to_expert,
+            use_packed_ue8m0=True,
+        )
+
+    bad_shape = _make_scale_factors(2, x_sf_row.shape[1] + 1, torch.float32)
+    with pytest.raises(ValueError, match="x_sf must have shape"):
+        _ref_moe_expand_to_fused_with_sf(
+            x,
+            bad_shape,
+            128,
+            token_topk_to_pos,
+            pos_to_expert,
+        )
+
+    bad_packed_dtype = _make_scale_factors(
+        2,
+        _sf_channels(576, 128, packed=True),
+        torch.float32,
+    )
+    with pytest.raises(ValueError, match="packed UE8M0 scale factors are represented"):
+        _ref_moe_expand_to_fused_with_sf(
+            x,
+            bad_packed_dtype,
+            128,
+            token_topk_to_pos,
+            pos_to_expert,
+            use_tma_aligned_col_major_sf=True,
+            use_packed_ue8m0=True,
+        )
+
+
+@pytest.mark.smoke
+def test_moe_expand_to_fused_reference_rejects_duplicate_positions() -> None:
+    x = torch.randn(2, 4, dtype=torch.float16)
+    token_topk_to_pos = torch.tensor([[0, 1], [1, 2]], dtype=torch.int32)
+    pos_to_expert = torch.tensor([0, 0, 1], dtype=torch.int32)
+
+    with pytest.raises(ValueError, match="duplicate valid positions"):
+        _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+
+@pytest.mark.smoke
+def test_moe_expand_to_fused_reference_rejects_padding_target() -> None:
+    x = torch.randn(1, 4, dtype=torch.float16)
+    token_topk_to_pos = torch.tensor([[0]], dtype=torch.int32)
+    pos_to_expert = torch.tensor([-1], dtype=torch.int32)
+
+    with pytest.raises(ValueError, match="padding position"):
+        _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+
+@pytest.mark.smoke
+def test_moe_expand_to_fused_reference_rejects_manifest_dtype_mismatch() -> None:
+    x = torch.randn(1, 4, dtype=torch.float16)
+    token_topk_to_pos = torch.tensor([[0]], dtype=torch.int64)
+    pos_to_expert = torch.tensor([0], dtype=torch.int32)
+
+    with pytest.raises(ValueError, match="token_topk_to_pos must be int32"):
+        _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.smoke
+def test_moe_expand_to_fused_kernel_matches_reference_smoke() -> None:
+    pytest.importorskip("tilelang")
+    module = pytest.importorskip("tileops.kernels.moe.expand_to_fused")
+    kernel_cls = getattr(module, "MoeExpandToFusedKernel")
+
+    x = _make_x(num_tokens=4, hidden=64, dtype=torch.bfloat16).cuda()
+    token_topk_to_pos, pos_to_expert = _make_manifest_mapping(
+        num_tokens=4,
+        topk=2,
+        num_expanded_tokens=12,
+    )
+    token_topk_to_pos = token_topk_to_pos.cuda()
+    pos_to_expert = pos_to_expert.cuda()
+
+    ref = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+    kernel = kernel_cls(
+        hidden=x.shape[1],
+        num_topk=token_topk_to_pos.shape[1],
+        dtype=x.dtype,
+    )
+    out = kernel(x, token_topk_to_pos, pos_to_expert)
+    torch.cuda.synchronize()
+
+    _assert_exact(out, ref)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.smoke
+def test_moe_expand_to_fused_op_matches_reference_smoke() -> None:
+    module = pytest.importorskip("tileops.ops.moe.expand_to_fused")
+    op_cls = getattr(module, "MoeExpandToFusedFwdOp")
+
+    x = _make_x(num_tokens=4, hidden=64, dtype=torch.bfloat16).cuda()
+    token_topk_to_pos, pos_to_expert = _make_manifest_mapping(
+        num_tokens=4,
+        topk=2,
+        num_expanded_tokens=12,
+    )
+    token_topk_to_pos = token_topk_to_pos.cuda()
+    pos_to_expert = pos_to_expert.cuda()
+
+    ref = _ref_moe_expand_to_fused(x, token_topk_to_pos, pos_to_expert)
+    out = op_cls()(x, token_topk_to_pos, pos_to_expert)
+
+    _assert_exact(out, ref)
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-vvs"])

tileops/kernels/moe/expand_to_fused.py

@@ -0,0 +1,432 @@
+"""MoE expand-to-fused kernels: scatter token activations into the fused layout.
+
+Each token row of ``x`` is broadcast to the ``num_topk`` expanded positions the
+routing plan assigned to it, producing the row-contiguous, expert-major input
+the fused expert GEMM consumes. Positions are numbered so that all rows of one
+expert are adjacent, so the plain scatter below already lands the data in
+expert-major order -- the kernel never sorts, it only follows the plan.
+
+One block per position (``max(num_tokens, num_expanded_tokens)`` blocks):
+  - Blocks below ``num_expanded_tokens`` zero-fill their output row when the
+    position is unassigned (``pos_to_expert[p] < 0``).
+  - Blocks below ``num_tokens`` load their token row into a fragment and
+    scatter it to every non-negative position in ``token_topk_to_pos[t]``.
+    A negative entry marks a dropped token-expert slot and is skipped.
+
+``num_tokens`` and ``num_expanded_tokens`` are TileLang dynamic symbols, so one
+compiled kernel serves every batch size at a given static configuration.
+
+Quantized activations
+---------------------
+The same program carries a per-block scale factor (SF) tensor alongside the
+activations, moved by the identical scatter. ``num_per_channels`` (32 or 128)
+sets the block width, giving ``hidden_sf = ceil(hidden / num_per_channels)``
+scale factors per row. Two SF encodings are supported:
+
+  - **float32**, row-major ``[num_expanded_tokens, hidden_sf]``.
+  - **Packed UE8M0 as int32**: four UE8M0 exponent bytes packed per int32 word,
+    so ``hidden_sf`` shrinks by a further factor of 4. Always paired with the
+    TMA-aligned column-major layout.
+
+With ``use_tma_aligned_col_major_sf`` the SF output is written transposed as
+``[hidden_sf, num_expanded_tokens]`` over a token dimension padded to a
+multiple of 4, which is the layout the downstream TMA-based GEMM expects. The
+kernel takes it as a ``T.StridedTensor`` with a runtime leading stride, so the
+caller can hand it a slice of the padded allocation.
+
+Because the activation payload is only ever copied, never interpreted, the
+element dtype is opaque to the kernel: fp8 (``float8_e4m3fn``) and 2-way-packed
+fp4 (carried in a ``uint8``/``int8`` buffer of ``hidden / 2`` columns) both work
+through the same code path as bf16/fp16.
+
+Inputs:
+  x                 [num_tokens, hidden]           activations
+  x_sf              [num_tokens, hidden_sf]        scale factors (SF path only)
+  token_topk_to_pos [num_tokens, num_topk]         int32 (token, slot) -> position
+  pos_to_expert     [num_expanded_tokens]          int32 position -> expert (-1 = unassigned)
+
+Outputs:
+  expanded_x        [num_expanded_tokens, hidden]  same dtype as x
+  expanded_x_sf     scale factors in the requested layout (SF path only)
+"""
+
+import functools
+from typing import Optional, Tuple
+
+import tilelang
+import tilelang.language as T
+import torch
+
+from tileops.kernels.kernel_base import Kernel
+
+__all__ = ["MoeExpandToFusedKernel", "MoeExpandToFusedWithSFKernel"]
+
+_NUM_THREADS = 64
+
+# Scale-factor block widths the downstream quantized GEMM understands.
+_SUPPORTED_NUM_PER_CHANNELS = (32, 128)
+
+# UE8M0 exponent bytes packed per int32 word.
+_UE8M0_PER_INT32 = 4
+
+# Token-dimension multiple the TMA-aligned column-major SF layout pads to.
+_SF_TOKEN_ALIGN = 4
+
+def _ceil_div(x: int, y: int) -> int:
+    return -(-x // y)
+
+def _align(x: int, y: int) -> int:
+    return _ceil_div(x, y) * y
+
+def _hidden_sf(hidden: int, num_per_channels: int, use_packed_ue8m0: bool) -> int:
+    """Return the scale-factor column count for one row of ``hidden`` elements."""
+    hidden_sf = _ceil_div(hidden, num_per_channels)
+    if use_packed_ue8m0:
+        hidden_sf = _ceil_div(hidden_sf, _UE8M0_PER_INT32)
+    return hidden_sf
+
+@functools.lru_cache(maxsize=32)
+def _expand_to_fused_kernel(
+    hidden: int,
+    num_topk: int,
+    num_per_channels: Optional[int],
+    use_tma_aligned_col_major_sf: Optional[bool],
+    use_packed_ue8m0: Optional[bool],
+    x_dtype: str,
+    sf_dtype: str,
+):
+    """Build the expand-to-fused prim_func for one static configuration.
+
+    Args:
+        hidden: Hidden dimension H, in elements of ``x_dtype``.
+        num_topk: Number of expert slots per token K.
+        num_per_channels: Channels per scale-factor block (32 or 128), or None
+            to build the unquantized path that carries no scale factors.
+        use_tma_aligned_col_major_sf: Whether the SF output is transposed to
+            ``[hidden_sf, num_expanded_tokens]``. Ignored when
+            ``num_per_channels`` is None.
+        use_packed_ue8m0: Whether scale factors are UE8M0 bytes packed four to
+            an int32. Ignored when ``num_per_channels`` is None.
+        x_dtype: TileLang dtype string for the activations.
+        sf_dtype: TileLang dtype string for the scale factors. Ignored when
+            ``num_per_channels`` is None; pass ``x_dtype``.
+
+    Returns:
+        A ``@tilelang.jit`` builder returning the compiled kernel.
+    """
+    # Round the element loop up to a whole number of threads so every thread
+    # takes the same trip count; the tail iterations index past `hidden` and
+    # are predicated out against the declared buffer extent.
+    hidden_aligned = _align(hidden, _NUM_THREADS)
+
+    if num_per_channels is not None:
+        hidden_sf = _hidden_sf(hidden, num_per_channels, bool(use_packed_ue8m0))
+        hidden_sf_aligned = _align(hidden_sf, _NUM_THREADS)
+    else:
+        # The unquantized path still declares the SF parameters so both paths
+        # share one program; the caller passes None and no SF code is emitted.
+        hidden_sf, hidden_sf_aligned = 1, 1
+
+    # Leading stride of the SF output. Runtime, not static: the caller may pass
+    # a column slice of an allocation padded out to _SF_TOKEN_ALIGN.
+    sf_stride = T.dynamic("sf_stride")
+    num_tokens = T.dynamic("num_tokens")
+    num_expanded_tokens = T.dynamic("num_expanded_tokens")
+    num_blocks = T.max(num_tokens, num_expanded_tokens)
+
+    sf_shape = (
+        (hidden_sf, num_expanded_tokens)
+        if use_tma_aligned_col_major_sf
+        else (num_expanded_tokens, hidden_sf)
+    )
+
+    @tilelang.jit(
+        out_idx=[],
+        pass_configs={tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True},
+    )
+    def _expand_to_fused():
+
+        @T.prim_func
+        def _expand_to_fused_main(
+            x: T.Tensor((num_tokens, hidden), x_dtype),
+            x_sf: T.Tensor((num_tokens, hidden_sf), sf_dtype),
+            expanded_x: T.Tensor((num_expanded_tokens, hidden), x_dtype),
+            expanded_x_sf: T.StridedTensor(sf_shape, (sf_stride, 1), sf_dtype),
+            token_topk_to_pos: T.Tensor((num_tokens, num_topk), "int32"),
+            pos_to_expert: T.Tensor((num_expanded_tokens,), "int32"),
+        ):
+            with T.Kernel(num_blocks, threads=_NUM_THREADS) as (pid_token,):
+                pos_local = T.alloc_local((num_topk,), "int32")
+
+                # Unassigned positions are never a scatter target, so they must
+                # be zeroed here or the output row stays uninitialized.
+                # Kept nested: TileLang lowers `and` to a non-short-circuiting
+                # predicate, which would read pos_to_expert out of bounds on the
+                # blocks that only exist to cover num_tokens.
+                if pid_token < num_expanded_tokens:  # noqa: SIM102
+                    if pos_to_expert[pid_token] < 0:
+                        for i in T.Parallel(hidden_aligned):
+                            expanded_x[pid_token, i] = 0
+                        if num_per_channels is not None:
+                            for i in T.Parallel(hidden_sf_aligned):
+                                if use_tma_aligned_col_major_sf:
+                                    expanded_x_sf[i, pid_token] = 0
+                                else:
+                                    expanded_x_sf[pid_token, i] = 0
+
+                # The grid covers max(num_tokens, num_expanded_tokens); blocks
+                # past the token count have no row to scatter.
+                if pid_token >= num_tokens:
+                    T.thread_return()
+                T.assume(pid_token < num_tokens)
+
+                x_fragment = T.alloc_fragment((hidden_aligned,), x_dtype)
+                x_sf_fragment = T.alloc_fragment((hidden_sf_aligned,), sf_dtype)
+
+                T.copy(token_topk_to_pos[pid_token, :], pos_local)
+                T.copy(x[pid_token, :], x_fragment[0:hidden])
+                if num_per_channels is not None:
+                    T.copy(x_sf[pid_token, :], x_sf_fragment[0:hidden_sf])
+
+                # One staged read, num_topk scattered writes: the row is
+                # broadcast to every expert that selected this token.
+                for k in T.serial(num_topk):
+                    T.assume(pos_local[k] < num_expanded_tokens)
+                    if pos_local[k] >= 0:
+                        for i in T.Parallel(hidden_aligned):
+                            expanded_x[pos_local[k], i] = x_fragment[i]
+                        if num_per_channels is not None:
+                            for i in T.Parallel(hidden_sf_aligned):
+                                if use_tma_aligned_col_major_sf:
+                                    expanded_x_sf[i, pos_local[k]] = x_sf_fragment[i]
+                                else:
+                                    expanded_x_sf[pos_local[k], i] = x_sf_fragment[i]
+
+        return _expand_to_fused_main
+
+    return _expand_to_fused
+
+class MoeExpandToFusedKernel(Kernel):
+    """Scatter token activations into the fused expert layout.
+
+    Args:
+        hidden: Hidden dimension H.
+        num_topk: Number of expert slots per token K.
+        dtype: Data type of the activations (bf16 or fp16).
+        config: Optional config dict.
+
+    Example:
+        >>> kernel = MoeExpandToFusedKernel(hidden=128, num_topk=2)
+        >>> expanded_x = kernel(x, token_topk_to_pos, pos_to_expert)
+    """
+
+    supported_archs: list[int] = [80, 86, 89, 90]
+
+    def __init__(
+        self,
+        hidden: int,
+        num_topk: int,
+        dtype: torch.dtype = torch.bfloat16,
+        config: Optional[dict] = None,
+    ) -> None:
+        super().__init__()
+        self.hidden = hidden
+        self.num_topk = num_topk
+        self.dtype = dtype
+
+        self._kernel_fn = _expand_to_fused_kernel(
+            hidden,
+            num_topk,
+            None,
+            None,
+            None,
+            self.dtype_str,
+            self.dtype_str,
+        )
+
+        self.init_config(config, tune=False)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        token_topk_to_pos: torch.Tensor,
+        pos_to_expert: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run expand-to-fused.
+
+        Args:
+            x: [num_tokens, hidden] token activations.
+            token_topk_to_pos: [num_tokens, num_topk] int32 map from a
+                (token, slot) pair to its expanded position; negative marks a
+                dropped slot.
+            pos_to_expert: [num_expanded_tokens] int32 map from an expanded
+                position to its expert; negative marks an unassigned position.
+
+        Returns:
+            expanded_x: [num_expanded_tokens, hidden] activations in fused
+                expert layout. Unassigned positions are zero-filled.
+        """
+        assert x.is_cuda and token_topk_to_pos.is_cuda and pos_to_expert.is_cuda
+        assert x.is_contiguous() and token_topk_to_pos.is_contiguous()
+        assert pos_to_expert.is_contiguous()
+        assert token_topk_to_pos.dtype == torch.int32
+        assert pos_to_expert.dtype == torch.int32
+        assert x.shape[1] == self.hidden
+        assert token_topk_to_pos.shape[1] == self.num_topk
+        assert token_topk_to_pos.shape[0] == x.shape[0]
+
+        num_expanded_tokens = pos_to_expert.shape[0]
+        expanded_x = torch.empty(
+            (num_expanded_tokens, self.hidden), dtype=x.dtype, device=x.device
+        )
+        # A zero-token batch leaves nothing to scatter, and the grid would be
+        # sized only by num_expanded_tokens whose rows are all unassigned.
+        if x.shape[0] > 0:
+            self._kernel_fn()(
+                x, None, expanded_x, None, token_topk_to_pos, pos_to_expert
+            )
+        return expanded_x
+
+class MoeExpandToFusedWithSFKernel(Kernel):
+    """Scatter quantized activations and their scale factors into fused layout.
+
+    Moves the per-block scale factors alongside the activations in the same
+    pass, so the expanded activation and its scale factor stay paired. The
+    activation payload is copied without interpretation, so any 1-byte
+    quantized encoding works -- ``float8_e4m3fn``, or fp4 packed two values per
+    ``uint8``/``int8`` byte with ``hidden`` given as the packed byte count.
+
+    Args:
+        hidden: Hidden dimension H, in elements of ``dtype``.
+        num_topk: Number of expert slots per token K.
+        num_per_channels: Channels per scale-factor block; 32 or 128.
+        dtype: Data type of the quantized activations.
+        sf_dtype: Scale-factor dtype. ``torch.float32`` for plain scale
+            factors, ``torch.int32`` for Packed UE8M0.
+        use_tma_aligned_col_major_sf: Emit the SF output transposed as
+            ``[hidden_sf, num_expanded_tokens]`` over a token dimension padded
+            to a multiple of 4. Required for Packed UE8M0.
+        config: Optional config dict.
+
+    Example:
+        >>> kernel = MoeExpandToFusedWithSFKernel(hidden=256, num_topk=2,
+        ...                                       num_per_channels=128)
+        >>> expanded_x, expanded_x_sf = kernel(x, x_sf, token_topk_to_pos,
+        ...                                    pos_to_expert)
+    """
+
+    supported_archs: list[int] = [80, 86, 89, 90]
+
+    def __init__(
+        self,
+        hidden: int,
+        num_topk: int,
+        num_per_channels: int,
+        dtype: torch.dtype = torch.float8_e4m3fn,
+        sf_dtype: torch.dtype = torch.float32,
+        use_tma_aligned_col_major_sf: bool = False,
+        config: Optional[dict] = None,
+    ) -> None:
+        super().__init__()
+        assert num_per_channels in _SUPPORTED_NUM_PER_CHANNELS
+        assert sf_dtype in (torch.float32, torch.int32)
+
+        self.hidden = hidden
+        self.num_topk = num_topk
+        self.num_per_channels = num_per_channels
+        self.dtype = dtype
+        self.sf_dtype = sf_dtype
+
+        # An int32 SF buffer is Packed UE8M0 by construction: it is the only
+        # encoding that puts four exponent bytes in one word, and the packed
+        # form is defined solely for the TMA-aligned column-major layout.
+        self.use_packed_ue8m0 = sf_dtype == torch.int32
+        if self.use_packed_ue8m0:
+            assert use_tma_aligned_col_major_sf, (
+                "Packed UE8M0 scale factors require "
+                "use_tma_aligned_col_major_sf=True"
+            )
+        self.use_tma_aligned_col_major_sf = use_tma_aligned_col_major_sf
+
+        self.hidden_sf = _hidden_sf(hidden, num_per_channels, self.use_packed_ue8m0)
+
+        self._kernel_fn = _expand_to_fused_kernel(
+            hidden,
+            num_topk,
+            num_per_channels,
+            use_tma_aligned_col_major_sf,
+            self.use_packed_ue8m0,
+            self.dtype_str,
+            self.dtype_to_str(sf_dtype),
+        )
+
+        self.init_config(config, tune=False)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_sf: torch.Tensor,
+        token_topk_to_pos: torch.Tensor,
+        pos_to_expert: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Run expand-to-fused on quantized activations.
+
+        Args:
+            x: [num_tokens, hidden] quantized token activations.
+            x_sf: [num_tokens, hidden_sf] scale factors for ``x``.
+            token_topk_to_pos: [num_tokens, num_topk] int32 map from a
+                (token, slot) pair to its expanded position; negative marks a
+                dropped slot.
+            pos_to_expert: [num_expanded_tokens] int32 map from an expanded
+                position to its expert; negative marks an unassigned position.
+
+        Returns:
+            A tuple ``(expanded_x, expanded_x_sf)``. ``expanded_x`` is
+            [num_expanded_tokens, hidden]; ``expanded_x_sf`` is
+            [num_expanded_tokens, hidden_sf], returned as a transposed view of
+            the column-major buffer when
+            ``use_tma_aligned_col_major_sf``. Unassigned positions are
+            zero-filled in both.
+        """
+        assert x.is_cuda and x_sf.is_cuda
+        assert token_topk_to_pos.is_cuda and pos_to_expert.is_cuda
+        assert x.is_contiguous() and x_sf.is_contiguous()
+        assert token_topk_to_pos.is_contiguous() and pos_to_expert.is_contiguous()
+        assert token_topk_to_pos.dtype == torch.int32
+        assert pos_to_expert.dtype == torch.int32
+        assert x.dtype == self.dtype and x_sf.dtype == self.sf_dtype
+        assert x.shape[1] == self.hidden
+        assert x_sf.shape == (x.shape[0], self.hidden_sf)
+        assert token_topk_to_pos.shape[1] == self.num_topk
+        assert token_topk_to_pos.shape[0] == x.shape[0]
+
+        num_expanded_tokens = pos_to_expert.shape[0]
+        expanded_x = torch.empty(
+            (num_expanded_tokens, self.hidden), dtype=x.dtype, device=x.device
+        )
+
+        if self.use_tma_aligned_col_major_sf:
+            # Allocate over the padded token dimension the TMA layout needs,
+            # then hand the kernel the exact column slice. The padding columns
+            # stay untouched and are never returned.
+            num_expanded_sf_tokens = _align(num_expanded_tokens, _SF_TOKEN_ALIGN)
+            expanded_x_sf = torch.empty(
+                (self.hidden_sf, num_expanded_sf_tokens),
+                dtype=x_sf.dtype,
+                device=x_sf.device,
+            )[:, :num_expanded_tokens]
+        else:
+            expanded_x_sf = torch.empty(
+                (num_expanded_tokens, self.hidden_sf),
+                dtype=x_sf.dtype,
+                device=x_sf.device,
+            )
+
+        if x.shape[0] > 0:
+            self._kernel_fn()(
+                x, x_sf, expanded_x, expanded_x_sf, token_topk_to_pos, pos_to_expert
+            )
+
+        if self.use_tma_aligned_col_major_sf:
+            expanded_x_sf = expanded_x_sf.T
+        return expanded_x, expanded_x_sf

tileops/manifest/moe.yaml

@@ -839,3 +839,106 @@ MoeTopKSumGroupIdxFwdOp:
     op: tileops/ops/moe/topk_sum_group_idx.py
     test: tests/ops/test_moe_topk_sum_group_idx.py
     bench: benchmarks/ops/bench_moe_topk_sum_group_idx.py
+
+MoeExpandToFusedFwdOp:
+  ref_api: "tile_kernels.torch.expand_to_fused"
+  family: moe
+  status: spec-only
+
+  signature:
+    inputs:
+      x: {dtype: "float16 | bfloat16", shape: "[T, H]"}
+      token_topk_to_pos: {dtype: "int32", shape: "[T, K]"}
+      pos_to_expert: {dtype: "int32", shape: "[X]"}
+    outputs:
+      expanded_x: {dtype: "same_as(x)", shape: "[X, H]"}
+    params: {}
+    shape_rules:
+    - "T > 0"
+    - "K > 0"
+    - "X >= 0"
+    - "expanded_x.shape == (pos_to_expert.shape[0], x.shape[1])"
+
+  workloads:
+  - {x_shape: [1, 7168], token_topk_to_pos_shape: [1, 8], pos_to_expert_shape: [16], dtypes: [bfloat16], label: "kimi-k2-decode"}
+  - {x_shape: [32, 7168], token_topk_to_pos_shape: [32, 8], pos_to_expert_shape: [256], dtypes: [bfloat16], label: "kimi-k2-small"}
+  - {x_shape: [512, 7168], token_topk_to_pos_shape: [512, 8], pos_to_expert_shape: [4096], dtypes: [bfloat16], label: "kimi-k2-medium"}
+  - {x_shape: [4096, 7168], token_topk_to_pos_shape: [4096, 8], pos_to_expert_shape: [32768], dtypes: [bfloat16], label: "kimi-k2-prefill"}
+  - {x_shape: [1, 3072], token_topk_to_pos_shape: [1, 8], pos_to_expert_shape: [16], dtypes: [bfloat16], label: "qwen3-30b-decode"}
+  - {x_shape: [32, 3072], token_topk_to_pos_shape: [32, 8], pos_to_expert_shape: [256], dtypes: [bfloat16], label: "qwen3-30b-small"}
+  - {x_shape: [512, 3072], token_topk_to_pos_shape: [512, 8], pos_to_expert_shape: [4096], dtypes: [bfloat16], label: "qwen3-30b-medium"}
+  - {x_shape: [4096, 3072], token_topk_to_pos_shape: [4096, 8], pos_to_expert_shape: [32768], dtypes: [bfloat16], label: "qwen3-30b-prefill"}
+
+  roofline:
+    vars:
+      T: "x.shape[0]"
+      H: "x.shape[1]"
+      K: "token_topk_to_pos.shape[1]"
+      X: "pos_to_expert.shape[0]"
+    flops: "0"
+    # x read + token/position mappings read + expanded output write.
+    bytes: "T * H * elem_bytes + T * K * 4 + X * 4 + X * H * elem_bytes"
+
+  source:
+    # Planned PR B locations; path existence is deferred while spec-only.
+    kernel: tileops/kernels/moe/expand_to_fused.py
+    op: tileops/ops/moe/expand_to_fused.py
+    test: tests/ops/test_moe_expand_to_fused.py
+    bench: benchmarks/ops/bench_moe_expand_to_fused.py
+
+MoeExpandToFusedWithSFFwdOp:
+  ref_api: "tile_kernels.torch.expand_to_fused_with_sf"
+  family: moe
+  status: spec-only
+  variant_of: MoeExpandToFusedFwdOp
+
+  signature:
+    inputs:
+      x: {dtype: "float8_e4m3fn", shape: "[T, H]"}
+      x_sf: {dtype: "float32 | int32", shape: "[T, S]"}
+      token_topk_to_pos: {dtype: "int32", shape: "[T, K]"}
+      pos_to_expert: {dtype: "int32", shape: "[X]"}
+    outputs:
+      expanded_x: {dtype: "same_as(x)", shape: "[X, H]"}
+      expanded_x_sf: {dtype: "same_as(x_sf)", shape: "[X, S]"}
+    params:
+      num_per_channels: {type: int}
+      use_tma_aligned_col_major_sf: {type: bool, default: false}
+    shape_rules:
+    - "T > 0"
+    - "K > 0"
+    - "S > 0"
+    - "X >= 0"
+    - "num_per_channels == 32 or num_per_channels == 128"
+    - "expanded_x.shape == (pos_to_expert.shape[0], x.shape[1])"
+    - "expanded_x_sf.shape == (pos_to_expert.shape[0], x_sf.shape[1])"
+    # For float32 scale factors S=ceil(H/num_per_channels). For packed
+    # UE8M0/int32 scale factors S=ceil(ceil(H/num_per_channels)/4), with
+    # use_tma_aligned_col_major_sf=True. PR B must validate both forms.
+
+  workloads:
+  - {x_shape: [1, 7168], x_sf_shape: [1, 56], token_topk_to_pos_shape: [1, 8], pos_to_expert_shape: [16], num_per_channels: 128, use_tma_aligned_col_major_sf: false, dtypes: [float8_e4m3fn], label: "kimi-k2-decode-sf-row-major"}
+  - {x_shape: [512, 7168], x_sf_shape: [512, 56], token_topk_to_pos_shape: [512, 8], pos_to_expert_shape: [4096], num_per_channels: 128, use_tma_aligned_col_major_sf: false, dtypes: [float8_e4m3fn], label: "kimi-k2-medium-sf-row-major"}
+  - {x_shape: [1, 7168], x_sf_shape: [1, 14], token_topk_to_pos_shape: [1, 8], pos_to_expert_shape: [16], num_per_channels: 128, use_tma_aligned_col_major_sf: true, dtypes: [float8_e4m3fn], label: "kimi-k2-decode-sf-packed"}
+  - {x_shape: [512, 7168], x_sf_shape: [512, 14], token_topk_to_pos_shape: [512, 8], pos_to_expert_shape: [4096], num_per_channels: 128, use_tma_aligned_col_major_sf: true, dtypes: [float8_e4m3fn], label: "kimi-k2-medium-sf-packed"}
+  - {x_shape: [32, 3072], x_sf_shape: [32, 24], token_topk_to_pos_shape: [32, 8], pos_to_expert_shape: [256], num_per_channels: 128, use_tma_aligned_col_major_sf: false, dtypes: [float8_e4m3fn], label: "qwen3-30b-small-sf-row-major"}
+  - {x_shape: [32, 3072], x_sf_shape: [32, 6], token_topk_to_pos_shape: [32, 8], pos_to_expert_shape: [256], num_per_channels: 128, use_tma_aligned_col_major_sf: true, dtypes: [float8_e4m3fn], label: "qwen3-30b-small-sf-packed"}
+
+  roofline:
+    vars:
+      T: "x.shape[0]"
+      H: "x.shape[1]"
+      K: "token_topk_to_pos.shape[1]"
+      S: "x_sf.shape[1]"
+      X: "pos_to_expert.shape[0]"
+    flops: "0"
+    # x/x_sf read + token/position mappings read + expanded x/x_sf writes.
+    # Scale factors are always 4-byte float32 or packed int32 values.
+    bytes: "T * H * elem_bytes + T * S * 4 + T * K * 4 + X * 4 + X * H * elem_bytes + X * S * 4"
+
+  source:
+    # Same planned PR B implementation as the primary entry.
+    kernel: tileops/kernels/moe/expand_to_fused.py
+    op: tileops/ops/moe/expand_to_fused.py
+    test: tests/ops/test_moe_expand_to_fused.py
+    bench: benchmarks/ops/bench_moe_expand_to_fused.py

tileops/ops/moe/__init__.py

@@ -1,5 +1,6 @@
 """MoE operator package."""
 
+from .expand_to_fused import MoeExpandToFusedFwdOp, MoeExpandToFusedWithSFFwdOp
 from .fused_moe import FusedMoe, FusedMoeFwdCbFwdOp, FusedMoeFwdOp
 from .fused_topk import FusedTopKOp
 from .permute_align import MoePermuteAlignFwdOp
@@ -29,6 +30,8 @@ __all__ = [
     "FusedMoeFwdOp",
     "FusedTopKOp",
     "MoEPrepareAndFinalizeNoDPEP",
+    "MoeExpandToFusedFwdOp",
+    "MoeExpandToFusedWithSFFwdOp",
     "MoeGroupedGemmNopad3WGFusedActFwdOp",
     "MoeGroupedGemmNopadFwdOp",
     "MoePermuteAlignFwdOp",

tileops/ops/moe/expand_to_fused.py

@@ -0,0 +1,514 @@
+"""MoE expand-to-fused ops: scatter tokens into the fused expert layout.
+
+Provides:
+  - MoeExpandToFusedFwdOp: expanded_x[p] = x[t] for every routed (t, k) -> p
+  - MoeExpandToFusedWithSFFwdOp: the same scatter for quantized activations,
+    carrying per-block scale factors alongside the data
+"""
+
+from typing import Dict, Optional, Tuple
+
+import torch
+
+from tileops.kernels.kernel_base import Kernel
+from tileops.kernels.moe.expand_to_fused import (
+    MoeExpandToFusedKernel,
+    MoeExpandToFusedWithSFKernel,
+)
+
+from ..op_base import Op
+
+__all__ = ["MoeExpandToFusedFwdOp", "MoeExpandToFusedWithSFFwdOp"]
+
+# Scale-factor block widths declared by the manifest shape rules.
+_SUPPORTED_NUM_PER_CHANNELS = (32, 128)
+
+# UE8M0 exponent bytes packed per int32 word.
+_UE8M0_PER_INT32 = 4
+
+def _ceil_div(x: int, y: int) -> int:
+    return -(-x // y)
+
+def _check_routing_inputs(
+    x: torch.Tensor,
+    token_topk_to_pos: torch.Tensor,
+    pos_to_expert: torch.Tensor,
+    tensors: Tuple[Tuple[str, torch.Tensor], ...],
+) -> Tuple[int, int, int, int]:
+    """Validate the routing-plan contract shared by both expand-to-fused ops.
+
+    Args:
+        x: Activation tensor, the device and token-count reference.
+        token_topk_to_pos: [T, K] int32 routing map.
+        pos_to_expert: [X] int32 position-to-expert map.
+        tensors: All (name, tensor) pairs to residency-check, in signature
+            order.
+
+    Returns:
+        ``(num_tokens, hidden, num_topk, num_expanded_tokens)``.
+
+    Raises:
+        ValueError: If any residency, rank, or shape rule is violated.
+    """
+    for name, t in tensors:
+        if not t.is_cuda:
+            raise ValueError(f"{name} must be a CUDA tensor")
+        if t.device != x.device:
+            raise ValueError(
+                f"Expected all inputs on the same device, got {x.device} "
+                f"for x and {t.device} for {name}"
+            )
+    if x.ndim != 2:
+        raise ValueError(f"Expected x to be 2D [T, H], got {x.ndim}D")
+    if token_topk_to_pos.ndim != 2:
+        raise ValueError(
+            f"Expected token_topk_to_pos to be 2D [T, K], "
+            f"got {token_topk_to_pos.ndim}D"
+        )
+    if pos_to_expert.ndim != 1:
+        raise ValueError(
+            f"Expected pos_to_expert to be 1D [X], got {pos_to_expert.ndim}D"
+        )
+
+    num_tokens, hidden = x.shape
+    if num_tokens <= 0:
+        raise ValueError(f"Expected x.shape[0] > 0, got {num_tokens}")
+    if token_topk_to_pos.shape[0] != num_tokens:
+        raise ValueError(
+            f"Expected token_topk_to_pos.shape[0] == x.shape[0] "
+            f"({num_tokens}), got {token_topk_to_pos.shape[0]}"
+        )
+    num_topk = token_topk_to_pos.shape[1]
+    if num_topk <= 0:
+        raise ValueError(
+            f"Expected token_topk_to_pos.shape[1] > 0, got {num_topk}"
+        )
+    return num_tokens, hidden, num_topk, pos_to_expert.shape[0]
+
+def _check_index_dtypes(
+    token_topk_to_pos: torch.Tensor, pos_to_expert: torch.Tensor
+) -> None:
+    """Raise if either routing map is not int32."""
+    if token_topk_to_pos.dtype != torch.int32:
+        raise ValueError(
+            f"Expected token_topk_to_pos.dtype torch.int32, "
+            f"got {token_topk_to_pos.dtype}"
+        )
+    if pos_to_expert.dtype != torch.int32:
+        raise ValueError(
+            f"Expected pos_to_expert.dtype torch.int32, got {pos_to_expert.dtype}"
+        )
+
+class MoeExpandToFusedFwdOp(Op):
+    """Expand token activations into the fused expert layout.
+
+    ``token_topk_to_pos`` maps each (token, expert-slot) pair to a row of the
+    expanded buffer; ``pos_to_expert`` labels each expanded row with its
+    expert. Both use a negative entry as the sentinel for a dropped slot /
+    unassigned position, so the expanded row count is data-dependent and is
+    taken from ``pos_to_expert`` rather than committed at construction.
+
+    Args:
+        dtype: Optional committed activation dtype (bf16 or fp16). Preferred
+            API infers it from ``x``.
+        kernel_map: Optional override for kernel dispatch.
+
+    Example:
+        >>> op = MoeExpandToFusedFwdOp()
+        >>> expanded_x = op(x, token_topk_to_pos, pos_to_expert)
+    """
+
+    _SUPPORTED_DTYPES = (torch.float16, torch.bfloat16)
+
+    # The manifest entry declares no static_dims: hidden and num_topk are read
+    # from the inputs at forward time, and the token counts are TileLang
+    # dynamic symbols. See _cache_key for the resulting kernel-cache
+    # projection.
+    _static_axes: frozenset[tuple[int, int]] = frozenset()
+
+    def __init__(
+        self,
+        *,
+        dtype: Optional[torch.dtype] = None,
+        kernel_map: Optional[Dict[str, Kernel]] = None,
+    ) -> None:
+        self.dtype = dtype
+        self._committed_dtype = dtype
+
+        self.dispatch_kernel(kernel_map)
+        self._kernel_cache: Dict[Tuple[int, int, torch.dtype], Kernel] = {}
+
+    @property
+    def default_kernel_map(self) -> Dict[str, Kernel]:
+        return {"expand_to_fused_kernel": MoeExpandToFusedKernel}
+
+    def _cache_key(
+        self,
+        x_shape: tuple,
+        token_topk_to_pos_shape: tuple,
+        pos_to_expert_shape: tuple,
+    ) -> Tuple[int, int]:
+        """Project input shapes onto what the compiled kernel depends on.
+
+        Only ``hidden`` and ``num_topk`` are baked into the prim_func; the two
+        token counts are ``T.dynamic`` symbols, so batch size does not
+        fragment the cache.
+        """
+        return (x_shape[1], token_topk_to_pos_shape[1])
+
+    def _infer_output_shapes(
+        self,
+        x_shape: tuple,
+        token_topk_to_pos_shape: tuple,
+        pos_to_expert_shape: tuple,
+    ) -> Dict[str, tuple]:
+        return {"expanded_x": (pos_to_expert_shape[0], x_shape[1])}
+
+    def _validate_dtypes(
+        self,
+        x: torch.Tensor,
+        token_topk_to_pos: torch.Tensor,
+        pos_to_expert: torch.Tensor,
+    ) -> None:
+        if x.dtype not in self._SUPPORTED_DTYPES:
+            names = ", ".join(str(dt) for dt in self._SUPPORTED_DTYPES)
+            raise ValueError(f"Expected x.dtype in [{names}], got {x.dtype}")
+        _check_index_dtypes(token_topk_to_pos, pos_to_expert)
+
+    def eval_roofline(self) -> tuple[int, int]:
+        if (
+            not hasattr(self, "x_shape")
+            or not hasattr(self, "token_topk_to_pos_shape")
+            or not hasattr(self, "pos_to_expert_shape")
+            or self.dtype is None
+        ):
+            raise ValueError(
+                "MoeExpandToFusedFwdOp.eval_roofline() requires a prior forward() "
+                "to bind x_shape, token_topk_to_pos_shape, pos_to_expert_shape, "
+                "and dtype"
+            )
+        num_tokens, hidden = self.x_shape
+        num_topk = self.token_topk_to_pos_shape[1]
+        num_expanded_tokens = self.pos_to_expert_shape[0]
+        elem_bytes = self.dtype.itemsize
+        nbytes = (
+            num_tokens * hidden * elem_bytes
+            + num_tokens * num_topk * 4
+            + num_expanded_tokens * 4
+            + num_expanded_tokens * hidden * elem_bytes
+        )
+        return 0, int(nbytes)
+
+    def _get_kernel(self, hidden: int, num_topk: int, dtype: torch.dtype) -> Kernel:
+        key = (hidden, num_topk, dtype)
+        if key not in self._kernel_cache:
+            self._kernel_cache[key] = self.kernel_map["expand_to_fused_kernel"](
+                hidden, num_topk, dtype
+            )
+        return self._kernel_cache[key]
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        token_topk_to_pos: torch.Tensor,
+        pos_to_expert: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run expand-to-fused.
+
+        Args:
+            x: [num_tokens, hidden] token activations (bf16/fp16).
+            token_topk_to_pos: [num_tokens, num_topk] int32 map from a
+                (token, slot) pair to its expanded position; negative marks a
+                dropped slot.
+            pos_to_expert: [num_expanded_tokens] int32 map from an expanded
+                position to its expert; negative marks an unassigned position.
+
+        Returns:
+            expanded_x: [num_expanded_tokens, hidden] activations in fused
+                expert layout. Unassigned positions are zero-filled.
+        """
+        self._validate_dtypes(x, token_topk_to_pos, pos_to_expert)
+        _, hidden, num_topk, _ = _check_routing_inputs(
+            x,
+            token_topk_to_pos,
+            pos_to_expert,
+            (
+                ("x", x),
+                ("token_topk_to_pos", token_topk_to_pos),
+                ("pos_to_expert", pos_to_expert),
+            ),
+        )
+        if self._committed_dtype is not None and x.dtype != self._committed_dtype:
+            raise ValueError(
+                f"Expected x.dtype {self._committed_dtype}, got {x.dtype}"
+            )
+
+        x = x.contiguous()
+        token_topk_to_pos = token_topk_to_pos.contiguous()
+        pos_to_expert = pos_to_expert.contiguous()
+
+        self.dtype = x.dtype
+        self.x_shape = tuple(x.shape)
+        self.token_topk_to_pos_shape = tuple(token_topk_to_pos.shape)
+        self.pos_to_expert_shape = tuple(pos_to_expert.shape)
+
+        kernel = self._get_kernel(hidden, num_topk, x.dtype)
+        return kernel(x, token_topk_to_pos, pos_to_expert)
+
+class MoeExpandToFusedWithSFFwdOp(Op):
+    """Expand quantized activations and their scale factors into fused layout.
+
+    The routing contract is identical to :class:`MoeExpandToFusedFwdOp`; this
+    variant additionally moves the per-block scale factors (SF) so each
+    expanded row keeps the scale that decodes it.
+
+    ``num_per_channels`` sets the SF block width, giving
+    ``S = ceil(H / num_per_channels)`` float32 scale factors per row. Passing
+    an int32 ``x_sf`` selects Packed UE8M0, where four UE8M0 exponent bytes
+    share one word and ``S`` shrinks by a further factor of 4; that encoding is
+    only defined for the TMA-aligned column-major SF layout.
+
+    Args:
+        num_per_channels: Channels per scale-factor block; 32 or 128.
+        use_tma_aligned_col_major_sf: Return ``expanded_x_sf`` as a transposed
+            view of a column-major buffer whose token dimension is padded to a
+            multiple of 4, which is what the downstream TMA-based grouped GEMM
+            reads. Required for Packed UE8M0.
+        dtype: Optional committed activation dtype. Preferred API infers it
+            from ``x``.
+        sf_dtype: Optional committed scale-factor dtype. Preferred API infers
+            it from ``x_sf``.
+        kernel_map: Optional override for kernel dispatch.
+
+    Example:
+        >>> op = MoeExpandToFusedWithSFFwdOp(num_per_channels=128)
+        >>> expanded_x, expanded_x_sf = op(x, x_sf, token_topk_to_pos,
+        ...                                pos_to_expert)
+    """
+
+    # The kernel copies the activation payload without interpreting it, so fp8
+    # and packed-fp4 byte storage both run through the same path unchanged.
+    _SUPPORTED_DTYPES = (torch.float8_e4m3fn, torch.uint8, torch.int8)
+    _SUPPORTED_SF_DTYPES = (torch.float32, torch.int32)
+
+    # As in the unquantized op: hidden, num_topk and the SF width are read from
+    # the inputs at forward time, and both token counts are dynamic symbols.
+    _static_axes: frozenset[tuple[int, int]] = frozenset()
+
+    def __init__(
+        self,
+        *,
+        num_per_channels: int,
+        use_tma_aligned_col_major_sf: bool = False,
+        dtype: Optional[torch.dtype] = None,
+        sf_dtype: Optional[torch.dtype] = None,
+        kernel_map: Optional[Dict[str, Kernel]] = None,
+    ) -> None:
+        if num_per_channels not in _SUPPORTED_NUM_PER_CHANNELS:
+            raise ValueError(
+                f"Expected num_per_channels in "
+                f"{list(_SUPPORTED_NUM_PER_CHANNELS)}, got {num_per_channels}"
+            )
+        self.num_per_channels = num_per_channels
+        self.use_tma_aligned_col_major_sf = use_tma_aligned_col_major_sf
+
+        self.dtype = dtype
+        self.sf_dtype = sf_dtype
+        self._committed_dtype = dtype
+        self._committed_sf_dtype = sf_dtype
+
+        self.dispatch_kernel(kernel_map)
+        self._kernel_cache: Dict[
+            Tuple[int, int, torch.dtype, torch.dtype], Kernel
+        ] = {}
+
+    @property
+    def default_kernel_map(self) -> Dict[str, Kernel]:
+        return {"expand_to_fused_with_sf_kernel": MoeExpandToFusedWithSFKernel}
+
+    def _expected_hidden_sf(self, hidden: int, sf_dtype: torch.dtype) -> int:
+        """Return the SF column count the manifest shape rules require."""
+        hidden_sf = _ceil_div(hidden, self.num_per_channels)
+        if sf_dtype == torch.int32:
+            hidden_sf = _ceil_div(hidden_sf, _UE8M0_PER_INT32)
+        return hidden_sf
+
+    def _cache_key(
+        self,
+        x_shape: tuple,
+        x_sf_shape: tuple,
+        token_topk_to_pos_shape: tuple,
+        pos_to_expert_shape: tuple,
+    ) -> Tuple[int, int]:
+        """Project input shapes onto what the compiled kernel depends on.
+
+        ``hidden`` and ``num_topk`` are baked into the prim_func; the SF width
+        is derived from ``hidden`` and the ctor-committed ``num_per_channels``,
+        so it adds no independent axis. Both token counts are ``T.dynamic``
+        symbols.
+        """
+        return (x_shape[1], token_topk_to_pos_shape[1])
+
+    def _infer_output_shapes(
+        self,
+        x_shape: tuple,
+        x_sf_shape: tuple,
+        token_topk_to_pos_shape: tuple,
+        pos_to_expert_shape: tuple,
+    ) -> Dict[str, tuple]:
+        num_expanded_tokens = pos_to_expert_shape[0]
+        return {
+            "expanded_x": (num_expanded_tokens, x_shape[1]),
+            "expanded_x_sf": (num_expanded_tokens, x_sf_shape[1]),
+        }
+
+    def _validate_dtypes(
+        self,
+        x: torch.Tensor,
+        x_sf: torch.Tensor,
+        token_topk_to_pos: torch.Tensor,
+        pos_to_expert: torch.Tensor,
+    ) -> None:
+        if x.dtype not in self._SUPPORTED_DTYPES:
+            names = ", ".join(str(dt) for dt in self._SUPPORTED_DTYPES)
+            raise ValueError(f"Expected x.dtype in [{names}], got {x.dtype}")
+        if x_sf.dtype not in self._SUPPORTED_SF_DTYPES:
+            names = ", ".join(str(dt) for dt in self._SUPPORTED_SF_DTYPES)
+            raise ValueError(f"Expected x_sf.dtype in [{names}], got {x_sf.dtype}")
+        _check_index_dtypes(token_topk_to_pos, pos_to_expert)
+
+    def eval_roofline(self) -> tuple[int, int]:
+        if (
+            not hasattr(self, "x_shape")
+            or not hasattr(self, "x_sf_shape")
+            or not hasattr(self, "token_topk_to_pos_shape")
+            or not hasattr(self, "pos_to_expert_shape")
+            or self.dtype is None
+        ):
+            raise ValueError(
+                "MoeExpandToFusedWithSFFwdOp.eval_roofline() requires a prior "
+                "forward() to bind x_shape, x_sf_shape, token_topk_to_pos_shape, "
+                "pos_to_expert_shape, and dtype"
+            )
+        num_tokens, hidden = self.x_shape
+        hidden_sf = self.x_sf_shape[1]
+        num_topk = self.token_topk_to_pos_shape[1]
+        num_expanded_tokens = self.pos_to_expert_shape[0]
+        elem_bytes = self.dtype.itemsize
+        # Scale factors are 4-byte float32 or packed int32 words either way.
+        nbytes = (
+            num_tokens * hidden * elem_bytes
+            + num_tokens * hidden_sf * 4
+            + num_tokens * num_topk * 4
+            + num_expanded_tokens * 4
+            + num_expanded_tokens * hidden * elem_bytes
+            + num_expanded_tokens * hidden_sf * 4
+        )
+        return 0, int(nbytes)
+
+    def _get_kernel(
+        self,
+        hidden: int,
+        num_topk: int,
+        dtype: torch.dtype,
+        sf_dtype: torch.dtype,
+    ) -> Kernel:
+        key = (hidden, num_topk, dtype, sf_dtype)
+        if key not in self._kernel_cache:
+            self._kernel_cache[key] = self.kernel_map[
+                "expand_to_fused_with_sf_kernel"
+            ](
+                hidden,
+                num_topk,
+                self.num_per_channels,
+                dtype,
+                sf_dtype,
+                self.use_tma_aligned_col_major_sf,
+            )
+        return self._kernel_cache[key]
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        x_sf: torch.Tensor,
+        token_topk_to_pos: torch.Tensor,
+        pos_to_expert: torch.Tensor,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Run expand-to-fused on quantized activations.
+
+        Args:
+            x: [num_tokens, hidden] quantized token activations.
+            x_sf: [num_tokens, hidden_sf] scale factors for ``x``; float32, or
+                int32 for Packed UE8M0.
+            token_topk_to_pos: [num_tokens, num_topk] int32 map from a
+                (token, slot) pair to its expanded position; negative marks a
+                dropped slot.
+            pos_to_expert: [num_expanded_tokens] int32 map from an expanded
+                position to its expert; negative marks an unassigned position.
+
+        Returns:
+            A tuple ``(expanded_x, expanded_x_sf)`` of shapes
+            [num_expanded_tokens, hidden] and [num_expanded_tokens, hidden_sf].
+            Unassigned positions are zero-filled in both. Under
+            ``use_tma_aligned_col_major_sf`` the SF result is a transposed view
+            of a column-major buffer, so it is not row-major contiguous.
+        """
+        self._validate_dtypes(x, x_sf, token_topk_to_pos, pos_to_expert)
+        _, hidden, num_topk, _ = _check_routing_inputs(
+            x,
+            token_topk_to_pos,
+            pos_to_expert,
+            (
+                ("x", x),
+                ("x_sf", x_sf),
+                ("token_topk_to_pos", token_topk_to_pos),
+                ("pos_to_expert", pos_to_expert),
+            ),
+        )
+        if x_sf.ndim != 2:
+            raise ValueError(f"Expected x_sf to be 2D [T, S], got {x_sf.ndim}D")
+        if x_sf.shape[0] != x.shape[0]:
+            raise ValueError(
+                f"Expected x_sf.shape[0] == x.shape[0] ({x.shape[0]}), "
+                f"got {x_sf.shape[0]}"
+            )
+        # Packed UE8M0 only exists in the TMA-aligned column-major layout: the
+        # packing groups four scale factors along the token-major axis.
+        if x_sf.dtype == torch.int32 and not self.use_tma_aligned_col_major_sf:
+            raise ValueError(
+                "int32 x_sf selects Packed UE8M0, which requires "
+                "use_tma_aligned_col_major_sf=True"
+            )
+        expected_hidden_sf = self._expected_hidden_sf(hidden, x_sf.dtype)
+        if x_sf.shape[1] != expected_hidden_sf:
+            raise ValueError(
+                f"Expected x_sf.shape[1] == {expected_hidden_sf} for hidden "
+                f"{hidden}, num_per_channels {self.num_per_channels} and "
+                f"x_sf.dtype {x_sf.dtype}, got {x_sf.shape[1]}"
+            )
+        if self._committed_dtype is not None and x.dtype != self._committed_dtype:
+            raise ValueError(
+                f"Expected x.dtype {self._committed_dtype}, got {x.dtype}"
+            )
+        if (
+            self._committed_sf_dtype is not None
+            and x_sf.dtype != self._committed_sf_dtype
+        ):
+            raise ValueError(
+                f"Expected x_sf.dtype {self._committed_sf_dtype}, got {x_sf.dtype}"
+            )
+
+        x = x.contiguous()
+        x_sf = x_sf.contiguous()
+        token_topk_to_pos = token_topk_to_pos.contiguous()
+        pos_to_expert = pos_to_expert.contiguous()
+
+        self.dtype = x.dtype
+        self.sf_dtype = x_sf.dtype
+        self.x_shape = tuple(x.shape)
+        self.x_sf_shape = tuple(x_sf.shape)
+        self.token_topk_to_pos_shape = tuple(token_topk_to_pos.shape)
+        self.pos_to_expert_shape = tuple(pos_to_expert.shape)
+
+        kernel = self._get_kernel(hidden, num_topk, x.dtype, x_sf.dtype)
+        return kernel(x, x_sf, token_topk_to_pos, pos_to_expert)