diff --git a/examples/neuron/device_diff.py b/examples/neuron/device_diff.py
new file mode 100644
index 0000000..a6a0ac1
--- /dev/null
+++ b/examples/neuron/device_diff.py
@@ -0,0 +1,216 @@
+"""Device-vs-host differential check for the Neuron kernels.
+
+Runs the tiny checkpoint on the NeuronCore through run_deepinfer_tiny.py and, for every call of a vendor
+kernel (deep_gemm, tile_kernels, flash_mla, deep_select) or a tilelang_ops/neuron kernel, re-runs the same
+pure-torch implementation on host copies of the arguments and compares the results, including arguments
+mutated in place. Prints one line per kernel with the worst relative error and the argument layouts of
+the worst call, so a kernel that misbehaves on the device (and only there) stands out.
+
+    NEURON_RT_VISIBLE_CORES=63 CUDA_VISIBLE_DEVICES=63 python examples/neuron/device_diff.py \\
+        --prompts short --out ~/di/neuron_diff [run_deepinfer_tiny.py args...]
+"""
+
+from __future__ import annotations
+
+import collections
+import functools
+import sys
+import types
+from pathlib import Path
+from typing import Any
+
+import torch
+
+HERE = Path(__file__).resolve().parent
+sys.path.insert(0, str(HERE))
+
+_STATS: dict[str, dict[str, Any]] = collections.OrderedDict()
+_ACTIVE = [True]
+
+
+def _to_host(x: Any) -> Any:
+    if isinstance(x, torch.Tensor):
+        # Keep strides: a few kernels depend on the caller's (e.g. column-major) layout.
+        return torch.empty_strided(x.shape, x.stride(), dtype=x.dtype, device="cpu").copy_(x.cpu()) if x.device.type != "cpu" else x.clone()
+    if isinstance(x, (list, tuple)):
+        return type(x)(_to_host(v) for v in x)
+    if isinstance(x, dict):
+        return {k: _to_host(v) for k, v in x.items()}
+    return x
+
+
+def _tensors(x: Any, prefix: str = "") -> list[tuple[str, torch.Tensor]]:
+    if isinstance(x, torch.Tensor):
+        return [(prefix, x)]
+    if isinstance(x, (list, tuple)):
+        return [t for i, v in enumerate(x) for t in _tensors(v, f"{prefix}[{i}]")]
+    if isinstance(x, dict):
+        return [t for k, v in x.items() for t in _tensors(v, f"{prefix}.{k}")]
+    return []
+
+
+def _rel(dev: torch.Tensor, host: torch.Tensor) -> float:
+    a = dev.detach().cpu()
+    b = host.detach()
+    if a.shape != b.shape:
+        return float("nan")
+    if a.dtype in (torch.float8_e4m3fn, torch.float8_e8m0fnu) or not a.is_floating_point():
+        if a.dtype in (torch.float8_e4m3fn, torch.float8_e8m0fnu):
+            a, b = a.view(torch.uint8), b.view(torch.uint8)
+        neq = (a != b).float().mean().item() if a.numel() else 0.0
+        return neq  # fraction of differing elements / bytes
+    a, b = a.float(), b.float()
+    finite = torch.isfinite(b) & torch.isfinite(a)
+    if not finite.all():
+        if not torch.equal(torch.isfinite(a), torch.isfinite(b)):
+            return float("inf")
+        a, b = a[finite], b[finite]
+    den = b.norm().item()
+    return (a - b).norm().item() / den if den > 0 else (a - b).norm().item()
+
+
+def _layout(args: Any, kwargs: Any) -> str:
+    parts = []
+    for name, t in _tensors(args, "a") + _tensors(kwargs, "k"):
+        parts.append(f"{name}:{str(t.dtype).replace('torch.', '')}{list(t.shape)}s{list(t.stride())}")
+    return " ".join(parts)[:400]
+
+
+def _record(name: str, rel: float, where: str, layout: str) -> None:
+    st = _STATS.setdefault(name, {"calls": 0, "worst": 0.0, "where": "", "layout": ""})
+    st["calls"] += 1
+    if not (rel <= st["worst"]):  # also catches nan/inf
+        st["worst"], st["where"], st["layout"] = rel, where, layout
+
+
+def wrap(name: str, fn: Any) -> Any:
+    if getattr(fn, "_device_diff", False):
+        return fn
+
+    @functools.wraps(fn)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        if not _ACTIVE[0]:
+            return fn(*args, **kwargs)
+        h_args, h_kwargs = _to_host(args), _to_host(kwargs)
+        layout = _layout(args, kwargs)
+        out = fn(*args, **kwargs)
+        _ACTIVE[0] = False  # nested kernels inside the host replay are not checked again
+        try:
+            h_out = fn(*h_args, **h_kwargs)
+        except Exception as e:  # noqa: BLE001
+            _record(name, float("nan"), f"host replay failed: {type(e).__name__}: {e}"[:200], layout)
+            return out
+        finally:
+            _ACTIVE[0] = True
+        worst, where = 0.0, ""
+        for (n, d), (_, h) in zip(_tensors(out, "out"), _tensors(h_out, "out")):
+            r = _rel(d, h)
+            if not (r <= worst):
+                worst, where = r, n
+        for (n, d), (_, h) in zip(_tensors(args, "a") + _tensors(kwargs, "k"), _tensors(h_args, "a") + _tensors(h_kwargs, "k")):
+            if d.device.type == "cpu":
+                continue
+            r = _rel(d, h)  # in-place written arguments
+            if not (r <= worst):
+                worst, where = r, f"inplace {n}"
+        _record(name, worst, where, layout)
+        return out
+
+    wrapper._device_diff = True
+    return wrapper
+
+
+def wrap_factory(name: str, factory: Any) -> Any:
+    """tilelang_ops: ``xxx_tl(static...)`` returns the kernel; wrap the kernel it returns."""
+
+    @functools.wraps(factory)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        return wrap(name, factory(*args, **kwargs))
+
+    return wrapper
+
+
+def install() -> None:
+    import deep_gemm
+    import deep_select.interface
+    import flash_mla
+    import flash_mla.flash_mla_interface
+    import tile_kernels
+    import tile_kernels.engram
+    import tile_kernels.moe
+    import tile_kernels.quant
+    import tile_kernels.transform
+
+    replaced: dict[int, Any] = {}
+
+    def patch_module(mod: types.ModuleType, prefix: str) -> None:
+        for attr, val in list(vars(mod).items()):
+            if attr.startswith("_") or not callable(val) or isinstance(val, type):
+                continue
+            if getattr(val, "__module__", "") is None or not str(getattr(val, "__module__", "")).startswith(prefix.split(".")[0]):
+                continue
+            w = replaced.get(id(val)) or wrap(f"{prefix}.{attr}", val)
+            replaced[id(val)] = w
+            setattr(mod, attr, w)
+
+    for mod, prefix in [
+        (deep_gemm, "deep_gemm"),
+        (flash_mla, "flash_mla"),
+        (flash_mla.flash_mla_interface, "flash_mla"),
+        (deep_select.interface, "deep_select"),
+        (tile_kernels.quant, "tile_kernels.quant"),
+        (tile_kernels.transform, "tile_kernels.transform"),
+        (tile_kernels.moe, "tile_kernels.moe"),
+        (tile_kernels.engram, "tile_kernels.engram"),
+    ]:
+        patch_module(mod, prefix)
+    for sub in ("expand_kernel", "head_compute_mix_kernel", "norm_fn_kernel", "post_kernel", "pre_apply_mix_kernel",
+                "pre_split_mixes_kernel", "sinkhorn_kernel", "pre_big_fuse_kernel"):
+        m = __import__(f"tile_kernels.mhc.{sub}", fromlist=["x"])
+        patch_module(m, "tile_kernels.mhc")
+
+    from deepinfer.model_executor.layers.tilelang_ops import ops
+
+    for attr in [a for a in dir(ops) if a.endswith("_tl")]:
+        setattr(ops, attr, wrap_factory(f"tl.{attr}", getattr(ops, attr)))
+
+    # Model modules bound some kernels with `from X import f`: rebind those names too.
+    import importlib
+
+    for modname in ["deepinfer.model_executor.models.v4_modules.mhc_ops", "deepinfer.model_executor.models.v4_modules.mhc_modeling",
+                    "deepinfer.model_executor.models.v4_modules.mlp", "deepinfer.model_executor.models.v4_modules.hc",
+                    "deepinfer.model_executor.models.v4_modules.attn", "deepinfer.model_executor.models.v4_modules.index_attn",
+                    "deepinfer.model_executor.layers.attention_backend", "deepinfer.model_executor.layers.rotary_embedding",
+                    "deepinfer.model_executor.layers.sampler", "deepinfer.model_executor.layers.sampler_utils"]:
+        mod = importlib.import_module(modname)
+        for attr, val in list(vars(mod).items()):
+            if id(val) in replaced:
+                setattr(mod, attr, replaced[id(val)])
+
+
+def report() -> None:
+    print("\n[device_diff] kernel                                         calls   worst-rel  where / layout of worst call")
+    for name, st in sorted(_STATS.items(), key=lambda kv: -(kv[1]["worst"] if kv[1]["worst"] == kv[1]["worst"] else 1e9)):
+        print(f"[device_diff] {name:48s} {st['calls']:6d}  {st['worst']:9.2e}  {st['where']}  {st['layout']}")
+
+
+def main() -> None:
+    import run_deepinfer_tiny as R
+
+    argv = sys.argv[1:]
+    if "--device" not in argv:
+        argv = ["--device", "neuron", *argv]
+    R.setup_env("neuron")
+    import deepinfer
+
+    deepinfer.PLATFORM  # noqa: B018 - initialize the platform (device mapping) before patching
+    install()
+    sys.argv = [sys.argv[0], *argv]
+    try:
+        R.main()
+    finally:
+        report()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/examples/neuron/probe_layer0.py b/examples/neuron/probe_layer0.py
index 3a8a493..acf1229 100644
--- a/examples/neuron/probe_layer0.py
+++ b/examples/neuron/probe_layer0.py
@@ -291,6 +291,10 @@ def run_deepinfer(args):
         print("[probe] using P-bf16 sparse attention (proposed fix)")
     wrap(mlp_mod, "moe_topk_gate", lambda o, a, k: (cap.put("topk_idx", o[0].float()), cap.put("topk_weights", o[1])))
 
+    # Fixed KV pool: the engine's profile run (a max_num_batched_tokens dummy prefill) is slow on the host and
+    # exhausts HBM on the NeuronCore with the pure-torch kernels.
+    R.install_workarounds(32)
+
     from deepinfer.entrypoints.llm import LLM
     from deepinfer.sampling_params import SamplingParams
 
diff --git a/tests/neuron/deep_ep/refs.py b/tests/neuron/deep_ep/refs.py
new file mode 100644
index 0000000..722fafa
--- /dev/null
+++ b/tests/neuron/deep_ep/refs.py
@@ -0,0 +1,250 @@
+import math
+import torch
+import torch.distributed as dist
+from typing import Optional, Tuple, Union
+
+def get_global_seed() -> int:
+    return 0
+
+
+def dispatch(x: Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]],
+             topk_idx: torch.Tensor, topk_weights: Optional[torch.Tensor],
+             num_max_tokens_per_rank: int, num_experts: int):
+    """
+    The reference implementation of dispatching tokens to experts across multiple ranks.
+
+    Not expanded. Sorted by rank and then by token within each rank (i.e. sorted by `src_token_global_idx`).
+
+    Arguments:
+        x: input tokens, `[num_tokens, hidden]` or (`[num_tokens, hidden]`, `[num_tokens, num_sf_packs]`).
+        topk_idx: top-k expert indices for each token, `[num_tokens, num_topk]`.
+        topk_weights: top-k weights for each token, can be None, `[num_tokens, num_topk]`.
+        num_max_tokens_per_rank: the maximum number of tokens per rank, must be `>=` the actual number of
+            tokens per rank and aligned with DeepEP's `num_max_tokens_per_rank`, since we calculate
+            `src_token_global_idx = src_rank_idx * num_max_tokens_per_rank + src_token_local_idx`.
+        num_experts: the total number of experts across all ranks.
+
+    Returns:
+        recv_x, recv_topk_idx, recv_topk_weights: received tokens, top-k indices and top-k weights after
+            dispatching. Out of range `recv_topk_idx` (i.e. that expert is not on the current rank) are
+            set to -1.
+        recv_src_token_idx: received `src_token_global_idx` for each received token.
+        num_recv_tokens_per_rank: number of received tokens from each rank, `[num_ranks]`.
+    """
+    rank_idx = dist.get_rank()
+    num_ranks = dist.get_world_size()
+
+    assert num_experts % num_ranks == 0
+    num_experts_per_rank = num_experts // num_ranks
+
+    # Unpack SF
+    use_fp8 = isinstance(x, tuple)
+    x, sf = x if use_fp8 else (x, None)
+
+    num_tokens, hidden = x.size()
+    num_tokens_, num_topk = topk_idx.size()
+    assert num_tokens == num_tokens_
+    num_sf_packs = 0
+    if sf is not None:
+        num_tokens__, num_sf_packs = sf.size()
+        assert num_tokens == num_tokens__
+    if topk_weights is not None:
+        num_tokens__, num_topk_ = topk_weights.size()
+        assert num_tokens == num_tokens__
+        assert num_topk == num_topk_
+
+    # Prepare per-peer send buffers
+    send_x_list = []
+    send_sf_list = []
+    send_topk_idx_list = []
+    send_topk_weights_list = []
+    send_src_token_idx_list = []
+    num_send_tokens_per_rank = torch.zeros((num_ranks, ), dtype=torch.int, device=x.device)
+    for dst_rank_idx in range(num_ranks):
+        expert_start_idx = dst_rank_idx * num_experts_per_rank
+        expert_end_idx = expert_start_idx + num_experts_per_rank
+
+        # Get the indices of tokens
+        mask_to_send = ((expert_start_idx <= topk_idx) & (topk_idx < expert_end_idx)).any(dim=1)
+        indices_to_send = mask_to_send.nonzero(as_tuple=True)[0]
+        num_send_tokens_per_rank[dst_rank_idx] = indices_to_send.numel()
+
+        # Select the data for tokens
+        x_to_send = x[indices_to_send]
+        sf_to_send = sf[indices_to_send] if use_fp8 else None
+        topk_idx_to_send = topk_idx[indices_to_send]
+        topk_weights_to_send = topk_weights[indices_to_send]
+        masked_topk_idx = torch.where((expert_start_idx <= topk_idx_to_send) & (topk_idx_to_send < expert_end_idx),
+                                      topk_idx_to_send, torch.full_like(topk_idx_to_send, -1))
+
+        send_x_list.append(x_to_send)
+        send_sf_list.append(sf_to_send)
+        send_topk_idx_list.append(masked_topk_idx)
+        send_topk_weights_list.append(topk_weights_to_send)
+        send_src_token_idx_list.append(indices_to_send)
+
+    send_x = torch.cat(send_x_list, dim=0)
+    send_sf = torch.cat(send_sf_list, dim=0) if use_fp8 else None
+    send_topk_idx = torch.cat(send_topk_idx_list, dim=0)
+    send_topk_weights = torch.cat(send_topk_weights_list, dim=0)
+    send_src_token_idx = torch.cat(send_src_token_idx_list, dim=0).to(torch.int)
+    send_src_token_idx += rank_idx * num_max_tokens_per_rank
+
+    # Exchange size
+    num_recv_tokens_per_rank = torch.empty((num_ranks, ), dtype=torch.int, device=x.device)
+    dist.all_to_all_single(num_recv_tokens_per_rank, num_send_tokens_per_rank)
+    num_recv_tokens = int(num_recv_tokens_per_rank.sum().item())
+
+    # Exchange main data
+    num_send_tokens_per_rank = num_send_tokens_per_rank.tolist()
+    num_recv_tokens_per_rank = num_recv_tokens_per_rank.tolist()
+    recv_x = torch.empty((num_recv_tokens, hidden), dtype=x.dtype, device=x.device)
+    recv_sf = torch.empty((num_recv_tokens, num_sf_packs), dtype=sf.dtype, device=x.device) if use_fp8 else None
+    recv_topk_idx = torch.empty((num_recv_tokens, num_topk), dtype=topk_idx.dtype, device=x.device)
+    recv_topk_weights = torch.empty((num_recv_tokens, num_topk), dtype=topk_weights.dtype, device=x.device)
+    recv_src_token_idx = torch.empty((num_recv_tokens, ), dtype=torch.int, device=x.device)
+    dist.all_to_all_single(recv_x, send_x, num_recv_tokens_per_rank, num_send_tokens_per_rank)
+    if use_fp8:
+        # gloo has no int16 collective: move the SF bytes (bit-exact, same semantics).
+        dist.all_to_all_single(recv_sf.view(torch.uint8), send_sf.contiguous().view(torch.uint8),
+                               num_recv_tokens_per_rank, num_send_tokens_per_rank)
+    dist.all_to_all_single(recv_topk_idx, send_topk_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank)
+    dist.all_to_all_single(recv_topk_weights, send_topk_weights, num_recv_tokens_per_rank, num_send_tokens_per_rank)
+    dist.all_to_all_single(recv_src_token_idx, send_src_token_idx, num_recv_tokens_per_rank, num_send_tokens_per_rank)
+
+    # Mask top-k indices
+    expert_start_idx = rank_idx * num_experts_per_rank
+    expert_end_idx = expert_start_idx + num_experts_per_rank
+    mask = (expert_start_idx <= recv_topk_idx) & (recv_topk_idx < expert_end_idx)
+    recv_topk_idx = recv_topk_idx - expert_start_idx
+    recv_topk_idx.masked_fill_(~mask, -1)
+
+    # Pack SF
+    recv_x = (recv_x, recv_sf) if use_fp8 else recv_x
+
+    return (recv_x, recv_topk_idx, recv_topk_weights,
+            recv_src_token_idx, torch.tensor(num_recv_tokens_per_rank, dtype=torch.int))
+
+
+def generate_pre_combine_data(src_token_global_idx: torch.Tensor,
+                              num_max_tokens_per_rank: int, num_topk: int, hidden: int) -> torch.Tensor:
+    """
+    Generate data needed for combine from `src_token_global_idx`.
+    Recall that `src_token_global_idx = src_rank_idx * num_max_tokens_per_rank + src_token_local_idx`.
+    The generated data (denoted as `y`) of the i-th token has a shape of `[num_topk, hidden]`, with
+
+    `y[j, k] = sin((token_seeds * P % max_seed + 1) / max_seed * (k + 1) + sin(seed))`
+
+    where `P=131071` is a large prime, `token_seeds` is calculated via
+    `token_seeds = src_token_global_idx[i] * num_topk + j`, and
+    `max_seed = num_ranks * num_max_tokens_per_rank * num_topk`.
+
+    Arguments:
+        src_token_global_idx: source token global indices, `[num_tokens]`.
+
+    Returns:
+        the generated data, `[num_tokens, num_topk, hidden]`.
+    """
+    num_ranks = dist.get_world_size()
+    token_seeds = (src_token_global_idx.unsqueeze(1) * num_topk +
+                   torch.arange(num_topk, device=src_token_global_idx.device).unsqueeze(0))  # [num_tokens, num_topk]
+    max_seed = num_ranks * num_max_tokens_per_rank * num_topk
+    result = torch.sin(
+        (((token_seeds * 131071 % max_seed).float() + 1) / max_seed).unsqueeze(-1) *
+        torch.arange(1, hidden + 1, device=src_token_global_idx.device, dtype=torch.float32).broadcast_to(1, 1, hidden) +
+        math.sin(float(get_global_seed()))
+    )
+    return result.to(torch.bfloat16)
+
+
+def ordered_accumulate(data: torch.Tensor, initial_value: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """
+    Accumulate `data` in order along the num_topk dimension.
+
+    Arguments:
+        data: data to be accumulated, `[num_tokens, num_topk, hidden]`.
+        initial_value: initial value for accumulation, `[num_tokens, hidden]`.
+
+    Returns:
+        the result, `[num_tokens, hidden]`.
+    """
+    num_topk = data.shape[1]
+    if initial_value is None:
+        result = torch.zeros((data.shape[0], data.shape[2]), dtype=torch.float32, device=data.device)
+    else:
+        result = initial_value.clone()
+    for i in range(num_topk):
+        result += data[:, i, :].float()
+    return result.to(data.dtype)
+
+
+def combine(y: torch.Tensor, topk_idx: torch.Tensor, num_experts: int,
+            bias: Optional[Union[Tuple[torch.Tensor, ...], torch.Tensor]],
+            reduce_in_local: bool) -> torch.Tensor:
+    """
+    The reference implementation of (possibly intra-rank reduced) combining tokens.
+
+    Arguments:
+        y: input tokens to be combined, `[num_tokens, num_topk, hidden]`.
+        topk_idx: `[num_tokens, num_topk]`.
+        num_experts: the total number of experts across all ranks.
+        bias: 0, 1 or 2 `[num_tokens, hidden]` final bias added to the output.
+        reduce_in_local: whether to do reduction within each rank first.
+            - `True`  -> reduce selections landing on the same rank before the global accumulation.
+            - `False` -> equivalent to `allow_multiple_reduction` being False.
+
+    Returns:
+        the combined result, `[num_tokens, hidden]`.
+    """
+    num_ranks = dist.get_world_size()
+    num_tokens, hidden = y.shape[0], y.shape[2]
+    num_topk = y.shape[1]
+
+    def grouped_reduce(data_to_reduce: torch.Tensor, group_id: torch.Tensor) -> torch.Tensor:
+        """
+        Perform in-place grouped reduction on `data_to_reduce` according to `group_id`.
+        The summation within each group is performed in strict order along the `num_topk` dimension.
+        The result for each group is stored at the rightmost token of that group, and other tokens are
+        set to zero.
+
+        Arguments:
+            data_to_reduce: data to be reduced, `[num_tokens, num_topk, hidden]`.
+            group_id: group IDs for each token, `[num_tokens, num_topk]`.
+        """
+        # Shuffle to make tokens with the same group_id contiguous
+        group_id, src_indices = torch.sort(group_id, dim=-1, stable=True)
+        # transformed_src_indices[i, j] = i * num_topk + src_indices[i, j]
+        transformed_src_indices = (
+            (src_indices + torch.arange(0, num_tokens, device=y.device).unsqueeze(-1) * num_topk).flatten())
+        data_to_reduce = data_to_reduce.view(-1, hidden)[transformed_src_indices].view(num_tokens, num_topk, hidden)
+        # Perform segmented reduce within each group
+        cur_accum_buf = torch.zeros((num_tokens, hidden), dtype=torch.float32, device=y.device)
+        for i in range(num_topk):
+            is_segment_break = torch.full((num_tokens, ), True, dtype=torch.bool, device=y.device) \
+                if i == num_topk - 1 else group_id[:, i] != group_id[:, i + 1]
+            cur_accum_buf += data_to_reduce[:, i, :].float()
+            # For one token, if `is_segment_break` is True,
+            # save the accumulated value and clear the buffer, otherwise, clear `data_to_reduce[:, i, :]`
+            segment_break_token_indices = torch.where(is_segment_break)[0]
+            data_to_reduce[segment_break_token_indices, i] = cur_accum_buf[segment_break_token_indices].to(data_to_reduce.dtype)
+            cur_accum_buf[segment_break_token_indices] = 0.0
+            non_segment_break_token_indices = torch.where(~is_segment_break)[0]
+            data_to_reduce[non_segment_break_token_indices, i] = 0.0
+        # Unshuffle
+        # noinspection PyShadowingNames
+        result = torch.empty_like(data_to_reduce)
+        result.view(-1, hidden)[transformed_src_indices] = data_to_reduce.view(-1, hidden)
+        return result.view(num_tokens, num_topk, hidden)
+
+    num_experts_per_rank = num_experts // num_ranks
+    if reduce_in_local:
+        y = grouped_reduce(y, topk_idx // num_experts_per_rank)
+    bias_sum = None
+    if isinstance(bias, tuple):
+        assert len(bias) <= 2
+        for bias_tensor in bias:
+            bias_sum = bias_tensor.float() if bias_sum is None else bias_sum + bias_tensor.float()
+    elif bias is not None:
+        bias_sum = bias.float()
+    result = ordered_accumulate(y, bias_sum)
+    return result
diff --git a/tests/neuron/deep_ep/test_bucket.py b/tests/neuron/deep_ep/test_bucket.py
new file mode 100644
index 0000000..09d0c82
--- /dev/null
+++ b/tests/neuron/deep_ep/test_bucket.py
@@ -0,0 +1,103 @@
+"""BucketBuffer / BucketSession checks (subset of DeepEP-Ascend tests/bucket/*, CPU + gloo).
+
+Covers what `model_executor/parallel_utils/cp.py` uses: a session, `session.allocate`,
+in-place and out-of-place `all_gather`, and the NotImplementedError stubs.
+"""
+
+import argparse
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
+import deep_ep  # noqa: E402
+
+# DEEP_EP_TEST_DEVICE=neuron runs every rank on its own NeuronCore (core DEEP_EP_TEST_CORE_BASE + rank)
+# with a "neuron" process group, exercising the real torch-neuronx collectives; default: CPU + gloo.
+DEVICE = os.environ.get('DEEP_EP_TEST_DEVICE', 'cpu')
+
+
+def worker(local_rank: int, num_ranks: int, args):
+    os.environ.setdefault('MASTER_ADDR', '127.0.0.1')
+    os.environ.setdefault('MASTER_PORT', str(args.port))
+    if DEVICE == 'neuron':
+        # Must be set before the Neuron runtime initializes (first device op) in this process.
+        os.environ['NEURON_RT_VISIBLE_CORES'] = str(int(os.environ.get('DEEP_EP_TEST_CORE_BASE', '0')) + local_rank)
+    dist.init_process_group('gloo', rank=local_rank, world_size=num_ranks)
+    group = dist.new_group(list(range(num_ranks)), backend='neuron') if DEVICE == 'neuron' else dist.group.WORLD
+    rank = group.rank()
+
+    alignment = deep_ep.get_num_allocation_alignment()
+    assert alignment == 2 << 20 and deep_ep.get_num_rdma_alignment() == 512
+
+    # Plan-based construction (upstream BufferAllocator path).
+    plan = deep_ep.BufferAllocator()
+    gathered = plan.allocate((num_ranks, 2, 8), torch.float32)
+    buffer = deep_ep.BucketBuffer(group, plan, explicitly_destroy=True, device=torch.device(DEVICE))
+    assert buffer.num_bytes == plan.num_bytes
+
+    src = torch.arange(16, dtype=torch.float32, device=DEVICE).view(2, 8) + rank * 100
+    ref = torch.stack([src - rank * 100 + r * 100 for r in range(num_ranks)])
+
+    # Out-of-place: registered destination.
+    gathered.fill_(float('nan'))
+    result = buffer.all_gather(src, dsts=gathered, group=group).wait()
+    assert result.ndim == 1 and torch.equal(result, ref.reshape(-1))
+
+    # In-place: source is this rank's shard of a registered region.
+    gathered.fill_(float('nan'))
+    gathered[rank].copy_(src)
+    result = buffer.all_gather(gathered[rank], group=group).wait()
+    assert torch.equal(result, ref.reshape(-1))
+    buffer.barrier(group=group)
+    buffer.destroy()
+
+    # Session path (what CPManager uses).
+    buffer = deep_ep.BucketBuffer(group, 4 * alignment, explicitly_destroy=True,
+                                 device=torch.device(DEVICE))
+    with buffer.session() as session:
+        region = session.allocate((num_ranks, 2, 8), torch.float32, group=group)
+        shard = region[rank]
+        shard.copy_(src)
+        handle = buffer.all_gather(shard)
+        out = handle.wait()
+        assert torch.equal(out.view(num_ranks, 2, 8).flatten(0, 1), ref.reshape(-1, 8))
+
+        # Unregistered source is staged into the session storage.
+        unregistered = (src + 1000).contiguous()
+        out = buffer.all_gather(unregistered).wait()
+        ref2 = torch.stack([src - rank * 100 + r * 100 + 1000 for r in range(num_ranks)])
+        assert torch.equal(out, ref2.reshape(-1))
+
+        # Batched
+        outs = buffer.all_gather([shard, shard]).wait()
+        assert len(outs) == 2 and all(torch.equal(o, ref.reshape(-1)) for o in outs)
+
+        for name in ('reduce_scatter', 'all_reduce'):
+            try:
+                getattr(buffer, name)(shard)
+            except NotImplementedError:
+                pass
+            else:
+                raise AssertionError(f'{name} should raise NotImplementedError')
+    assert 'all_gather' not in buffer.__dict__
+    buffer.destroy()
+
+    if rank == 0:
+        print(f'Bucket tests passed (world_size={num_ranks})', flush=True)
+    dist.destroy_process_group()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--num-processes', type=int, default=2)
+    parser.add_argument('--port', type=int, default=29551)
+    args = parser.parse_args()
+    if args.num_processes == 1:
+        worker(0, 1, args)
+    else:
+        mp.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes)
diff --git a/tests/neuron/deep_ep/test_engram.py b/tests/neuron/deep_ep/test_engram.py
new file mode 100644
index 0000000..53574a8
--- /dev/null
+++ b/tests/neuron/deep_ep/test_engram.py
@@ -0,0 +1,131 @@
+"""EngramBuffer correctness, ported from DeepEP-Ascend tests/engram/test_engram.py (CPU + gloo)."""
+
+import argparse
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
+import deep_ep  # noqa: E402
+from deep_ep.utils.math import ceil_div, per_token_cast_to_fp8  # noqa: E402
+
+# DEEP_EP_TEST_DEVICE=neuron runs every rank on its own NeuronCore (core DEEP_EP_TEST_CORE_BASE + rank)
+# with a "neuron" process group, exercising the real torch-neuronx collectives; default: CPU + gloo.
+DEVICE = os.environ.get('DEEP_EP_TEST_DEVICE', 'cpu')
+
+
+def run(rank: int, num_ranks: int, group, args, use_fp8: bool):
+    dtype = torch.float8_e4m3fn if use_fp8 else torch.bfloat16
+    num_entries_per_layer = [int(v) for v in args.num_entries_per_layer.split(',')]
+    num_layers = len(num_entries_per_layer)
+    num_sf_packs = 2 * ceil_div(args.hidden, 128) if use_fp8 else 0
+
+    num_bytes, num_rdma_bytes = deep_ep.EngramBuffer.get_storage_size_hint(
+        group, num_entries_per_layer, args.hidden, args.num_tokens,
+        args.num_entries_per_token, dtype, num_sf_packs)
+    num_qps, qp_depth, restrict = deep_ep.EngramBuffer.get_theoretical_config(
+        group, num_layers, args.num_tokens, args.num_entries_per_token)
+    assert (num_qps, qp_depth, restrict) == (0, 0, False)
+
+    buffer = deep_ep.EngramBuffer(
+        group, num_gpu_bytes=num_bytes, num_rdma_storage_size=num_rdma_bytes,
+        use_cpu_rdma_storage=True, num_allocated_qps=num_qps, qp_depth=qp_depth,
+        restrict_rd_atomic=restrict, allow_hybrid_mode=True, explicitly_destroy=True)
+    buffer.set_config(num_entries_per_layer, args.hidden, args.num_tokens,
+                      args.num_entries_per_token, dtype, num_sf_packs)
+
+    torch.manual_seed(7 + rank)
+    local_storages, sfs = [], ([] if use_fp8 else None)
+    for num_entries in num_entries_per_layer:
+        local_bf16 = torch.randn((num_entries, args.hidden), dtype=torch.bfloat16, device=DEVICE)
+        if use_fp8:
+            local_storage, local_sf = per_token_cast_to_fp8(local_bf16)
+            local_sf = local_sf.view(torch.int16).contiguous()
+            sf = torch.empty((num_ranks * num_entries, num_sf_packs), dtype=local_sf.dtype, device=DEVICE)
+            if num_ranks == 1:
+                sf.copy_(local_sf)
+            else:
+                # gloo has no int16 collective; gather the bytes instead (DeepInfer's
+                # NgramEmbedding.gather_sf gathers int16 directly on the accelerator).
+                dist.all_gather_into_tensor(sf.view(torch.uint8), local_sf.view(torch.uint8), group=group)
+            sfs.append(sf)
+        else:
+            local_storage = local_bf16
+        local_storages.append(local_storage.contiguous())
+    buffer.write(local_storages, sfs=sfs)
+
+    # Reference global tables
+    global_bits = []
+    for local_storage in local_storages:
+        bits = torch.empty((num_ranks * local_storage.shape[0], args.hidden * dtype.itemsize),
+                           dtype=torch.uint8, device=DEVICE)
+        local_bits = local_storage.view(torch.uint8).view(local_storage.shape[0], -1).contiguous()
+        if num_ranks == 1:
+            bits.copy_(local_bits)
+        else:
+            dist.all_gather_into_tensor(bits, local_bits, group=group)
+        global_bits.append(bits)
+
+    def generate_indices():
+        return torch.stack([
+            torch.randint(0, num_ranks * num_entries, (args.num_tokens, args.num_entries_per_token),
+                          dtype=torch.int, device=DEVICE)
+            for num_entries in num_entries_per_layer])
+
+    for col_major in ((False, True) if use_fp8 else (False,)):
+        for layer_order in (range(num_layers), reversed(range(num_layers))):
+            indices = generate_indices()
+            ref_data = [global_bits[l][indices[l].reshape(-1).long()].reshape(args.num_tokens, -1)
+                        for l in range(num_layers)]
+            ref_sfs = [sfs[l][indices[l].reshape(-1).long()].reshape(args.num_tokens, -1)
+                       for l in range(num_layers)] if use_fp8 else None
+
+            hooks = buffer.fetch(indices, use_tma_aligned_col_major_sf=col_major)
+            assert len(hooks) == num_layers
+            for l in layer_order:
+                fetched = hooks[l]()
+                data, fetched_sf = fetched if use_fp8 else (fetched, None)
+                assert data.shape == (args.num_tokens, args.num_entries_per_token * args.hidden)
+                assert torch.equal(ref_data[l], data.reshape(args.num_tokens, -1).contiguous().view(torch.uint8)), \
+                    f'data mismatch (layer {l}, fp8={use_fp8}, col_major={col_major})'
+                if use_fp8:
+                    expected_stride = (1, args.num_tokens) if col_major else (fetched_sf.size(1), 1)
+                    assert fetched_sf.stride() == expected_stride, (fetched_sf.stride(), expected_stride)
+                    assert torch.equal(ref_sfs[l], fetched_sf), f'sf mismatch (layer {l})'
+    buffer.destroy()
+
+
+def worker(local_rank: int, num_ranks: int, args):
+    os.environ.setdefault('MASTER_ADDR', '127.0.0.1')
+    os.environ.setdefault('MASTER_PORT', str(args.port))
+    if DEVICE == 'neuron':
+        # Must be set before the Neuron runtime initializes (first device op) in this process.
+        os.environ['NEURON_RT_VISIBLE_CORES'] = str(int(os.environ.get('DEEP_EP_TEST_CORE_BASE', '0')) + local_rank)
+    dist.init_process_group('gloo', rank=local_rank, world_size=num_ranks)
+    group = dist.new_group(list(range(num_ranks)), backend='neuron') if DEVICE == 'neuron' else dist.group.WORLD
+    for use_fp8 in (False, True):
+        run(group.rank(), num_ranks, group, args, use_fp8)
+        if group.rank() == 0:
+            print(f'  ok: fp8={use_fp8}', flush=True)
+    if group.rank() == 0:
+        print(f'Engram tests passed (world_size={num_ranks})', flush=True)
+    dist.destroy_process_group()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--num-processes', type=int, default=2)
+    parser.add_argument('--num-entries-per-layer', type=str, default='37,64')
+    parser.add_argument('--hidden', type=int, default=128)
+    parser.add_argument('--num-tokens', type=int, default=11)
+    parser.add_argument('--num-entries-per-token', type=int, default=3)
+    parser.add_argument('--port', type=int, default=29541)
+    args = parser.parse_args()
+    if args.num_processes == 1:
+        worker(0, 1, args)
+    else:
+        mp.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes)
diff --git a/tests/neuron/deep_ep/test_ep.py b/tests/neuron/deep_ep/test_ep.py
new file mode 100644
index 0000000..b72bdc3
--- /dev/null
+++ b/tests/neuron/deep_ep/test_ep.py
@@ -0,0 +1,301 @@
+"""EPBuffer dispatch/combine correctness, ported from DeepEP-Ascend tests/ep/test_ep.py.
+
+CPU + gloo, no Neuron device. Run: python test_ep.py --num-processes {1,2,4}
+"""
+
+import argparse
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
+import deep_ep  # noqa: E402
+from deep_ep.utils.math import align, per_token_cast_to_fp8  # noqa: E402
+from refs import combine as ref_combine  # noqa: E402
+from refs import dispatch as ref_dispatch  # noqa: E402
+from refs import generate_pre_combine_data  # noqa: E402
+
+# DEEP_EP_TEST_DEVICE=neuron runs every rank on its own NeuronCore (core DEEP_EP_TEST_CORE_BASE + rank)
+# with a "neuron" process group, exercising the real torch-neuronx collectives; default: CPU + gloo.
+DEVICE = os.environ.get('DEEP_EP_TEST_DEVICE', 'cpu')
+
+
+def enumerate_modes():
+    for expert_alignment in (128, 1):
+        for use_fp8 in (True, False):
+            for num_bias in (0, 1, 2):
+                for col_major in ((False, True) if use_fp8 else (False,)):
+                    yield expert_alignment, num_bias, use_fp8, col_major
+
+
+def launch(buffer, name, defer_epilogue, **params):
+    result = getattr(buffer, name)(defer_epilogue=defer_epilogue, **params)
+    if defer_epilogue:
+        event = result
+        return *event.current_stream_wait(release_handle=True), event
+    result[-1].current_stream_wait(release_handle=True)
+    return result
+
+
+def check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync):
+    expert_alignment, num_bias, use_fp8, col_major = mode
+    tag = f'mode={mode} defer={defer_epilogue} cpu_sync={do_cpu_sync}'
+    cap = args.num_tokens
+    num_tokens = max(1, args.num_tokens - rank)
+    hidden, num_topk, num_experts = args.hidden, args.num_topk, args.num_experts
+    num_local_experts = num_experts // num_ranks
+
+    torch.manual_seed(1234 + rank)
+    scores = torch.randn((num_tokens, num_experts), dtype=torch.float32, device=DEVICE)
+    topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False)
+    topk_idx = topk_idx.to(torch.int64)
+    if args.masked_ratio > 0:
+        rand_mask = torch.rand_like(topk_idx, dtype=torch.float)
+        topk_idx = topk_idx.masked_fill(rand_mask < args.masked_ratio, -1)
+        topk_weights = topk_weights.masked_fill(topk_idx < 0, 0)
+    topk_idx, topk_weights = topk_idx.contiguous(), topk_weights.contiguous()
+
+    x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=DEVICE)
+    if use_fp8:
+        x_data, sf = per_token_cast_to_fp8(x)
+        sf = sf.view(torch.int16)
+        num_sf_packs = sf.size(1)
+        if col_major:
+            sf_storage = torch.empty((num_tokens * num_sf_packs,), dtype=torch.int16, device=DEVICE)
+            input_sf = torch.as_strided(sf_storage, (num_tokens, num_sf_packs), (1, num_tokens))
+            input_sf.copy_(sf)
+        else:
+            input_sf = sf
+        x_in = (x_data, input_sf)
+    else:
+        x_in = x
+    bias = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=DEVICE) if num_bias == 1 else None
+    if num_bias == 2:
+        bias = tuple(torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=DEVICE) for _ in range(2))
+
+    stats = torch.zeros((num_local_experts,), dtype=torch.int32, device=DEVICE)
+    dispatch_args = dict(
+        x=x_in, topk_idx=topk_idx, topk_weights=topk_weights,
+        num_sms=7, num_max_tokens_per_rank=cap, num_experts=num_experts,
+        expert_alignment=expert_alignment, async_with_compute_stream=True,
+        do_handle_copy=False, do_cpu_sync=do_cpu_sync, do_expand=True,
+        do_zero_padding=True, use_tma_aligned_col_major_sf=col_major)
+
+    ref_x = (x_data.view(torch.uint8), input_sf) if use_fp8 else x
+    (ref_recv_x, ref_recv_topk_idx, ref_recv_topk_weights,
+     ref_recv_src_token_idx, ref_num_recv_tokens_per_rank) = ref_dispatch(
+        ref_x, topk_idx, topk_weights, cap, num_experts)
+
+    recv_x, recv_topk_idx, recv_weights, handle, _ = launch(
+        buffer, 'dispatch', defer_epilogue, cumulative_local_expert_recv_stats=stats, **dispatch_args)
+    assert recv_topk_idx is None
+    assert recv_weights is not None
+    assert handle.num_sms == 7
+
+    recv_x_data, recv_sf = recv_x if use_fp8 else (recv_x, None)
+    check_x = recv_x_data.view(torch.uint8) if use_fp8 else recv_x_data
+
+    num_recv_tokens = handle.psum_num_recv_tokens_per_rank[-1].item()
+    num_expanded_tokens = align(handle.psum_num_recv_tokens_per_expert[-1].item(), expert_alignment)
+    expected_num_recv = num_recv_tokens if do_cpu_sync else num_ranks * cap
+    assert handle.num_recv_tokens == expected_num_recv, (handle.num_recv_tokens, expected_num_recv)
+    assert handle.recv_src_metadata.size(0) == expected_num_recv
+    if use_fp8:
+        expected_stride = (1, recv_x_data.size(0)) if col_major else (num_sf_packs, 1)
+        assert recv_sf.stride() == expected_stride, (recv_sf.stride(), expected_stride)
+
+    ref_recv_x_data, ref_recv_sf = ref_recv_x if use_fp8 else (ref_recv_x, None)
+    assert num_recv_tokens == ref_recv_x_data.size(0), (num_recv_tokens, ref_recv_x_data.size(0))
+
+    metadata = handle.recv_src_metadata[:num_recv_tokens]
+    recv_src_token_idx = metadata[:, -2]
+    recv_expanded_slots = metadata[:, :-2]
+    ref_order = torch.argsort(ref_recv_src_token_idx)
+    metadata_order = torch.argsort(recv_src_token_idx)
+    assert torch.equal(recv_src_token_idx[metadata_order], ref_recv_src_token_idx[ref_order])
+
+    def reorder(src):
+        out = torch.empty_like(src)
+        out[metadata_order] = src[ref_order]
+        return out
+
+    ref_x_by_meta = reorder(ref_recv_x_data)
+    ref_sf_by_meta = reorder(ref_recv_sf) if ref_recv_sf is not None else None
+    ref_idx_by_meta = reorder(ref_recv_topk_idx)
+    ref_w_by_meta = reorder(ref_recv_topk_weights)
+
+    recv_valid_mask = recv_expanded_slots >= 0
+    assert torch.equal(recv_valid_mask, ref_idx_by_meta >= 0)
+
+    valid_slots = recv_expanded_slots[recv_valid_mask].long()
+    assert valid_slots.numel() == handle.num_unaligned_recv_tokens_per_expert.sum().item()
+    if valid_slots.numel() > 0:
+        assert valid_slots.min().item() >= 0 and valid_slots.max().item() < num_expanded_tokens
+        assert torch.unique(valid_slots).numel() == valid_slots.numel()
+        expected_x = ref_x_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask]
+        assert torch.equal(check_x[valid_slots], expected_x)
+        if ref_sf_by_meta is not None:
+            expected_sf = ref_sf_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask]
+            assert torch.equal(recv_sf[valid_slots], expected_sf)
+        assert torch.equal(recv_weights[valid_slots], ref_w_by_meta[recv_valid_mask])
+
+    # Master top-k slot column
+    topk_range = torch.arange(num_topk, device=DEVICE, dtype=metadata.dtype)
+    master = torch.where(recv_valid_mask, topk_range.view(1, num_topk),
+                         torch.full_like(recv_expanded_slots, -1)).max(dim=1).values
+    assert torch.equal(metadata[:, -1], master)
+
+    # Per-rank and per-expert counts
+    psum_rank = [0] + handle.psum_num_recv_tokens_per_rank.tolist()
+    for r in range(num_ranks):
+        assert psum_rank[r + 1] - psum_rank[r] == ref_num_recv_tokens_per_rank[r].item()
+    psum_expert = [0] + handle.psum_num_recv_tokens_per_expert.tolist()
+    for e in range(num_local_experts):
+        count = psum_expert[e + 1] - align(psum_expert[e], expert_alignment)
+        ref_count = (ref_recv_topk_idx == e).sum().item()
+        assert count == ref_count, (e, count, ref_count)
+        assert stats[e].item() == ref_count
+        assert handle.num_unaligned_recv_tokens_per_expert[e].item() == ref_count
+        if do_cpu_sync:
+            assert handle.num_recv_tokens_per_expert_list[e] == align(ref_count, expert_alignment)
+        expert_start = align(psum_expert[e], expert_alignment)
+        expert_end = psum_expert[e + 1]
+        expert_slots = recv_expanded_slots[ref_idx_by_meta == e]
+        assert expert_slots.numel() == ref_count
+        if ref_count > 0:
+            assert torch.equal(torch.sort(expert_slots).values,
+                               torch.arange(expert_start, expert_end, device=DEVICE,
+                                            dtype=expert_slots.dtype))
+        padding_end = align(expert_end, expert_alignment)
+        assert torch.count_nonzero(recv_x_data[expert_end:padding_end].view(torch.uint8)).item() == 0
+        if recv_sf is not None:
+            assert torch.count_nonzero(recv_sf[expert_end:padding_end]).item() == 0
+
+    # --- Combine ---
+    local_y = generate_pre_combine_data(metadata[:, -2], cap, num_topk, hidden)
+    input_for_combine = torch.empty((recv_x_data.size(0) + 1, hidden), dtype=torch.bfloat16, device=DEVICE)
+    input_for_combine[metadata[:, :-2].long().flatten()] = local_y.view(-1, hidden)
+    input_for_combine = input_for_combine[:-1].contiguous()
+
+    combined_x, combined_weights, _ = launch(
+        buffer, 'combine', defer_epilogue, x=input_for_combine, handle=handle,
+        topk_weights=recv_weights, bias=bias, num_sms=7, async_with_compute_stream=True)
+
+    ref_y = generate_pre_combine_data(rank * cap + torch.arange(num_tokens, device=DEVICE), cap, num_topk, hidden)
+    ref_y[topk_idx == -1] = 0
+    ref_combined_x = ref_combine(ref_y, topk_idx, num_experts, bias, reduce_in_local=True)
+    max_diff = (combined_x.float() - ref_combined_x.float()).abs().max().item()
+    assert torch.equal(combined_x, ref_combined_x), f'{tag} {max_diff=}'
+    assert torch.equal(combined_weights, topk_weights)
+
+    if num_bias == 0 and not col_major:
+        combined_x2, combined_weights2, _ = launch(
+            buffer, 'combine', defer_epilogue, x=input_for_combine, handle=handle,
+            bias=bias, num_sms=7, async_with_compute_stream=True)
+        assert torch.equal(combined_x2, ref_combined_x)
+        assert combined_weights2 is None
+
+    # --- Cached dispatch replay ---
+    cached_args = dict(dispatch_args, num_sms=0, do_cpu_sync=False, handle=handle)
+    del cached_args['topk_idx']
+    cached_tensors = (handle.psum_num_recv_tokens_per_rank, handle.psum_num_recv_tokens_per_expert,
+                      handle.num_unaligned_recv_tokens_per_expert, handle.recv_src_metadata)
+    cached_copies = [t.clone() for t in cached_tensors]
+    padding_mask = torch.ones(num_expanded_tokens, dtype=torch.bool, device=DEVICE)
+    padding_mask[valid_slots] = False
+    for send_weights in (True, False):
+        updated_x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device=DEVICE) * (8 if send_weights else 0.125)
+        if use_fp8:
+            updated_data, updated_sf = per_token_cast_to_fp8(updated_x)
+            updated_sf = updated_sf.view(torch.int16)
+            if col_major:
+                updated_sf = updated_sf.t().contiguous().t()
+            updated_in = (updated_data, updated_sf)
+        else:
+            updated_in = updated_x
+        updated_weights = torch.randn_like(topk_weights).masked_fill(topk_idx < 0, 0)
+        ref_updated = (updated_data.view(torch.uint8), updated_sf) if use_fp8 else updated_x
+        ref_cached_x, _, _, _, _ = ref_dispatch(ref_updated, topk_idx, updated_weights, cap, num_experts)
+        cached_x, cached_topk_idx, cached_w, cached_handle, _ = launch(
+            buffer, 'dispatch', defer_epilogue,
+            **dict(cached_args, x=updated_in,
+                   topk_weights=updated_weights if send_weights else None))
+        assert cached_handle is handle and cached_topk_idx is None
+        for tensor, saved in zip(cached_tensors, cached_copies):
+            assert torch.equal(tensor, saved)
+        cached_data, cached_sf = cached_x if use_fp8 else (cached_x, None)
+        ref_cached_data, ref_cached_sf = ref_cached_x if use_fp8 else (ref_cached_x, None)
+        cached_check = cached_data.view(torch.uint8) if use_fp8 else cached_data
+        ref_cached_by_meta = reorder(ref_cached_data)
+        assert torch.equal(cached_check[valid_slots],
+                           ref_cached_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask])
+        assert torch.count_nonzero(cached_check[:num_expanded_tokens][padding_mask]).item() == 0
+        if cached_sf is not None:
+            assert cached_sf.stride() == expected_stride
+            ref_cached_sf_by_meta = reorder(ref_cached_sf)
+            assert torch.equal(cached_sf[valid_slots],
+                               ref_cached_sf_by_meta.unsqueeze(1).expand(-1, num_topk, -1)[recv_valid_mask])
+        if send_weights:
+            assert torch.equal(cached_w[valid_slots], reorder(updated_weights_ref(
+                updated_weights, topk_idx, cap, num_experts))[recv_valid_mask])
+        else:
+            assert cached_w is None
+
+
+def updated_weights_ref(updated_weights, topk_idx, cap, num_experts):
+    _, _, ref_w, _, _ = ref_dispatch(
+        torch.zeros((topk_idx.size(0), 8), dtype=torch.bfloat16, device=DEVICE),
+        topk_idx, updated_weights, cap, num_experts)
+    return ref_w
+
+
+def worker(local_rank: int, num_ranks: int, args: argparse.Namespace):
+    os.environ.setdefault('MASTER_ADDR', '127.0.0.1')
+    os.environ.setdefault('MASTER_PORT', str(args.port))
+    if DEVICE == 'neuron':
+        # Must be set before the Neuron runtime initializes (first device op) in this process.
+        os.environ['NEURON_RT_VISIBLE_CORES'] = str(int(os.environ.get('DEEP_EP_TEST_CORE_BASE', '0')) + local_rank)
+    dist.init_process_group('gloo', rank=local_rank, world_size=num_ranks)
+    group = dist.new_group(list(range(num_ranks)), backend='neuron') if DEVICE == 'neuron' else dist.group.WORLD
+    rank = group.rank()
+
+    buffer = deep_ep.EPBuffer(
+        group, num_max_tokens_per_rank=args.num_tokens, hidden=args.hidden,
+        num_topk=args.num_topk, use_fp8_dispatch=True, allow_hybrid_mode=True,
+        allow_multiple_reduction=True, explicitly_destroy=True)
+    assert buffer.num_bytes > 0
+    assert deep_ep.get_physical_domain_size(group) == (1, num_ranks)
+
+    for mode in enumerate_modes():
+        for defer_epilogue in (True, False):
+            for do_cpu_sync in (True, False):
+                check_one(buffer, args, rank, num_ranks, mode, defer_epilogue, do_cpu_sync)
+        if rank == 0:
+            print(f'  ok: expert_alignment={mode[0]} num_bias={mode[1]} fp8={mode[2]} col_major={mode[3]}',
+                  flush=True)
+
+    buffer.destroy()
+    if rank == 0:
+        print(f'EP tests passed (world_size={num_ranks})', flush=True)
+    dist.destroy_process_group()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--num-processes', type=int, default=2)
+    parser.add_argument('--num-tokens', type=int, default=17)
+    parser.add_argument('--hidden', type=int, default=256)
+    parser.add_argument('--num-topk', type=int, default=4)
+    parser.add_argument('--num-experts', type=int, default=16)
+    parser.add_argument('--masked-ratio', type=float, default=0.2)
+    parser.add_argument('--port', type=int, default=29531)
+    args = parser.parse_args()
+    assert args.num_experts % args.num_processes == 0
+    if args.num_processes == 1:
+        worker(0, 1, args)
+    else:
+        mp.spawn(worker, args=(args.num_processes, args), nprocs=args.num_processes)
diff --git a/tests/neuron/deep_gemm/test_deep_gemm_neuron.py b/tests/neuron/deep_gemm/test_deep_gemm_neuron.py
new file mode 100644
index 0000000..03733db
--- /dev/null
+++ b/tests/neuron/deep_gemm/test_deep_gemm_neuron.py
@@ -0,0 +1,693 @@
+"""CPU correctness tests for the pure-torch deep_gemm port.
+
+Reference implementations and data generators are ported from
+DeepGEMM-Ascend/tests (generators.py, test_layout.py, test_attention.py,
+test_hyperconnection.py), adapted to CPU and small shapes.
+"""
+
+import math
+import os
+import sys
+import random
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
+
+import torch
+
+import deep_gemm
+
+torch.manual_seed(0)
+random.seed(0)
+
+FAILURES = []
+
+# bf16 outputs: the dominant error is rounding D to bf16 (~1e-6 in calc_diff terms).
+BF16_TOL = 1e-5
+
+
+def check(name, cond, extra=""):
+    if cond:
+        print(f"  [ok] {name}")
+    else:
+        print(f"  [FAIL] {name} {extra}")
+        FAILURES.append(name)
+
+
+def calc_diff(x, y):
+    x, y = x.double(), y.double()
+    denom = (x * x + y * y).sum().item()
+    if denom == 0:
+        return 0.0
+    return 1 - 2 * (x * y).sum().item() / denom
+
+
+def ceil_div(x, y):
+    return (x + y - 1) // y
+
+
+def aligned(x, a):
+    return ceil_div(x, a) * a
+
+
+# --------------------------------------------------------------------------- #
+# generators (ported from DeepGEMM-Ascend/tests/generators.py)
+# --------------------------------------------------------------------------- #
+
+def ceil_to_ue8m0(x):
+    bits = x.abs().float().view(torch.int32)
+    exp = ((bits // 2 ** 23) % 256) + (bits % 2 ** 23).bool().int()
+    return (exp.clamp(1, 254) << 23).view(torch.float32)
+
+
+def quantize_to_fp4_raw(x):
+    ax = x.abs()
+    code = torch.zeros_like(x, dtype=torch.uint8)
+    for boundary in (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0):
+        code += (ax > boundary).to(torch.uint8)
+    sign = (x < 0) & (code != 0)
+    values = torch.tensor((0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0), dtype=torch.float32)
+    repr_f32 = values[code.long()]
+    repr_f32 = torch.where(sign, -repr_f32, repr_f32)
+    code = code | (sign.to(torch.uint8) * 8)
+    packed = (code[..., ::2] & 0x0F) | ((code[..., 1::2] & 0x0F) << 4)
+    return packed.contiguous().view(torch.int8), repr_f32
+
+
+def quantize_last_dim(x, quant_dtype, gran_k=32):
+    """Quantize along the last dim in gran_k blocks; returns (data, sf_fp32, ref_f32)."""
+    k = x.shape[-1]
+    assert k % gran_k == 0
+    blocks = x.float().reshape(*x.shape[:-1], k // gran_k, gran_k)
+    sf_div = 6.0 if quant_dtype == "e2m1" else 448.0
+    sf = ceil_to_ue8m0(blocks.abs().amax(dim=-1).clamp(1e-4) / sf_div)
+    scaled = (blocks / sf[..., None]).reshape_as(x)
+    if quant_dtype == "e2m1":
+        data, repr_f32 = quantize_to_fp4_raw(scaled)
+    else:
+        data = scaled.to(torch.float8_e4m3fn)
+        repr_f32 = data.float()
+    ref = (repr_f32.reshape(*x.shape[:-1], k // gran_k, gran_k) * sf[..., None]).reshape_as(x)
+    return data, sf.contiguous(), ref
+
+
+def pack_sf_row_major(sf_fp32):
+    """Packed UE8M0 int16 in K-major order (what per_token_cast emits without the col-major view)."""
+    exponent = ((sf_fp32.view(torch.int32) // (1 << 23)) & 0xFF).to(torch.uint8)
+    if exponent.shape[-1] % 2:
+        exponent = torch.nn.functional.pad(exponent, (0, 1))
+    return exponent.contiguous().view(torch.int16)
+
+
+def pack_sf_reference(sf_fp32, mn, gran_mn):
+    """test_layout.py::reference - fp32 sf -> canonical int16 MN-major."""
+    exponent = ((sf_fp32.view(torch.int32) // (1 << 23)) & 0xFF).to(torch.uint8)
+    if exponent.shape[-1] % 2:
+        exponent = torch.nn.functional.pad(exponent, (0, 1))
+    sf = exponent.contiguous().view(torch.int16)
+    return (
+        sf.transpose(-2, -1)
+        .repeat_interleave(gran_mn, -1)[..., :mn]
+        .contiguous()
+        .transpose(-2, -1)
+    )
+
+
+# --------------------------------------------------------------------------- #
+# 1. SF layout
+# --------------------------------------------------------------------------- #
+
+def make_sf(shape, dtype):
+    if dtype == torch.float32:
+        exponent = torch.randint(1, 255, shape, dtype=torch.int32)
+        return (exponent * (1 << 23)).view(torch.float32)
+    return torch.randint(-32768, 32768, shape, dtype=torch.int16)
+
+
+def with_major(sf, mn_major):
+    if mn_major and sf.size(-2) > 1 and sf.size(-1) > 1:
+        sf = sf.transpose(-2, -1).contiguous().transpose(-2, -1)
+        assert sf.stride(-2) == 1
+    else:
+        assert sf.stride(-1) == 1
+    return sf
+
+
+def test_sf_layout():
+    print("test_sf_layout")
+    for dtype in (torch.float32, torch.int16):
+        for mn_major in (False, True):
+            for gran_mn in (1, 32):
+                for num_groups in (None, 3):
+                    mn, k = 128, 256
+                    sf_k_div = 32 if dtype == torch.float32 else 64
+                    shape = (ceil_div(mn, gran_mn), ceil_div(k, sf_k_div))
+                    if num_groups is not None:
+                        shape = (num_groups, *shape)
+                    sf = with_major(make_sf(shape, dtype), mn_major)
+                    out = deep_gemm.transform_sf_into_required_layout(
+                        sf, mn, k, (gran_mn, gran_mn, 32), num_groups=num_groups, is_sfa=True
+                    )
+                    expected = pack_sf_reference(sf, mn, gran_mn) if dtype == torch.float32 else (
+                        sf.transpose(-2, -1).repeat_interleave(gran_mn, -1)[..., :mn].contiguous().transpose(-2, -1)
+                    )
+                    tag = f"dtype={dtype}, mn_major={mn_major}, gran_mn={gran_mn}, groups={num_groups}"
+                    ok = (
+                        out.dtype == torch.int16
+                        and out.shape == expected.shape
+                        and out.stride(-2) == 1
+                        and out.stride(-1) == out.size(-2)
+                        and torch.equal(out, expected)
+                    )
+                    check(f"sf layout ({tag})", ok, f"{out.shape=} {expected.shape=} {out.stride()=}")
+    # Canonical input is returned unchanged (no copy).
+    mn, k = 64, 128
+    canonical = torch.zeros((ceil_div(k, 64), mn), dtype=torch.int16).transpose(0, 1)
+    out = deep_gemm.transform_sf_into_required_layout(canonical, mn, k, (1, 1, 32), is_sfa=True)
+    check("canonical sf passthrough", out.data_ptr() == canonical.data_ptr())
+    # Ascend/Neuron col-major sf is unpadded, so get_tma_aligned_size is the identity.
+    check("get_tma_aligned_size (unpadded)", deep_gemm.get_tma_aligned_size(5, 2) == 5 and deep_gemm.get_tma_aligned_size(9, 4) == 9)
+    # Unpadded col-major int16 sf (strides (1, MN)) is the canonical layout.
+    mn2, k2 = 37, 256
+    col_major = torch.zeros((ceil_div(k2, 64), mn2), dtype=torch.int16).transpose(0, 1)
+    out2 = deep_gemm.transform_sf_into_required_layout(col_major, mn2, k2, (1, 1, 32), is_sfa=False)
+    check("unpadded col-major sf passthrough", out2.data_ptr() == col_major.data_ptr() and out2.stride() == (1, mn2))
+
+
+# --------------------------------------------------------------------------- #
+# 2. dense GEMMs
+# --------------------------------------------------------------------------- #
+
+def test_bf16_gemm():
+    print("test_bf16_gemm")
+    m, n, k = 65, 128, 256
+    a = torch.randn(m, k, dtype=torch.bfloat16)
+    b = torch.randn(n, k, dtype=torch.bfloat16)
+    ref = a.float() @ b.float().T
+    d = torch.empty(m, n, dtype=torch.bfloat16)
+    deep_gemm.bf16_gemm_nt(a, b, d, None)
+    check("bf16_gemm_nt", calc_diff(d.float(), ref) < BF16_TOL, calc_diff(d.float(), ref))
+
+    d32 = torch.empty(m, n, dtype=torch.float32)
+    deep_gemm.cublaslt_gemm_nt(a, b, d32, None)
+    check("cublaslt_gemm_nt == bf16 ref", calc_diff(d32, ref) < 1e-10)
+
+    # accumulate: c aliases d
+    acc = torch.randn(m, n, dtype=torch.float32)
+    d_acc = acc.clone()
+    deep_gemm.bf16_gemm_nt(a, b, d_acc, d_acc)
+    check("bf16_gemm_nt accumulate", calc_diff(d_acc, ref + acc) < 1e-10)
+
+    # separate [1, N] bias
+    bias = torch.randn(1, n, dtype=torch.float32)
+    d_bias = torch.empty(m, n, dtype=torch.float32)
+    deep_gemm.bf16_gemm_nt(a, b, d_bias, bias)
+    check("bf16_gemm_nt bias", calc_diff(d_bias, ref + bias) < 1e-10)
+
+    # tn / nn variants
+    at = a.T.contiguous()
+    d_tn = torch.empty(m, n, dtype=torch.float32)
+    deep_gemm.bf16_gemm_tn(at, b.T.contiguous(), d_tn, None)
+    check("bf16_gemm_tn", calc_diff(d_tn, ref) < 1e-10)
+    d_nn = torch.empty(m, n, dtype=torch.float32)
+    deep_gemm.bf16_gemm_nn(a, b.T.contiguous(), d_nn, None)
+    check("bf16_gemm_nn", calc_diff(d_nn, ref) < 1e-10)
+    # strided output view
+    storage = torch.empty(m, n + 7, dtype=torch.bfloat16)
+    view = storage[:, :n]
+    deep_gemm.bf16_gemm_nt(a, b, view, None)
+    check("bf16_gemm_nt strided out", calc_diff(view.float(), ref) < BF16_TOL, calc_diff(view.float(), ref))
+
+
+def run_fp8_fp4_gemm_case(a_fmt, b_fmt, sf_mode, m=64, n=128, k=256, acc=False, gran_b=1):
+    a_bf16 = torch.randn(m, k, dtype=torch.bfloat16)
+    b_bf16 = torch.randn(n, k, dtype=torch.bfloat16)
+    a_data, a_sf, a_ref = quantize_last_dim(a_bf16, a_fmt)
+    b_data, b_sf, b_ref = quantize_last_dim(b_bf16, b_fmt)
+    if sf_mode == "packed":
+        a_sf_in = pack_sf_reference(a_sf, m, 1)
+        b_sf_in = pack_sf_reference(b_sf, n, 1)
+    elif sf_mode == "packed_row_major":
+        a_sf_in = pack_sf_row_major(a_sf)
+        b_sf_in = pack_sf_row_major(b_sf)
+    else:
+        a_sf_in, b_sf_in = a_sf, b_sf
+    ref = a_ref @ b_ref.T
+    d = torch.empty(m, n, dtype=torch.bfloat16)
+    c = None
+    if acc:
+        base = torch.randn(m, n, dtype=torch.bfloat16)
+        d.copy_(base)
+        c = d
+        ref = ref + base.float()
+    deep_gemm.fp8_fp4_gemm_nt((a_data, a_sf_in), (b_data, b_sf_in), d, c, recipe=(1, 1, 32))
+    return calc_diff(d.float(), ref)
+
+
+def test_fp8_fp4_gemm():
+    print("test_fp8_fp4_gemm")
+    for a_fmt, b_fmt in (("e4m3", "e4m3"), ("e4m3", "e2m1"), ("e2m1", "e2m1")):
+        for sf_mode in ("fp32", "packed", "packed_row_major"):
+            diff = run_fp8_fp4_gemm_case(a_fmt, b_fmt, sf_mode)
+            check(f"fp8_fp4_gemm_nt ({a_fmt}x{b_fmt}, sf={sf_mode})", diff < BF16_TOL, f"{diff=}")
+    diff = run_fp8_fp4_gemm_case("e4m3", "e2m1", "packed", acc=True)
+    check("fp8_fp4_gemm_nt accumulate", diff < BF16_TOL, f"{diff=}")
+
+
+# --------------------------------------------------------------------------- #
+# 3. m-grouped contiguous GEMMs
+# --------------------------------------------------------------------------- #
+
+def make_psum_layout(expected_m, num_groups, alignment):
+    layout, psum = [], 0
+    for _ in range(num_groups):
+        size = max(int(random.uniform(0.7, 1.3) * expected_m), 0)
+        layout.append(psum + size)
+        psum = aligned(psum + size, alignment)
+    return psum, layout
+
+
+def iter_psum_slice(layout, alignment):
+    last = 0
+    for ptr in layout:
+        yield slice(last, ptr)
+        last = aligned(ptr, alignment)
+
+
+def test_m_grouped():
+    print("test_m_grouped")
+    alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout()
+    check("mk alignment == 256", alignment == 256)
+    deep_gemm.set_mk_alignment_for_contiguous_layout(alignment)
+    num_groups, expected_m, n, k = 3, 200, 64, 128
+    total_m, layout = make_psum_layout(expected_m, num_groups, alignment)
+    grouped = torch.tensor(layout, dtype=torch.int32)
+    spans = list(iter_psum_slice(layout, alignment))
+
+    # bf16
+    a = torch.randn(total_m, k, dtype=torch.bfloat16)
+    for span in spans:  # zero the inter-group padding
+        pass
+    pad_mask = torch.ones(total_m, dtype=torch.bool)
+    for span in spans:
+        pad_mask[span] = False
+    a[pad_mask] = 0
+    b = torch.randn(num_groups, n, k, dtype=torch.bfloat16)
+    ref = torch.zeros(total_m, n, dtype=torch.float32)
+    for i, span in enumerate(spans):
+        ref[span] = a[span].float() @ b[i].float().T
+    d = torch.empty(total_m, n, dtype=torch.bfloat16)
+    deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped, use_psum_layout=True, expected_m_for_psum_layout=expected_m)
+    check("m_grouped_bf16_gemm_nt_contiguous", calc_diff(d.float(), ref) < BF16_TOL, calc_diff(d.float(), ref))
+    check("m_grouped bf16 zero padding", bool((d.float()[pad_mask] == 0).all()))
+
+    # fp8 x fp4 (the DeepInfer mxfp4 MoE path)
+    a_bf16 = torch.randn(total_m, k, dtype=torch.bfloat16)
+    a_bf16[pad_mask] = 0
+    a_data, a_sf, a_ref = quantize_last_dim(a_bf16, "e4m3")
+    b_bf16 = torch.randn(num_groups, n, k, dtype=torch.bfloat16)
+    b_data, b_sf, b_ref = quantize_last_dim(b_bf16, "e2m1")
+    a_sf_packed = pack_sf_reference(a_sf, total_m, 1)
+    b_sf_packed = pack_sf_reference(b_sf, n, 1)
+    ref = torch.zeros(total_m, n, dtype=torch.float32)
+    for i, span in enumerate(spans):
+        ref[span] = a_ref[span] @ b_ref[i].T
+    d = torch.empty(total_m, n, dtype=torch.bfloat16)
+    deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(
+        (a_data, a_sf_packed), (b_data, b_sf_packed), d, grouped,
+        recipe_a=(1, 32), recipe_b=(1, 32), use_psum_layout=True, expected_m_for_psum_layout=expected_m,
+    )
+    diff = calc_diff(d.float(), ref)
+    check("m_grouped_fp8_fp4_gemm_nt_contiguous", diff < BF16_TOL, f"{diff=}")
+    check("m_grouped fp8xfp4 zero padding", bool((d.float()[pad_mask] == 0).all()))
+
+
+# --------------------------------------------------------------------------- #
+# 4. einsum
+# --------------------------------------------------------------------------- #
+
+def test_einsum():
+    print("test_einsum")
+    b_dim, h, r, dd = 7, 4, 128, 64
+    a_bf16 = torch.randn(b_dim, h, r, dtype=torch.bfloat16)
+    w_bf16 = torch.randn(h, dd, r, dtype=torch.bfloat16)
+    ref = torch.einsum("bhr,hdr->bhd", a_bf16.float(), w_bf16.float())
+    out = torch.empty(b_dim, h, dd, dtype=torch.bfloat16)
+    deep_gemm.einsum("bhr,hdr->bhd", a_bf16, w_bf16, out, None)
+    check("einsum bf16 bhr,hdr->bhd", calc_diff(out.float(), ref) < BF16_TOL, calc_diff(out.float(), ref))
+
+    a_data, a_sf, a_ref = quantize_last_dim(a_bf16, "e4m3")
+    w_data, w_sf, w_ref = quantize_last_dim(w_bf16, "e4m3")
+    a_sf_packed = pack_sf_reference(a_sf, a_sf.shape[-2], 1)  # [b, h, r/64] canonical per b
+    w_sf_packed = pack_sf_reference(w_sf, w_sf.shape[-2], 1)
+    ref = torch.einsum("bhr,hdr->bhd", a_ref, w_ref)
+    out = torch.empty(b_dim, h, dd, dtype=torch.bfloat16)
+    deep_gemm.fp8_einsum("bhr,hdr->bhd", (a_data, a_sf_packed), (w_data, w_sf_packed), out, recipe=(1, 1, 32))
+    diff = calc_diff(out.float(), ref)
+    check("fp8_einsum bhr,hdr->bhd", diff < BF16_TOL, f"{diff=}")
+    # the fp32-sf input form must agree too
+    out2 = torch.empty_like(out)
+    deep_gemm.fp8_einsum("bhr,hdr->bhd", (a_data, a_sf), (w_data, w_sf), out2, recipe=(1, 1, 32))
+    check("fp8_einsum fp32 sf == packed sf", torch.equal(out, out2))
+
+
+# --------------------------------------------------------------------------- #
+# 5. MQA logits (ported from tests/test_attention.py)
+# --------------------------------------------------------------------------- #
+
+def ceil_to_ue8m0_with_byte(x):
+    bits = x.abs().float().view(torch.int32)
+    exp = (((bits >> 23) & 0xFF) + (bits & 0x7FFFFF).bool().int()).clamp(1, 254)
+    return (exp << 23).view(torch.float32), exp.to(torch.uint8)
+
+
+def quantize_mqa_tensor(x, use_fp4):
+    sf_div = 6.0 if use_fp4 else 448.0
+    sf_shape = (*x.shape[:-1], x.shape[-1] // 32, 32)
+    sf, sf_byte = ceil_to_ue8m0_with_byte(x.float().reshape(sf_shape).abs().amax(dim=-1).clamp(1e-4) / sf_div)
+    x_scaled = (x.float().reshape(sf_shape) / sf[..., None]).reshape_as(x)
+    if use_fp4:
+        x_quant, x_repr = quantize_to_fp4_raw(x_scaled)
+    else:
+        x_quant = x_scaled.to(torch.float8_e4m3fn)
+        x_repr = x_quant.float()
+    x_ref = (x_repr.reshape(sf_shape) * sf[..., None]).reshape_as(x)
+    sf_packed = sf_byte[..., ::2].to(torch.int16) | (sf_byte[..., 1::2].to(torch.int16) << 8)
+    return x_quant, sf_packed.contiguous(), x_ref
+
+
+def ref_mqa_logits(q, kv, weights, ks, ke):
+    seq_len, num_heads, _ = q.shape
+    seq_len_kv = kv.shape[0]
+    score = torch.einsum("qhd,kd->qhk", q.float(), kv.float())
+    out = (score.relu() * weights.float()[:, :, None]).sum(dim=1)
+    pos = torch.arange(seq_len_kv)
+    mask = (pos[None, :] >= ks[:, None]) & (pos[None, :] < ke[:, None])
+    return out.masked_fill(~mask, float("-inf"))
+
+
+def test_mqa_logits(use_fp4):
+    name = "fp4" if use_fp4 else "fp8"
+    seq_len, seq_len_kv, num_heads, head_dim = 9, 512, 8, 128
+    q_bf16 = torch.randn(seq_len, num_heads, head_dim, dtype=torch.bfloat16)
+    kv_bf16 = torch.randn(seq_len_kv, head_dim, dtype=torch.bfloat16)
+    weights = torch.randn(seq_len, num_heads, dtype=torch.bfloat16)
+    a = torch.randint(0, seq_len_kv + 1, (seq_len,), dtype=torch.int32)
+    b = torch.randint(0, seq_len_kv + 1, (seq_len,), dtype=torch.int32)
+    ks, ke = torch.minimum(a, b), torch.maximum(a, b)
+    max_seqlen_k = max(1, int((ke - ks).max()))
+
+    q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
+    kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)
+    ref = ref_mqa_logits(q_ref, kv_ref, weights, ks, ke).to(torch.bfloat16)
+
+    out = deep_gemm.fp8_fp4_mqa_logits((q, q_sf), (kv, kv_sf), weights, ks, ke, max_seqlen_k)
+    check(f"mqa logits {name} dtype/stride", out.dtype == torch.bfloat16
+          and out.shape == (seq_len, max_seqlen_k)
+          and out.stride(0) * out.element_size() % 512 == 0 and out.stride(1) == 1)
+    dense = torch.full((seq_len, seq_len_kv), float("-inf"), dtype=torch.bfloat16)
+    for i, (s, e) in enumerate(zip(ks.tolist(), ke.tolist())):
+        dense[i, s:e] = out[i, : e - s]
+    neg = torch.isinf(ref.float()) & (ref.float() < 0)
+    check(f"mqa logits {name} -inf mask", torch.equal(torch.isinf(dense.float()) & (dense.float() < 0), neg))
+    diff = calc_diff(dense.float().masked_fill(neg, 0), ref.float().masked_fill(neg, 0))
+    check(f"mqa logits {name} values", diff < 5e-4, f"{diff=}")
+    return q, q_sf, q_ref, kv, kv_sf, kv_ref, weights, ks, ke
+
+
+def make_fused_pages(kv, kv_sf, num_pages, page_kv, use_fp4):
+    data_bytes = 64 if use_fp4 else 128
+    sf_bytes = 4
+    token_bytes = data_bytes + sf_bytes
+    data = kv.contiguous().view(torch.uint8).reshape(num_pages, page_kv * data_bytes)
+    sf = kv_sf.contiguous().view(torch.uint8).reshape(num_pages, page_kv * sf_bytes)
+    fused = torch.cat([data, sf], dim=1)
+    return fused.reshape(num_pages, page_kv, 1, token_bytes)
+
+
+def ref_paged_mqa_logits(q, kv, weights, context_lens, block_table, max_context_len):
+    num_q = q.shape[0]
+    page_size = kv.shape[1]
+    out = torch.zeros(num_q, max_context_len)
+    for q_idx, (context_len,) in enumerate(context_lens.tolist()):
+        pages = ceil_div(context_len, page_size)
+        keys = kv[block_table[q_idx, :pages].long()].flatten(0, 1).float()
+        scores = torch.einsum("hd,kd->hk", q[q_idx, 0].float(), keys)
+        logits = (scores.relu() * weights[q_idx, :, None].float()).sum(0)
+        out[q_idx, :context_len] = logits[:context_len]
+    return out
+
+
+def test_paged_mqa_logits(use_fp4):
+    name = "fp4" if use_fp4 else "fp8"
+    batch, num_heads, head_dim, page_kv = 3, 8, 128, 64
+    context_lens_host = [70, 128, 33]
+    pages_per_request = [ceil_div(c, page_kv) for c in context_lens_host]
+    num_pages = sum(pages_per_request)
+    q_bf16 = torch.randn(batch, 1, num_heads, head_dim, dtype=torch.bfloat16)
+    kv_bf16 = torch.randn(num_pages, page_kv, head_dim, dtype=torch.bfloat16)
+    weights = torch.randn(batch, num_heads, dtype=torch.bfloat16)
+    q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
+    kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)
+    fused = make_fused_pages(kv, kv_sf, num_pages, page_kv, use_fp4)
+
+    perm = torch.randperm(num_pages, dtype=torch.int32)
+    max_pages = max(pages_per_request)
+    block_table = torch.zeros(batch, max_pages, dtype=torch.int32)
+    offset = 0
+    for i, count in enumerate(pages_per_request):
+        block_table[i, :count] = perm[offset:offset + count]
+        offset += count
+    context_lens = torch.tensor(context_lens_host, dtype=torch.int32).view(-1, 1)
+    indices = torch.arange(batch, dtype=torch.int32)
+    max_context_len = max(context_lens_host)
+
+    # Reference uses the (permuted) page order through the block table.
+    ref = ref_paged_mqa_logits(q_ref, kv_ref, weights, context_lens, block_table, max_context_len)
+    meta = deep_gemm.get_paged_mqa_logits_metadata(context_lens, num_heads, indices)
+    out = deep_gemm.fp8_fp4_paged_mqa_logits(
+        (q, q_sf), fused, weights, context_lens, block_table, meta, max_context_len, indices, False
+    )
+    valid = torch.arange(max_context_len)[None] < context_lens.view(-1, 1)
+    diff = calc_diff(out.float().masked_fill(~valid, 0), ref.masked_fill(~valid, 0))
+    check(f"paged mqa logits {name}", diff < 5e-4, f"{diff=}")
+    check(f"paged mqa logits {name} stride", out.stride(0) * out.element_size() % 512 == 0 and out.stride(1) == 1)
+    return (q, q_sf), fused, weights, context_lens, block_table, indices, max_context_len, out
+
+
+def test_sparse_mqa_logits(use_fp4):
+    """Sparse logits must match the dense logits at the selected tokens."""
+    name = "fp4" if use_fp4 else "fp8"
+    num_q, num_kv, num_heads, head_dim = 6, 512, 8, 128
+    sparse_block_kv, capacity = 8, 16
+    q_bf16 = torch.randn(num_q, num_heads, head_dim, dtype=torch.bfloat16)
+    kv_bf16 = torch.randn(num_kv, head_dim, dtype=torch.bfloat16)
+    weights = torch.randn(num_q, num_heads, dtype=torch.bfloat16)
+    q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
+    kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)
+
+    rng = random.Random(1)
+    starts, ends, blocks = [], [], []
+    for _ in range(num_q):
+        start = rng.randrange(0, 4) * sparse_block_kv + rng.randrange(0, sparse_block_kv)
+        end = rng.randrange(start + 1, num_kv)
+        first_block = start // sparse_block_kv
+        available = ceil_div(max(0, end - start), sparse_block_kv)
+        count = min(available, capacity)
+        chosen = sorted(rng.sample(range(first_block, first_block + available), count))
+        chosen += [chosen[-1] if chosen else 0] * (capacity - count)
+        starts.append(start)
+        ends.append(end)
+        blocks.append(chosen)
+    ks = torch.tensor(starts, dtype=torch.int32)
+    ke = torch.tensor(ends, dtype=torch.int32)
+    sparse_blocks = torch.tensor(blocks, dtype=torch.int32)
+
+    meta = deep_gemm.get_sparse_mqa_logits_metadata(ks, ke, num_kv, sparse_blocks, q.dtype, sparse_block_kv)
+    sparse = deep_gemm.fp8_fp4_sparse_mqa_logits((q, q_sf), (kv, kv_sf), weights, meta, capacity, sparse_block_kv)
+    check(f"sparse mqa {name} shape/stride",
+          sparse.shape == (num_q, capacity * sparse_block_kv) and sparse.dtype == torch.bfloat16
+          and sparse.stride(0) * sparse.element_size() % 512 == 0)
+
+    full = deep_gemm.fp8_fp4_mqa_logits(
+        (q, q_sf), (kv, kv_sf), weights, ks, ke, max(1, int((ke - ks).max()))
+    )
+    offsets = torch.arange(sparse_block_kv, dtype=torch.int32)
+    token_indices = (sparse_blocks.unsqueeze(-1) * sparse_block_kv + (ks % sparse_block_kv)[:, None, None] + offsets).flatten(1).long()
+    num_blocks = torch.tensor([min(capacity, ceil_div(max(0, e - s), sparse_block_kv)) for s, e in zip(starts, ends)])
+    valid = torch.arange(token_indices.shape[1])[None] < (num_blocks * sparse_block_kv)[:, None]
+    valid &= token_indices >= ks[:, None]
+    valid &= token_indices < ke[:, None]
+    full_cols = (token_indices - ks[:, None]).clamp(0, full.shape[1] - 1)
+    check(f"sparse mqa {name} == dense logits",
+          torch.equal(sparse[valid], full.gather(1, full_cols)[valid]))
+
+    # Determinism / reuse
+    sparse2 = deep_gemm.fp8_fp4_sparse_mqa_logits((q, q_sf), (kv, kv_sf), weights, meta, capacity, sparse_block_kv)
+    check(f"sparse mqa {name} repeatable", torch.equal(sparse[valid], sparse2[valid]))
+
+
+def test_paged_sparse_mqa_logits(use_fp4):
+    name = "fp4" if use_fp4 else "fp8"
+    num_heads, head_dim, page_kv = 8, 128, 64
+    sparse_block_kv, capacity = 8, 16
+    context_lens_host = [70, 128, 33]
+    batch = len(context_lens_host)
+    pages_per_request = [ceil_div(c, page_kv) for c in context_lens_host]
+    num_pages = sum(pages_per_request)
+    q_bf16 = torch.randn(batch, 1, num_heads, head_dim, dtype=torch.bfloat16)
+    kv_bf16 = torch.randn(num_pages, page_kv, head_dim, dtype=torch.bfloat16)
+    weights = torch.randn(batch, num_heads, dtype=torch.bfloat16)
+    q, q_sf, q_ref = quantize_mqa_tensor(q_bf16, use_fp4)
+    kv, kv_sf, kv_ref = quantize_mqa_tensor(kv_bf16, use_fp4)
+    fused = make_fused_pages(kv, kv_sf, num_pages, page_kv, use_fp4)
+
+    perm = torch.randperm(num_pages, dtype=torch.int32)
+    max_pages = max(pages_per_request)
+    block_table = torch.zeros(batch, max_pages, dtype=torch.int32)
+    offset = 0
+    for i, count in enumerate(pages_per_request):
+        block_table[i, :count] = perm[offset:offset + count]
+        offset += count
+    context_lens = torch.tensor(context_lens_host, dtype=torch.int32)
+    indices = torch.arange(batch, dtype=torch.int32)
+
+    rng = random.Random(2)
+    blocks = []
+    for context_len in context_lens_host:
+        available = ceil_div(context_len, sparse_block_kv)
+        count = min(available, capacity)
+        chosen = sorted(rng.sample(range(available), count))
+        blocks.append(chosen + [chosen[-1]] * (capacity - count))
+    sparse_blocks = torch.tensor(blocks, dtype=torch.int32)
+
+    meta = deep_gemm.get_paged_sparse_mqa_logits_metadata(
+        context_lens, block_table, indices, page_kv, sparse_blocks, q.dtype, sparse_block_kv
+    )
+    sparse = deep_gemm.fp8_fp4_paged_sparse_mqa_logits((q, q_sf), fused, weights, meta, capacity, sparse_block_kv)
+    check(f"paged sparse {name} shape/stride",
+          sparse.shape == (batch, capacity * sparse_block_kv) and sparse.dtype == torch.bfloat16
+          and sparse.stride(0) * sparse.element_size() % 512 == 0)
+
+    dense_meta = deep_gemm.get_paged_mqa_logits_metadata(context_lens.view(-1, 1), num_heads, indices)
+    full = deep_gemm.fp8_fp4_paged_mqa_logits(
+        (q, q_sf), fused, weights, context_lens.view(-1, 1), block_table, dense_meta,
+        max(context_lens_host), indices, False
+    )
+    offsets = torch.arange(sparse_block_kv, dtype=torch.int32)
+    token_indices = (sparse_blocks.unsqueeze(-1) * sparse_block_kv + offsets).flatten(1).long()
+    num_blocks = torch.tensor([min(capacity, ceil_div(c, sparse_block_kv)) for c in context_lens_host])
+    valid = torch.arange(token_indices.shape[1])[None] < (num_blocks * sparse_block_kv)[:, None]
+    valid &= token_indices < context_lens[:, None]
+    full_cols = token_indices.clamp(0, full.shape[1] - 1)
+    check(f"paged sparse {name} == dense logits",
+          torch.equal(sparse[valid], full.gather(1, full_cols)[valid]))
+
+    # REUSE: the KV snapshot taken on the first call survives a KV overwrite.
+    first = sparse.clone()
+    fused.zero_()
+    again = deep_gemm.fp8_fp4_paged_sparse_mqa_logits((q, q_sf), fused, weights, meta, capacity, sparse_block_kv)
+    check(f"paged sparse {name} reuses the KV snapshot", torch.equal(first[valid], again[valid]))
+
+
+# --------------------------------------------------------------------------- #
+# 6. hyperconnection prenorm
+# --------------------------------------------------------------------------- #
+
+def test_hc_prenorm():
+    print("test_hc_prenorm")
+    m, n, k = 65, 24, 256
+    a = torch.randn(m, k, dtype=torch.bfloat16)
+    b = torch.randn(n, k, dtype=torch.float32)
+    d = torch.empty(m, n, dtype=torch.float32)
+    sqr_sum = torch.empty(m, dtype=torch.float32)
+    deep_gemm.tf32_hc_prenorm_gemm(a, b, d, sqr_sum)
+    ref_d = a.float() @ b.T
+    ref_sqr = a.float().square().sum(-1)
+    check("tf32_hc_prenorm_gemm d", calc_diff(d, ref_d) < 1e-6, calc_diff(d, ref_d))
+    check("tf32_hc_prenorm_gemm sqr_sum", calc_diff(sqr_sum, ref_sqr) < 1e-6)
+    d_split = torch.empty(1, m, n, dtype=torch.float32)
+    sqr_split = torch.empty(1, m, dtype=torch.float32)
+    deep_gemm.tf32_hc_prenorm_gemm(a, b, d_split, sqr_split, num_splits=1)
+    check("tf32_hc_prenorm_gemm num_splits=1", torch.equal(d_split[0], d) and torch.equal(sqr_split[0], sqr_sum))
+
+
+# --------------------------------------------------------------------------- #
+# 7. API surface
+# --------------------------------------------------------------------------- #
+
+def test_api_surface():
+    print("test_api_surface")
+    required = [
+        "SymmBuffer", "bf16_gemm_nt", "bf16_mega_gate", "bf16_mega_moe", "cublaslt_gemm_nt",
+        "fp8_einsum", "fp8_fp4_gemm_nt", "fp8_fp4_mega_moe", "fp8_fp4_mqa_logits",
+        "fp8_fp4_paged_mqa_logits", "fp8_fp4_paged_sparse_mqa_logits", "fp8_fp4_sparse_mqa_logits",
+        "get_block_m_for_mega_moe", "get_mk_alignment_for_contiguous_layout", "get_num_sms",
+        "get_paged_mqa_logits_metadata", "get_paged_sparse_mqa_logits_metadata",
+        "get_sparse_mqa_logits_metadata", "get_theoretical_mk_alignment_for_contiguous_layout",
+        "get_tma_aligned_size", "m_grouped_bf16_gemm_nt_contiguous",
+        "m_grouped_fp8_fp4_gemm_nt_contiguous", "mega_mhc", "set_mk_alignment_for_contiguous_layout",
+        "set_num_sms", "set_pdl", "transform_sf_into_required_layout",
+        "transform_weights_for_mega_moe", "use_deterministic_algorithms", "tf32_hc_prenorm_gemm",
+    ]
+    missing = [name for name in required if not hasattr(deep_gemm, name)]
+    check("all required symbols exported", not missing, f"missing={missing}")
+    for name in ("bf16_mega_gate", "bf16_mega_moe", "fp8_fp4_mega_moe", "mega_mhc",
+                 "transform_weights_for_mega_moe", "get_block_m_for_mega_moe"):
+        try:
+            getattr(deep_gemm, name)()
+            check(f"{name} raises NotImplementedError", False)
+        except NotImplementedError:
+            check(f"{name} raises NotImplementedError", True)
+        except Exception as exc:  # noqa: BLE001
+            check(f"{name} raises NotImplementedError", False, repr(exc))
+    prev = deep_gemm.get_num_sms()
+    check("default num_sms == 8 (tile_kernels.get_device_num_sms)", prev == 8, f"{prev=}")
+    deep_gemm.set_num_sms(prev + 4)
+    check("set/get_num_sms", deep_gemm.get_num_sms() == prev + 4)
+    deep_gemm.set_num_sms(prev)
+    deep_gemm.use_deterministic_algorithms(True)
+    check("use_deterministic_algorithms", deep_gemm.get_deterministic_algorithms() is True)
+    deep_gemm.use_deterministic_algorithms(False)
+    deep_gemm.set_pdl(True)
+
+    class _Group:
+        def size(self):
+            return 2
+
+        def rank(self):
+            return 0
+
+    buf = deep_gemm.SymmBuffer(
+        _Group(), num_experts=8, num_max_tokens_per_rank=16, num_topk=4, hidden=256,
+        intermediate_hidden=128, num_shared_experts=1, mma_type="fp8xfp4", activation="swiglu",
+        device="cpu",
+    )
+    check("SymmBuffer views", buf.x.shape == (16, 256) and buf.x_sf.shape == (16, 4)
+          and buf.topk_idx.shape == (16, 4) and buf.num_bytes > 0)
+
+
+if __name__ == "__main__":
+    test_sf_layout()
+    test_bf16_gemm()
+    test_fp8_fp4_gemm()
+    test_m_grouped()
+    test_einsum()
+    print("test_mqa_logits")
+    test_mqa_logits(False)
+    test_mqa_logits(True)
+    print("test_paged_mqa_logits")
+    test_paged_mqa_logits(False)
+    test_paged_mqa_logits(True)
+    print("test_sparse_mqa_logits")
+    test_sparse_mqa_logits(False)
+    test_sparse_mqa_logits(True)
+    print("test_paged_sparse_mqa_logits")
+    test_paged_sparse_mqa_logits(False)
+    test_paged_sparse_mqa_logits(True)
+    test_hc_prenorm()
+    test_api_surface()
+    print()
+    if FAILURES:
+        print(f"FAILED ({len(FAILURES)}): {FAILURES}")
+        sys.exit(1)
+    print("ALL TESTS PASSED")
diff --git a/tests/neuron/deep_gemm/test_interop_tile_kernels.py b/tests/neuron/deep_gemm/test_interop_tile_kernels.py
new file mode 100644
index 0000000..feaf136
--- /dev/null
+++ b/tests/neuron/deep_gemm/test_interop_tile_kernels.py
@@ -0,0 +1,52 @@
+"""Cross-check: tile_kernels.quant.per_token_cast output feeds deep_gemm's FP8 GEMM.
+
+This is the exact ReplicatedLinear.apply_weight_fp8 path
+(deepinfer/model_executor/layers/linear.py) with CPU tensors.
+"""
+
+import os
+import sys
+
+sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "deepinfer", "neuron_ops", "vendor"))
+
+import torch
+
+import deep_gemm
+import tile_kernels
+
+CAST_PARAMS = dict(fmt="e4m3", num_per_channels=32, round_sf=True,
+                   use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
+
+
+def calc_diff(x, y):
+    x, y = x.double(), y.double()
+    denom = (x * x + y * y).sum().item()
+    return 0.0 if denom == 0 else 1 - 2 * (x * y).sum().item() / denom
+
+
+m, n, k = 48, 128, 256
+x = torch.randn(m, k, dtype=torch.bfloat16)
+w = torch.randn(n, k, dtype=torch.bfloat16)
+
+xq, x_sf = tile_kernels.quant.per_token_cast(x, **CAST_PARAMS)
+wq, w_sf = tile_kernels.quant.per_token_cast(w, **CAST_PARAMS)
+print(f"{xq.dtype=} {x_sf.dtype=} {x_sf.shape=} {x_sf.stride()=}")
+print(f"{wq.dtype=} {w_sf.dtype=} {w_sf.shape=} {w_sf.stride()=}")
+
+out = torch.empty(m, n, dtype=torch.bfloat16)
+deep_gemm.fp8_fp4_gemm_nt((xq, x_sf), (wq, w_sf), out, None, recipe=(1, 1, 32))
+
+ref = x.float() @ w.float().T
+diff = calc_diff(out.float(), ref)
+print(f"diff vs bf16 reference: {diff:.3e}")
+assert diff < 2e-3, diff
+
+# The dequantized operands must agree with the GEMM result exactly.
+from deep_gemm._quant import dequantize  # noqa: E402
+
+x_deq = dequantize(xq, deep_gemm.transform_sf_into_required_layout(x_sf, m, k, (1, 1, 32), is_sfa=True))
+w_deq = dequantize(wq, deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 1, 32), is_sfa=False))
+exact = (x_deq @ w_deq.T).to(torch.bfloat16)
+print("exact match with dequantized matmul:", torch.equal(out, exact))
+print(f"quantization error (dequant vs bf16): {calc_diff(exact.float(), ref):.3e}")
+print("OK")
diff --git a/tests/neuron/device/test_device_primitives.py b/tests/neuron/device/test_device_primitives.py
new file mode 100644
index 0000000..f24190c
--- /dev/null
+++ b/tests/neuron/device/test_device_primitives.py
@@ -0,0 +1,128 @@
+"""Native torch-neuronx behaviours the Neuron port depends on (or works around).
+
+Run on a NeuronCore (no DeepInfer import, plain torch-neuronx):
+
+    NEURON_RT_VISIBLE_CORES=0 python tests/neuron/device/test_device_primitives.py
+
+Each check compares the device result with the host. ``expect="ok"`` checks must pass; ``expect="known"``
+checks document an observed torch-neuronx 2.14 limitation the port works around -- if one of those starts
+passing, the workaround can go. Exit status is non-zero only when an ``ok`` check fails.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+DEV = torch.device("neuron", 0)
+RESULTS: list[tuple[str, str, str]] = []
+
+
+def check(name: str, expect: str, fn) -> None:
+    try:
+        ok = bool(fn())
+        status = "ok" if ok else "DIFF"
+    except Exception as e:  # noqa: BLE001
+        status = f"EXC {type(e).__name__}: {str(e).splitlines()[0][:100]}"
+    RESULTS.append((name, expect, status))
+
+
+def same(dev: torch.Tensor, host: torch.Tensor) -> bool:
+    d = dev.cpu()
+    if d.dtype in (torch.float8_e4m3fn, torch.float8_e8m0fnu):
+        d, host = d.view(torch.uint8), host.view(torch.uint8)
+    return torch.equal(d, host)
+
+
+def main() -> int:
+    torch.manual_seed(0)
+    x = torch.randn(11, 512, dtype=torch.bfloat16)
+    xd = x.to(DEV)
+
+    # dtypes used by the kernels
+    f8 = torch.randn(64, 64).clamp(-400, 400).to(torch.float8_e4m3fn)
+    check("fp8 e4m3 H2D/D2H bit-exact", "ok", lambda: same(f8.to(DEV), f8))
+    check("fp8 -> bf16 upcast on device", "ok", lambda: same(f8.to(DEV).to(torch.bfloat16), f8.to(torch.bfloat16)))
+    e8 = torch.randint(100, 140, (8, 8), dtype=torch.uint8).view(torch.float8_e8m0fnu)
+    check("e8m0 H2D bit-exact", "ok", lambda: same(e8.to(DEV), e8))
+    b = torch.randint(-128, 127, (4, 16), dtype=torch.int8)
+    check("int8 nibble unpack", "ok", lambda: same(b.to(DEV) & 0xF, b & 0xF) and same((b.to(DEV) >> 4) & 0xF, (b >> 4) & 0xF))
+    s = torch.randint(118, 130, (4, 16), dtype=torch.uint8)
+    check("ue8m0 bits -> fp32 (<<23, view)", "ok",
+          lambda: same((s.to(DEV).to(torch.int32) << 23).view(torch.float32), (s.to(torch.int32) << 23).view(torch.float32)))
+    u = torch.tensor([1, 2**40 + 7, 2**63 + 5], dtype=torch.uint64)
+    check("uint64 round trip", "known", lambda: same(u.to(DEV), u))
+    check("uint64 add", "known", lambda: same(u.to(DEV) + 3, u + 3))
+
+    # views / in-place semantics
+    check("column-slice view -> host", "ok", lambda: same(xd[:, 256:], x[:, 256:]))
+    check("transpose view -> host", "ok", lambda: same(xd.t(), x.t()))
+
+    def inplace_slice():
+        o, oc = torch.zeros(11, 1024, dtype=torch.bfloat16, device=DEV), torch.zeros(11, 1024, dtype=torch.bfloat16)
+        o[:, 512:].copy_(xd)
+        oc[:, 512:].copy_(x)
+        o[:, :512] += xd
+        oc[:, :512] += x
+        return same(o, oc)
+
+    check("copy_ / += into column slice", "ok", inplace_slice)
+    buf = torch.randint(0, 256, (2 * 902880,), dtype=torch.uint8)
+    bufd = buf.to(DEV)
+    blk, row = torch.tensor([0, 0, 1, 1, 0]), torch.tensor([0, 5, 7, 128, 3])
+
+    def paged(t: torch.Tensor) -> torch.Tensor:  # KV-cache style strided record view
+        return t.as_strided((2, 129, 512), (902880, 528, 1))
+
+    check("uint8 strided view advanced indexing", "ok",
+          lambda: same(paged(bufd)[blk.to(DEV), row.to(DEV)], paged(buf)[blk, row]))
+    check("fp8 strided view advanced indexing", "known",
+          lambda: same(paged(bufd).view(torch.float8_e4m3fn)[blk.to(DEV), row.to(DEV)], paged(buf).view(torch.float8_e4m3fn)[blk, row]))
+    idx = torch.tensor([1, 1, 4])
+    check("index_add_ (duplicates)", "ok", lambda: same(torch.zeros(11, 512, dtype=torch.bfloat16, device=DEV).index_add_(0, idx.to(DEV), xd[:3]),
+                                                         torch.zeros(11, 512, dtype=torch.bfloat16).index_add_(0, idx, x[:3])))
+
+    def inference_version_counter():
+        with torch.inference_mode():
+            t = torch.topk(xd.float(), 4).indices
+            return t.is_inference() or t._version >= 0
+
+    check("inference-mode output has is_inference() or a version counter", "known", inference_version_counter)
+
+    # numerics
+    a, w = torch.randn(64, 256), torch.randn(256, 128)
+    ref = a.double() @ w.double()
+    rel = lambda r: ((r.cpu().double() - ref).norm() / ref.norm()).item()  # noqa: E731
+    check("fp32 matmul rel err < 1e-5 (vs fp64)", "ok", lambda: rel(a.to(DEV) @ w.to(DEV)) < 1e-5)
+    check("bf16 topk indices == host", "ok", lambda: torch.equal(
+        torch.topk(xd[0], 64).indices.cpu().sort().values, torch.topk(x[0], 64).indices.sort().values))
+
+    # single-rank process group (EP=1 bring-up)
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29761")
+    dist.init_process_group("gloo", rank=0, world_size=1)
+    g = dist.new_group([0], backend="neuron")
+    check("neuron PG size 1: all_reduce", "ok", lambda: (lambda t: (dist.all_reduce(t, group=g), same(t, x))[1])(xd.clone()))
+    check("neuron PG size 1: all_gather_into_tensor", "ok",
+          lambda: (lambda o: (dist.all_gather_into_tensor(o, xd, group=g), same(o, x))[1])(torch.empty_like(xd)))
+    check("neuron PG size 1: all_to_all_single", "known",
+          lambda: (lambda o: (dist.all_to_all_single(o, xd, group=g), same(o, x))[1])(torch.empty_like(xd)))
+    dist.destroy_process_group()
+
+    failed = 0
+    for name, expect, status in RESULTS:
+        if expect == "ok":
+            tag = "PASS" if status == "ok" else "FAIL"
+            failed += tag == "FAIL"
+        else:
+            tag = "KNOWN-ISSUE" if status != "ok" else "FIXED-UPSTREAM?"
+        print(f"{tag:16s} {name:62s} {status}")
+    print(f"\n{failed} unexpected failure(s)")
+    return 1 if failed else 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tests/neuron/flash_mla/ref_attn.py b/tests/neuron/flash_mla/ref_attn.py
new file mode 100644
index 0000000..29d8df9
--- /dev/null
+++ b/tests/neuron/flash_mla/ref_attn.py
@@ -0,0 +1,89 @@
+"""Reference sparse attention ported from vendor/FlashMLA/tests/ref.py (no TestParam wrappers)."""
+
+import torch
+
+
+def ref_sparse_attn_fwd(q, kv, indices, sm_scale, topk_length=None, attn_sink=None, d_v=512, rms=None):
+    """q [s_q,h_q,d], kv [s_kv,1,d], indices [s_q,1,topk] -> (out bf16, max_logits, lse)."""
+    s_q, h_q, d_qk = q.shape
+    s_kv = kv.shape[0]
+    topk = indices.shape[-1]
+    idx = indices.clone().squeeze(1)
+    if topk_length is not None:
+        mask = torch.arange(topk).unsqueeze(0).broadcast_to(s_q, topk) >= topk_length.unsqueeze(1)
+        idx.masked_fill_(mask, -1)
+    invalid = (idx < 0) | (idx >= s_kv)
+    idx = idx.masked_fill(invalid, 0)
+
+    qf = q.float()
+    gathered = kv.view(s_kv, d_qk).index_select(0, idx.flatten().long()).reshape(s_q, topk, d_qk).float()
+    P = qf @ gathered.transpose(1, 2)
+    if rms is not None:
+        P = P * (sm_scale * rms.unsqueeze(-1))
+    else:
+        P = P * sm_scale
+    P.masked_fill_(invalid.unsqueeze(1).broadcast_to(P.shape), float("-inf"))
+
+    orig_lse = torch.logsumexp(P, dim=-1)
+    max_logits = P.max(dim=-1).values
+
+    lse_for_o = orig_lse
+    if attn_sink is not None:
+        lse_for_o = torch.logsumexp(
+            torch.stack([orig_lse.view(s_q, h_q), attn_sink.broadcast_to(s_q, h_q)], dim=0), dim=0
+        )
+    lse_for_o = lse_for_o.clone()
+    lse_for_o.masked_fill_(lse_for_o == float("-inf"), float("+inf"))
+    s_for_o = torch.exp(P - lse_for_o.unsqueeze(-1))
+    out = s_for_o @ gathered[..., :d_v]
+
+    orig_lse = orig_lse.clone()
+    orig_lse.masked_fill_(orig_lse == float("-inf"), float("+inf"))
+    return out.to(torch.bfloat16), max_logits, orig_lse
+
+
+def ref_sparse_attn_decode(q, scopes, sm_scale, attn_sink=None, d_v=512, rms=None):
+    """q [b,s_q,h_q,d]; scopes = list of (blocked_k_bf16 [n_tokens, d], indices [b,s_q,topk], topk_length|None).
+
+    Returns (out bf16 [b,s_q,h_q,d_v], lse [b,h_q,s_q]).
+    """
+    b, s_q, h_q, d_qk = q.shape
+    gathered_all, invalid_all = [], []
+    for blocked_k, indices, topk_length in scopes:
+        topk = indices.shape[-1]
+        fixed = torch.clamp_min(indices, 0)
+        g = blocked_k.view(-1, d_qk).index_select(0, fixed.reshape(-1).long()).view(b, s_q, topk, d_qk)
+        invalid = indices == -1
+        if topk_length is not None:
+            invalid = invalid | (
+                torch.arange(0, topk).view(1, 1, topk).broadcast_to(b, s_q, topk) >= topk_length.view(b, 1, 1)
+            )
+        gathered_all.append(g)
+        invalid_all.append(invalid)
+    gathered = torch.cat(gathered_all, dim=2)
+    invalid = torch.cat(invalid_all, dim=2)
+
+    gathered = gathered.reshape(b * s_q, -1, d_qk).float()
+    gathered.masked_fill_(gathered != gathered, 0.0)
+    qf = q.float().view(b * s_q, h_q, d_qk)
+    w = qf @ gathered.transpose(-1, -2)
+    if rms is not None:
+        w = w * (sm_scale * rms.view(b * s_q, h_q).unsqueeze(-1))
+    else:
+        w = w * sm_scale
+    w[invalid.view(b * s_q, 1, -1).broadcast_to(b * s_q, h_q, invalid.size(-1))] = float("-inf")
+    lse = w.logsumexp(dim=-1)
+    w = torch.exp(w - lse.unsqueeze(-1))
+    out = w @ gathered[..., :d_v]
+    out = out.view(b, s_q, h_q, d_v)
+    lse = lse.view(b, s_q, h_q)
+
+    if attn_sink is not None:
+        out = out * (1.0 / (1.0 + torch.exp(attn_sink.view(1, 1, h_q) - lse))).unsqueeze(-1)
+
+    lonely = lse == float("-inf")
+    out = out.clone()
+    out[lonely.unsqueeze(-1).broadcast_to(b, s_q, h_q, d_v)] = 0.0
+    lse = lse.clone()
+    lse[lonely] = float("+inf")
+    return out.to(torch.bfloat16), lse.transpose(1, 2)
diff --git a/tests/neuron/flash_mla/ref_quant.py b/tests/neuron/flash_mla/ref_quant.py
new file mode 100644
index 0000000..aa1e2c4
--- /dev/null
+++ b/tests/neuron/flash_mla/ref_quant.py
@@ -0,0 +1,89 @@
+"""Reference quant/dequant ported from vendor/FlashMLA/tests/quant.py (CPU, no kernelkit)."""
+
+import torch
+
+META = {
+    "v41_fp8": (512, 448, 64, 32, 16),
+    "v41_fp4": (512, 448, 64, 16, 32),
+}
+BPT = {"v41_fp8": 528, "v41_fp4": 288}
+
+_E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32)
+
+
+def _cast_scale_inv_to_ue8m0(scales_inv, out_dtype=torch.float32):
+    return torch.pow(2, torch.clamp_min(scales_inv, 1e-4).log2().ceil()).to(out_dtype)
+
+
+def _quantize_to_e2m1(x):
+    x = x.float()
+    mags = _E2M1
+    sign = (torch.signbit(x)).to(torch.uint8) << 3
+    a = torch.nan_to_num(x.abs(), nan=0.0, posinf=6.0).clamp_max(6.0)
+    mids = (mags[:-1] + mags[1:]) / 2
+    code = torch.bucketize(a, mids, right=True)
+    on_tie = (a.unsqueeze(-1) == mids).any(dim=-1)
+    tie_code = torch.bucketize(a, mids, right=False)
+    code = torch.where(on_tie, tie_code + (tie_code & 1), code)
+    return (sign | code.to(torch.uint8)).to(torch.uint8)
+
+
+def quantize_k_cache(input_k_cache, fmt):
+    """input_k_cache: (num_blocks, block_size, 1, 512) bf16 -> uint8 (nb, bs, 1, bpt)."""
+    d, d_nope, d_rope, tile_size, num_tiles = META[fmt]
+    num_blocks, block_size, h_k, _ = input_k_cache.shape
+    assert h_k == 1
+    x = input_k_cache.squeeze(2)
+    bpt = BPT[fmt]
+    result = torch.empty((num_blocks, block_size, bpt), dtype=torch.uint8)
+
+    if fmt == "v41_fp8":
+        data = result[..., : d_nope + d_rope].view(torch.float8_e4m3fn)
+        sf = result[..., d_nope + d_rope :]
+        for t in range(16):
+            sl = x[..., t * tile_size : (t + 1) * tile_size]
+            inv = torch.abs(sl).max(dim=-1).values.float() / 448.0
+            inv = _cast_scale_inv_to_ue8m0(inv)
+            # ue8m0 byte = bf16 bits >> 7 (Ascend path of tests/quant.py)
+            sf[:, :, t] = (inv.view(torch.int32) // 2**23).to(torch.uint8)
+            data[:, :, t * tile_size : (t + 1) * tile_size] = (sl.float() / inv.unsqueeze(-1)).to(
+                torch.float8_e4m3fn
+            )
+    else:
+        data = result[..., : d // 2]
+        sf = result[..., d // 2 :].view(torch.float8_e4m3fn)
+        xf = x.float()
+        amax = torch.nan_to_num(xf.abs(), nan=float("inf")).view(num_blocks, block_size, num_tiles, tile_size).amax(-1)
+        scale = torch.clamp(amax / 6.0, 2.0**-9, 448.0).to(torch.float8_e4m3fn)
+        normalized = xf.view(num_blocks, block_size, num_tiles, tile_size) / scale.float().unsqueeze(-1)
+        codes = _quantize_to_e2m1(normalized).view(num_blocks, block_size, d)
+        data[...] = codes[..., 0::2] | (codes[..., 1::2] << 4)
+        sf[...] = scale
+    return result.view(num_blocks, block_size, 1, bpt)
+
+
+def dequantize_k_cache(quant_k_cache, fmt):
+    d, d_nope, d_rope, tile_size, num_tiles = META[fmt]
+    num_blocks, block_size, h_k, _ = quant_k_cache.shape
+    assert h_k == 1
+    out = torch.empty((num_blocks, block_size, d), dtype=torch.bfloat16)
+    q = quant_k_cache.view(num_blocks, block_size, -1)
+    if fmt == "v41_fp8":
+        raw = q[..., : d_nope + d_rope].view(torch.float8_e4m3fn)
+        sf = q[..., d_nope + d_rope :]
+        for t in range(16):
+            cur = raw[..., t * tile_size : (t + 1) * tile_size].to(torch.bfloat16)
+            scales = (sf[:, :, t].to(torch.int32) * 2**7).to(torch.int16).view(torch.bfloat16)
+            out[..., t * tile_size : (t + 1) * tile_size] = cur * scales.unsqueeze(-1)
+    else:
+        data = q[..., : d // 2]
+        sf = q[..., d // 2 :].view(torch.float8_e4m3fn)
+        codes = torch.empty((num_blocks, block_size, d), dtype=torch.uint8)
+        codes[..., 0::2] = data & 0xF
+        codes[..., 1::2] = data >> 4
+        val = _E2M1[(codes & 7).long()]
+        val = torch.where((codes & 8) != 0, -val, val).view(num_blocks, block_size, num_tiles, tile_size)
+        out[...] = (val.to(torch.bfloat16) * sf.to(torch.bfloat16).unsqueeze(-1)).view(
+            num_blocks, block_size, d
+        )
+    return out.view(num_blocks, block_size, 1, d)
diff --git a/tests/neuron/flash_mla/test_deep_select.py b/tests/neuron/flash_mla/test_deep_select.py
new file mode 100644
index 0000000..f1e6178
--- /dev/null
+++ b/tests/neuron/flash_mla/test_deep_select.py
@@ -0,0 +1,119 @@
+"""CPU correctness tests for the pure-torch deep_select vendor package.
+
+The checks are ported from vendor/DeepSelect/tests/test.py (NaN guard, index range, uniqueness,
+gathered value, topk condition, sorted index / sorted value).
+"""
+
+import itertools
+import sys
+
+import torch
+
+from deep_select.interface import get_stride_requirement, topk as ds_topk
+
+torch.manual_seed(0)
+FAILS = []
+NAN_SENTINEL = 0x3F3F3F3F
+
+
+def chk(name, cond):
+    print(f"[{'ok ' if cond else 'FAIL'}] {name}")
+    if not cond:
+        FAILS.append(name)
+
+
+def run_case(b, vocab, k, dtype, idx_dtype, sorted_value, sorted_index, return_value,
+             use_end, use_offset, with_nan=False, use_out=False):
+    tag = (f"b{b} v{vocab} k{k} {str(dtype).split('.')[-1]} {str(idx_dtype).split('.')[-1]} "
+           f"sv{int(sorted_value)} si{int(sorted_index)} rv{int(return_value)} "
+           f"end{int(use_end)} off{int(use_offset)} nan{int(with_nan)} out{int(use_out)}")
+    x = (torch.randn(b, vocab) * 3).to(dtype)
+    end = torch.randint(1, vocab + 1, (b,), dtype=torch.int32) if use_end else None
+    if use_end:
+        end[0] = min(k // 2 + 1, vocab)  # a short row, fewer valid entries than k
+    offset = torch.randint(0, 1000, (b,), dtype=torch.int32) if use_offset else None
+    nan_rows = torch.zeros(b, dtype=torch.bool)
+    if with_nan:
+        nan_rows[0] = True
+        x[0, 0] = float("nan")
+    out_buf = torch.empty((b, k), dtype=idx_dtype) if use_out else None
+
+    val, idx = ds_topk(
+        x, k, sorted=sorted_value, end=end, indices_type=idx_dtype, sorted_index=sorted_index,
+        output_idx=out_buf, output_idx_offset=offset, idx_oob_fill_value=-1,
+        value_oob_fill_value=float("-inf"), return_value=return_value, abort_when_nan_found=False,
+    )
+    chk(f"{tag}: dtypes", idx.dtype is idx_dtype and (val is None) == (not return_value)
+        and (val is None or val.dtype is dtype))
+    chk(f"{tag}: shapes", tuple(idx.shape) == (b, k) and (val is None or tuple(val.shape) == (b, k)))
+    if use_out:
+        chk(f"{tag}: writes output_idx in place", idx.data_ptr() == out_buf.data_ptr())
+
+    valid_len = (torch.full((b,), vocab, dtype=torch.int64) if end is None else end.to(torch.int64))
+    counts = torch.clamp(valid_len, max=k)
+    sel = torch.arange(k).unsqueeze(0) < counts.unsqueeze(1)
+
+    # NaN guard: whole row filled with the sentinel (Ascend behaviour)
+    if with_nan:
+        chk(f"{tag}: NaN row sentinel", bool((idx[nan_rows] == NAN_SENTINEL).all()))
+    sel = sel & ~nan_rows.unsqueeze(1)
+
+    raw = idx.to(torch.int64)
+    if offset is not None:
+        raw = raw - offset.to(torch.int64).unsqueeze(1)
+    in_range = (raw >= 0) & (raw < valid_len.unsqueeze(1))
+    chk(f"{tag}: index range", bool((in_range | ~sel).all()))
+    chk(f"{tag}: padding fill", bool((idx[~sel & ~nan_rows.unsqueeze(1)] == -1).all()))
+
+    srt = torch.sort(torch.where(sel, raw, torch.iinfo(torch.int64).max), dim=1).values
+    dup = (srt[:, 1:] == srt[:, :-1]) & sel[:, 1:]
+    chk(f"{tag}: unique index", not bool(dup.any()))
+
+    safe = torch.where(sel & in_range, raw, torch.zeros_like(raw))
+    gathered = x.gather(1, safe)
+    rows = ~nan_rows
+    if rows.any():
+        if val is not None:
+            exp = gathered[rows].masked_fill(~sel[rows], float("-inf"))
+            chk(f"{tag}: gathered value bitwise", torch.equal(val[rows], exp))
+        gv = gathered[rows].float()
+        smin = gv.masked_fill(~sel[rows], float("inf")).amin(dim=1)
+        unsel = x[rows].float().clone()
+        unsel.scatter_(1, safe[rows], float("-inf"))
+        col = torch.arange(vocab).unsqueeze(0)
+        unsel = torch.where(col >= valid_len[rows].unsqueeze(1), torch.full_like(unsel, float("-inf")), unsel)
+        unsel = torch.where(unsel != unsel, torch.full_like(unsel, float("-inf")), unsel)
+        chk(f"{tag}: topk condition", bool((smin >= unsel.amax(dim=1)).all()))
+
+    ordered = sel[:, 1:]
+    if sorted_index:
+        ok = (idx[:, 1:].to(torch.int64) >= idx[:, :-1].to(torch.int64)) | ~ordered
+        chk(f"{tag}: sorted index", bool(ok.all()))
+    if sorted_value and val is not None:
+        ok = (val[:, :-1].float() >= val[:, 1:].float()) | ~ordered
+        chk(f"{tag}: sorted value", bool(ok.all()))
+
+
+if __name__ == "__main__":
+    chk("stride requirement == (32, 32)", get_stride_requirement() == (32, 32))
+
+    combos = [(False, False, False), (False, False, True), (True, False, True),
+              (False, True, True), (False, True, False)]
+    for dtype in (torch.float32, torch.bfloat16):
+        for idx_dtype in (torch.int32, torch.int64):
+            for sv, si, rv in combos:
+                if sv and dtype is torch.bfloat16:
+                    continue  # sorted_value is fp32-only upstream
+                run_case(5, 97, 8, dtype, idx_dtype, sv, si, rv, use_end=True, use_offset=True)
+    for use_end, use_offset in itertools.product((False, True), repeat=2):
+        run_case(3, 64, 16, torch.float32, torch.int32, False, True, False, use_end, use_offset)
+    run_case(4, 33, 8, torch.float32, torch.int32, False, True, False, True, True, with_nan=True)
+    run_case(4, 33, 8, torch.float32, torch.int32, True, False, True, False, False, with_nan=True)
+    run_case(3, 16, 8, torch.float32, torch.int32, False, True, False, True, True, use_out=True)
+    # k larger than the vocab / every row short
+    run_case(2, 6, 8, torch.float32, torch.int32, False, True, False, True, False)
+    # DeepInfer indexer call shape: big-ish k, end, offset, sorted_index, no values
+    run_case(2, 512, 128, torch.float32, torch.int32, False, True, False, True, True)
+
+    print("\nFAILURES:", FAILS if FAILS else "none")
+    sys.exit(1 if FAILS else 0)
diff --git a/tests/neuron/flash_mla/test_flash_mla.py b/tests/neuron/flash_mla/test_flash_mla.py
new file mode 100644
index 0000000..fdd6458
--- /dev/null
+++ b/tests/neuron/flash_mla/test_flash_mla.py
@@ -0,0 +1,316 @@
+"""CPU correctness tests for the pure-torch flash_mla vendor package.
+
+Run: python -I test_flash_mla.py   (PYTHONPATH must contain the vendor dir)
+"""
+
+import sys
+
+import torch
+
+import flash_mla
+from flash_mla import flash_mla_sparse_fwd, flash_mla_with_kvcache, get_mla_metadata
+from flash_mla.flash_mla_interface import FlashMLASchedMeta
+
+import ref_attn
+import ref_quant
+
+torch.manual_seed(0)
+FAILS = []
+
+
+def check(name, got, exp, atol=2e-2, rtol=2e-2, exact=False):
+    assert got.shape == exp.shape, f"{name}: shape {tuple(got.shape)} != {tuple(exp.shape)}"
+    g, e = got.float(), exp.float()
+    if exact:
+        ok = torch.equal(g, e)
+        msg = "bitwise"
+    else:
+        same_inf = ((g == e) | ((g != g) & (e != e))).float()
+        finite = torch.isfinite(g) & torch.isfinite(e)
+        ok = bool(torch.all(torch.where(finite, (g - e).abs() <= atol + rtol * e.abs(), same_inf.bool())))
+        msg = f"max abs diff {(g - e).abs()[finite].max().item() if finite.any() else 0.0:.3e}"
+    print(f"[{'ok ' if ok else 'FAIL'}] {name} ({msg})")
+    if not ok:
+        FAILS.append(name)
+
+
+def make_indices(rows, topk, s_kv, invalid_frac=0.2, invalid_value=-1):
+    idx = torch.randint(0, s_kv, (rows, topk), dtype=torch.int32)
+    mask = torch.rand(rows, topk) < invalid_frac
+    idx = torch.where(mask, torch.full_like(idx, invalid_value), idx)
+    return idx
+
+
+# ---------------------------------------------------------------- prefill
+def test_prefill(h_q=64, s_q=7, s_kv=41, topk=24, with_sink=True, with_len=True, d=512):
+    q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
+    kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
+    idx = make_indices(s_q, topk, s_kv).unsqueeze(1)
+    # also exercise the ">= s_kv" invalid convention of the prefill kernel
+    idx[0, 0, 0] = s_kv + 3
+    sink = torch.randn(h_q) if with_sink else None
+    tl = torch.randint(1, topk + 1, (s_q,), dtype=torch.int32) if with_len else None
+    sm = d**-0.5
+
+    out, ml, lse = flash_mla_sparse_fwd(q, kv, idx, sm, 512, sink, tl)
+    r_out, r_ml, r_lse = ref_attn.ref_sparse_attn_fwd(q, kv, idx, sm, tl, sink)
+    tag = f"prefill h{h_q} sink={with_sink} len={with_len}"
+    check(f"{tag} out", out, r_out)
+    check(f"{tag} max_logits", ml, r_ml)
+    check(f"{tag} lse", lse, r_lse)
+    assert out.dtype == torch.bfloat16 and lse.dtype == torch.float32 and ml.dtype == torch.float32
+
+
+def test_prefill_lonely():
+    h_q, s_q, s_kv, topk, d = 64, 3, 16, 8, 512
+    q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
+    kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
+    idx = torch.full((s_q, 1, topk), -1, dtype=torch.int32)
+    out, ml, lse = flash_mla_sparse_fwd(q, kv, idx, d**-0.5)
+    check("prefill lonely out==0", out, torch.zeros_like(out), exact=True)
+    check("prefill lonely lse==+inf", lse, torch.full_like(lse, float("inf")), exact=True)
+    check("prefill lonely max_logits==-inf", ml, torch.full_like(ml, float("-inf")), exact=True)
+
+
+# ---------------------------------------------------------------- kv cache
+def test_kv_dequant():
+    for fmt in ("v41_fp8", "v41_fp4"):
+        raw = (torch.randn(5, 8, 1, 512) / 4).to(torch.bfloat16)
+        cache = ref_quant.quantize_k_cache(raw, fmt)
+        exp = ref_quant.dequantize_k_cache(cache, fmt).view(5, 8, 512)
+        got = flash_mla.kv_cache.dequantize_records(cache.view(5, 8, -1), fmt)
+        check(f"dequant {fmt}", got, exp, exact=True)
+        # relative error against the bf16 source
+        err = ((got.float() - raw.view(5, 8, 512).float()).abs() / raw.view(5, 8, 512).float().abs().clamp_min(1e-3)).median()
+        print(f"       {fmt} median rel err vs bf16 source: {err.item():.4f}")
+        # byte-view equivalence: int8 / float8 views of the same bytes must dequantize identically
+        got_i8 = flash_mla.kv_cache.dequantize_records(cache.view(5, 8, -1).view(torch.int8), fmt)
+        check(f"dequant {fmt} int8 view", got_i8, exp, exact=True)
+
+
+# ---------------------------------------------------------------- decode
+def test_decode(h_q=64, b=5, topk=20, block=4, nblocks=9, with_sink=True, with_len=True, extra=False, strided=False):
+    d, s_q = 512, 1
+    raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
+    cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
+    if strided:
+        # non-contiguous block stride, like deepinfer/worker/cache_engine's as_strided views
+        pad = torch.zeros(nblocks, block * 528 + 64, dtype=torch.uint8)
+        pad[:, : block * 528] = cache.reshape(nblocks, -1)
+        cache = pad.as_strided((nblocks, block, 1, 528), (pad.stride(0), 528, 528, 1))
+    kv_bf16 = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)
+
+    q = (torch.randn(b, s_q, h_q, d) / 8).to(torch.bfloat16)
+    idx = make_indices(b * s_q, topk, nblocks * block).view(b, s_q, topk)
+    sink = torch.randn(h_q) if with_sink else None
+    tl = torch.randint(1, topk + 1, (b,), dtype=torch.int32) if with_len else None
+    sm = d**-0.5
+
+    scopes = [(kv_bf16, idx, tl)]
+    e_cache = e_idx = e_tl = None
+    if extra:
+        e_raw = (torch.randn(6, block, 1, d) / 8).to(torch.bfloat16)
+        e_cache = ref_quant.quantize_k_cache(e_raw, "v41_fp4")
+        e_kv = ref_quant.dequantize_k_cache(e_cache, "v41_fp4").view(-1, d)
+        e_topk = 12
+        e_idx = make_indices(b * s_q, e_topk, 6 * block).view(b, s_q, e_topk)
+        e_tl = torch.randint(1, e_topk + 1, (b,), dtype=torch.int32)
+        scopes.append((e_kv, e_idx, e_tl))
+
+    meta, none = get_mla_metadata()
+    assert none is None and isinstance(meta, FlashMLASchedMeta)
+    out, lse = flash_mla_with_kvcache(
+        q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
+        tile_scheduler_metadata=meta, softmax_scale=sm, causal=False, is_fp8_kvcache=True,
+        indices=idx, attn_sink=sink, extra_k_cache=e_cache, extra_indices_in_kvcache=e_idx,
+        topk_length=tl, extra_topk_length=e_tl,
+    )
+    r_out, r_lse = ref_attn.ref_sparse_attn_decode(q, scopes, sm, sink)
+    tag = f"decode h{h_q} sink={with_sink} len={with_len} extra={extra} strided={strided}"
+    check(f"{tag} out", out, r_out)
+    check(f"{tag} lse", lse, r_lse)
+    assert meta.have_initialized and meta.config.topk == topk and meta.config.page_block_size == block
+    # reuse with a different topk must be rejected
+    try:
+        flash_mla_with_kvcache(
+            q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
+            tile_scheduler_metadata=meta, softmax_scale=sm, indices=idx[..., :-1],
+        )
+        print("[FAIL] sched_meta reuse check")
+        FAILS.append("sched_meta reuse")
+    except AssertionError:
+        print("[ok ] sched_meta reuse rejected on topk change")
+
+
+def test_decode_multi_token():
+    """s_q > 1 (MTP/DFlash style): topk_length is per request, broadcast over s_q."""
+    d, h_q, b, s_q, topk, block, nblocks = 512, 64, 3, 4, 16, 8, 7
+    raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
+    cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
+    kv_bf16 = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)
+    q = (torch.randn(b, s_q, h_q, d) / 8).to(torch.bfloat16)
+    idx = make_indices(b * s_q, topk, nblocks * block).view(b, s_q, topk)
+    tl = torch.randint(1, topk + 1, (b,), dtype=torch.int32)
+    sm = d**-0.5
+    out, lse = flash_mla_with_kvcache(
+        q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
+        tile_scheduler_metadata=FlashMLASchedMeta(), softmax_scale=sm, indices=idx, topk_length=tl,
+    )
+    r_out, r_lse = ref_attn.ref_sparse_attn_decode(q, [(kv_bf16, idx, tl)], sm)
+    check("decode s_q=4 out", out, r_out)
+    check("decode s_q=4 lse", lse, r_lse)
+
+
+def test_decode_lonely():
+    d, h_q, b, topk, block, nblocks = 512, 64, 2, 8, 4, 4
+    raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
+    cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
+    q = (torch.randn(b, 1, h_q, d) / 8).to(torch.bfloat16)
+    idx = torch.full((b, 1, topk), -1, dtype=torch.int32)
+    out, lse = flash_mla_with_kvcache(
+        q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
+        tile_scheduler_metadata=FlashMLASchedMeta(), softmax_scale=d**-0.5, indices=idx,
+        attn_sink=torch.randn(h_q),
+    )
+    check("decode lonely out==0", out, torch.zeros_like(out), exact=True)
+    check("decode lonely lse==+inf", lse, torch.full_like(lse, float("inf")), exact=True)
+
+
+def test_sink_infinities():
+    d, h_q, s_q, s_kv, topk = 512, 64, 4, 20, 8
+    q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
+    kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
+    idx = torch.randint(0, s_kv, (s_q, 1, topk), dtype=torch.int32)
+    sink = torch.full((h_q,), float("-inf"))
+    sink[: h_q // 2] = float("inf")
+    out, _, lse = flash_mla_sparse_fwd(q, kv, idx, d**-0.5, 512, sink)
+    base, _, base_lse = flash_mla_sparse_fwd(q, kv, idx, d**-0.5, 512, None)
+    check("sink=+inf -> out 0", out[:, : h_q // 2], torch.zeros_like(out[:, : h_q // 2]), exact=True)
+    check("sink=-inf -> unchanged", out[:, h_q // 2 :], base[:, h_q // 2 :], atol=1e-2, rtol=1e-2)
+    check("sink does not change lse", lse, base_lse, exact=True)
+
+
+def test_dense_unsupported():
+    try:
+        flash_mla.flash_attn_varlen_func(None, None, None, None, None, 0, 0)
+        FAILS.append("dense should raise")
+        print("[FAIL] dense entry point should raise")
+    except RuntimeError as e:
+        print(f"[ok ] dense entry point raises: {str(e)[:60]}...")
+
+
+def test_fused_gated():
+    import flash_mla.fused_norm_rope_attn_rope_cast as fused
+
+    for fn, args in ((fused.permute_q_b_proj, ((None, None), 64, 512)), (fused.permute_wv_proj, ((None, None), 1, 512))):
+        try:
+            fn(*args)
+            FAILS.append("permute should raise")
+        except RuntimeError as e:
+            print(f"[ok ] {fn.__name__} raises: {e}")
+
+    # the fused attention core itself (norm + rope + attn + rope), checked against the test ref
+    d, h_q, s_q, s_kv, topk = 512, 64, 5, 32, 16
+    q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
+    kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
+    idx = make_indices(s_q, topk, s_kv).unsqueeze(1)
+    positions = torch.randint(0, 1024, (s_q,), dtype=torch.int32)
+    cache = torch.randn(1024, 64)
+    rms = fused._rms_norm_scale_factor(q, 1e-4)
+    q_roped = fused._rope(q, False, cache, positions, 64).to(torch.bfloat16)
+    out, ml, lse = flash_mla_sparse_fwd(q_roped, kv, idx, d**-0.5, 512, None, None, _rms_norm_scale_factor=rms)
+    out = fused._rope(out, True, cache, positions, 64)
+    r_out, r_ml, r_lse = ref_attn.ref_sparse_attn_fwd(q_roped, kv, idx, d**-0.5, None, None, rms=rms)
+    # reference O-RoPE, computed independently
+    exp = r_out.float().clone()
+    cos = cache[positions.long(), :32].unsqueeze(1)
+    sin = -cache[positions.long(), 32:].unsqueeze(1)
+    rp = exp[..., -64:]
+    x0, x1 = rp[..., 0::2].clone(), rp[..., 1::2].clone()
+    rp[..., 0::2] = x0 * cos - x1 * sin
+    rp[..., 1::2] = x0 * sin + x1 * cos
+    check("fused norm+rope+attn+rope out", out, exp, atol=3e-2, rtol=3e-2)
+    check("fused lse", lse, r_lse)
+
+
+def test_extra_fp8_cache():
+    """SWA main cache + a second fp8 (528 B) cache as `extra`, the non-fp4 DeepInfer variant."""
+    d, h_q, b, topk, block = 512, 64, 3, 12, 4
+    raw = (torch.randn(5, block, 1, d) / 8).to(torch.bfloat16)
+    cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
+    kv = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)
+    e_raw = (torch.randn(7, block, 1, d) / 8).to(torch.bfloat16)
+    e_cache = ref_quant.quantize_k_cache(e_raw, "v41_fp8")
+    e_kv = ref_quant.dequantize_k_cache(e_cache, "v41_fp8").view(-1, d)
+    q = (torch.randn(b, 1, h_q, d) / 8).to(torch.bfloat16)
+    idx = make_indices(b, topk, 5 * block).view(b, 1, topk)
+    e_idx = make_indices(b, 9, 7 * block).view(b, 1, 9)
+    out, lse = flash_mla_with_kvcache(
+        q=q, k_cache=cache, block_table=None, cache_seqlens=None, head_dim_v=512,
+        tile_scheduler_metadata=FlashMLASchedMeta(), softmax_scale=d**-0.5, indices=idx,
+        extra_k_cache=e_cache, extra_indices_in_kvcache=e_idx,
+    )
+    r_out, r_lse = ref_attn.ref_sparse_attn_decode(q, [(kv, idx, None), (e_kv, e_idx, None)], d**-0.5)
+    check("decode extra fp8 cache out", out, r_out)
+    check("decode extra fp8 cache lse", lse, r_lse)
+
+
+def test_attention_backend_integration():
+    """Drive deepinfer/model_executor/layers/attention_backend.py (head padding path included)."""
+    import importlib.util
+    import os
+
+    path = os.path.join(os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..")), "deepinfer", "model_executor", "layers", "attention_backend.py")
+    spec = importlib.util.spec_from_file_location("ab", path)
+    ab = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(ab)
+
+    d, h_q, s_q, s_kv, topk = 512, 40, 6, 30, 16  # 40 heads -> padded to 64
+    q = (torch.randn(s_q, h_q, d) / 8).to(torch.bfloat16)
+    kv = (torch.randn(s_kv, 1, d) / 8).to(torch.bfloat16)
+    idx = make_indices(s_q, topk, s_kv, invalid_value=-1)
+    sink = torch.randn(h_q)
+    ilen = torch.randint(1, topk + 1, (s_q,), dtype=torch.int32)
+    out, lse = ab.prefill_attn_varlen_sparse(q, kv, idx, d**-0.5, attn_sink=sink, index_length=ilen)
+    pq = torch.cat([q, q.new_zeros(s_q, 24, d)], dim=1)
+    ps = torch.cat([sink, sink.new_zeros(24)])
+    r_out, _, r_lse = ref_attn.ref_sparse_attn_fwd(pq, kv, idx.unsqueeze(1), d**-0.5, ilen, ps)
+    check("attention_backend prefill out (head pad 40->64)", out, r_out[:, :h_q])
+    check("attention_backend prefill lse", lse, r_lse[:, :h_q])
+
+    b, block, nblocks = 4, 8, 6
+    raw = (torch.randn(nblocks, block, 1, d) / 8).to(torch.bfloat16)
+    cache = ref_quant.quantize_k_cache(raw, "v41_fp8")
+    kvb = ref_quant.dequantize_k_cache(cache, "v41_fp8").view(-1, d)
+    qd = (torch.randn(b, h_q, d) / 8).to(torch.bfloat16)
+    didx = make_indices(b, topk, nblocks * block).view(b, topk)
+    dlen = torch.randint(1, topk + 1, (b,), dtype=torch.int32)
+    out, lse = ab.decode_mla_mqa_sparse(
+        qd, cache, 512, didx, FlashMLASchedMeta(), d**-0.5, attn_sink=sink, index_length=dlen
+    )
+    pqd = torch.cat([qd, qd.new_zeros(b, 24, d)], dim=1).view(b, 1, 64, d)
+    r_out, r_lse = ref_attn.ref_sparse_attn_decode(pqd, [(kvb, didx.view(b, 1, topk), dlen)], d**-0.5, ps)
+    check("attention_backend decode out (head pad)", out, r_out.view(b, 64, d)[:, :h_q])
+    check("attention_backend decode lse", lse, r_lse.view(b, 64)[:, :h_q])
+
+
+if __name__ == "__main__":
+    test_kv_dequant()
+    for h in (64, 128):
+        test_prefill(h_q=h)
+    test_prefill(with_sink=False, with_len=False)
+    test_prefill_lonely()
+    test_decode()
+    test_decode(with_sink=False, with_len=False)
+    test_decode(extra=True)
+    test_decode(strided=True)
+    test_decode(h_q=128)
+    test_decode_multi_token()
+    test_decode_lonely()
+    test_sink_infinities()
+    test_extra_fp8_cache()
+    test_attention_backend_integration()
+    test_dense_unsupported()
+    test_fused_gated()
+    print("\nFAILURES:", FAILS if FAILS else "none")
+    sys.exit(1 if FAILS else 0)
diff --git a/tests/neuron/flash_mla/test_tilelang.py b/tests/neuron/flash_mla/test_tilelang.py
new file mode 100644
index 0000000..618fd2b
--- /dev/null
+++ b/tests/neuron/flash_mla/test_tilelang.py
@@ -0,0 +1,38 @@
+"""Checks the tilelang stub provides exactly what platform-independent DeepInfer code uses."""
+
+import sys
+
+import torch
+
+import tilelang
+import tilelang.language as T
+
+FAILS = []
+
+
+def chk(name, cond):
+    print(f"[{'ok ' if cond else 'FAIL'}] {name}")
+    if not cond:
+        FAILS.append(name)
+
+
+chk("PassConfigKey.TL_ENABLE_FAST_MATH", isinstance(tilelang.PassConfigKey.TL_ENABLE_FAST_MATH, str))
+chk("T.float8_e4m3fn", T.float8_e4m3fn == "float8_e4m3fn")
+chk("dtype map", tilelang.TL_DTYPE_TO_TORCH[T.float8_e4m3fn] is torch.float8_e4m3fn)
+chk("norm.py dtype strings", T.float32 == "float" and T.bfloat16 == "bfloat16")
+chk("max_value e4m3", T.max_value(T.float8_e4m3fn) == 448.0)
+chk("min_value fp32", T.min_value(T.float32) == torch.finfo(torch.float32).min)
+
+# the pass-config dict pattern of tilelang_ops/compress.py
+cfg = {}
+cfg[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] = True
+chk("pass_configs assignment", cfg[tilelang.PassConfigKey.TL_ENABLE_FAST_MATH] is True)
+
+try:
+    T.some_dsl_primitive
+    FAILS.append("unknown attr should raise")
+except AttributeError as e:
+    chk("unknown T attribute raises AttributeError", "Neuron stub" in str(e))
+
+print("\nFAILURES:", FAILS if FAILS else "none")
+sys.exit(1 if FAILS else 0)
diff --git a/tests/neuron/nie_skip.py b/tests/neuron/nie_skip.py
new file mode 100644
index 0000000..9e59e2a
--- /dev/null
+++ b/tests/neuron/nie_skip.py
@@ -0,0 +1,17 @@
+"""pytest plugin: report kernels stubbed out in Neuron basic mode as skipped, not failed.
+
+The Neuron backend exports every asc kernel factory, but those reachable only from DSpark/MTP, ZOS, CP or
+VL raise ``NotImplementedError("<name>: not supported on Neuron (basic mode)")``. The upstream interface
+tests still parametrize over them; with ``-p nie_skip`` they show up as skips.
+"""
+
+import pytest
+
+
+@pytest.hookimpl(hookwrapper=True)
+def pytest_runtest_makereport(item, call):
+    outcome = yield
+    report = outcome.get_result()
+    if call.excinfo is not None and call.excinfo.errisinstance(NotImplementedError) and "basic mode" in str(call.excinfo.value):
+        report.outcome = "skipped"
+        report.longrepr = (str(item.path), item.location[1] or 0, f"Neuron basic mode: {call.excinfo.value}")
diff --git a/tests/neuron/run_all.sh b/tests/neuron/run_all.sh
new file mode 100755
index 0000000..7fc2a59
--- /dev/null
+++ b/tests/neuron/run_all.sh
@@ -0,0 +1,114 @@
+#!/bin/bash
+# Neuron port test suite. Run after tests/neuron/setup_trn2.sh.
+#
+#   bash tests/neuron/run_all.sh cpu              # host-only: kernel unit tests + end-to-end vs golden (~1 h)
+#   CORES=0-7 bash tests/neuron/run_all.sh device # NeuronCore tests on the given logical cores
+#   CORES=0-7 bash tests/neuron/run_all.sh all
+#   bash tests/neuron/run_all.sh cpu deep_gemm flash_mla   # only selected stages
+#
+# Device stages use the first core of CORES for single-core runs and CORES for multi-rank deep_ep
+# (world sizes 2/4/8 limited by the core count). They refuse to start if those cores are in use.
+# Logs: $WORK/test_logs/<timestamp>/<stage>.log; a PASS/FAIL table is printed at the end.
+set -uo pipefail
+
+REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+WORK="${WORK:-$HOME/di}"
+MODE="${1:-cpu}"; shift || true
+ONLY=("$@")
+CORES="${CORES:-0-7}"
+LOGDIR="$WORK/test_logs/$(date +%Y%m%d-%H%M%S)"
+mkdir -p "$LOGDIR"
+# shellcheck disable=SC1091
+source "$WORK/venv/bin/activate"
+cd "$REPO" || exit 1
+VENDOR="$REPO/deepinfer/neuron_ops/vendor"
+T="$REPO/tests/neuron"
+export DEEPINFER_ACCELERATOR=neuron
+declare -a SUMMARY=()
+
+first_core() { echo "${CORES%%[-,]*}"; }
+core_count() { python3 - "$CORES" <<'EOF'
+import sys
+n = 0
+for part in sys.argv[1].split(","):
+    a, _, b = part.partition("-")
+    n += (int(b) - int(a) + 1) if b else 1
+print(n)
+EOF
+}
+selected() { [ ${#ONLY[@]} -eq 0 ] && return 0; for s in "${ONLY[@]}"; do [ "$s" = "$1" ] && return 0; done; return 1; }
+
+run() {  # run <stage> <command...>
+  local stage="$1"; shift
+  selected "$stage" || return 0
+  echo "=== $stage"
+  local t0=$SECONDS
+  if "$@" > "$LOGDIR/$stage.log" 2>&1; then r=PASS; else r=FAIL; fi
+  SUMMARY+=("$(printf '%-6s %-28s %5ss  %s' "$r" "$stage" $((SECONDS - t0)) "$LOGDIR/$stage.log")")
+  echo "    $r ($((SECONDS - t0))s)"
+  [ "$r" = FAIL ] && tail -5 "$LOGDIR/$stage.log" | sed 's/^/    | /'
+}
+
+cores_free() {
+  local busy
+  busy=$(neuron-ls 2>/dev/null | grep -cE "python|vLLM|Worker|torchrun")
+  [ "$busy" = "0" ] || { echo "NeuronCores busy ($busy processes); not starting device stages"; return 1; }
+}
+
+cpu_stages() {
+  export DEEPINFER_NEURON_DEVICE=cpu
+  run deep_gemm      env PYTHONPATH="$VENDOR" python "$T/deep_gemm/test_deep_gemm_neuron.py"
+  run deep_gemm_tk   env PYTHONPATH="$VENDOR" python "$T/deep_gemm/test_interop_tile_kernels.py"
+  run flash_mla      bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_flash_mla.py"
+  run deep_select    bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_deep_select.py"
+  run tilelang_stub  bash -c "cd $T/flash_mla && PYTHONPATH=$VENDOR python test_tilelang.py"
+  run tile_kernels   python -m pytest -q -p no:cacheprovider "$T/tile_kernels/test_tile_kernels.py"
+  for ws in 1 2 4; do
+    run "deep_ep_ws$ws"     python "$T/deep_ep/test_ep.py" --num-processes $ws --port $((29531 + ws))
+    run "engram_ws$ws"      python "$T/deep_ep/test_engram.py" --num-processes $ws --port $((29541 + ws))
+    run "bucket_ws$ws"      python "$T/deep_ep/test_bucket.py" --num-processes $ws --port $((29551 + ws))
+  done
+  run sampling_kat   env PYTHONPATH="$REPO" python "$T/sampling/kat.py"
+  run sampling_check env PYTHONPATH="$REPO" python "$T/sampling/check2.py"
+  # Upstream interface tests of tilelang_ops (compare the Neuron backend against the *_ref implementations).
+  for m in norm store_kv_cache load_paged_kv_cache fused_transform compress state_cache get_kv_index kv_index_cat \
+           swa_index candidate_blocks flash_mla_utils gather prepare_v4_decode_metadata random_sampling fused_sample \
+           sampler sum_smaller_probs structure_output; do
+    run "tl_$m" env PYTHONPATH="$T" python -m pytest -q -p no:cacheprovider --noconftest -p nie_skip \
+      "deepinfer/model_executor/layers/tilelang_ops/$m.py" -k "not benchmark and not 129280-4096"
+  done
+  run e2e_cpu python examples/neuron/run_deepinfer_tiny.py --device cpu --threads 32 --num-kv-blocks 32 \
+    --out "$WORK/e2e_cpu" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt" --compare
+}
+
+device_stages() {
+  cores_free || return 0
+  local c0 n
+  c0="$(first_core)"; n="$(core_count)"
+  export DEEPINFER_NEURON_DEVICE=neuron
+  run device_primitives env NEURON_RT_VISIBLE_CORES="$c0" python "$T/device/test_device_primitives.py"
+  for ws in 2 4 8; do
+    [ "$ws" -le "$n" ] || continue
+    run "deep_ep_neuron_ws$ws" env DEEP_EP_TEST_DEVICE=neuron DEEP_EP_TEST_CORE_BASE="$c0" \
+      python "$T/deep_ep/test_ep.py" --num-processes $ws --port $((29631 + ws))
+    run "engram_neuron_ws$ws" env DEEP_EP_TEST_DEVICE=neuron DEEP_EP_TEST_CORE_BASE="$c0" \
+      python "$T/deep_ep/test_engram.py" --num-processes $ws --port $((29641 + ws))
+  done
+  run e2e_neuron env NEURON_RT_VISIBLE_CORES="$c0" CUDA_VISIBLE_DEVICES="$c0" MASTER_PORT=29733 \
+    python examples/neuron/run_deepinfer_tiny.py --device neuron --num-kv-blocks 32 \
+    --out "$WORK/e2e_neuron" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt" --compare
+  run device_diff env NEURON_RT_VISIBLE_CORES="$c0" CUDA_VISIBLE_DEVICES="$c0" MASTER_PORT=29734 \
+    python examples/neuron/device_diff.py --prompts short --num-kv-blocks 32 \
+    --out "$WORK/device_diff" --golden-dir "$WORK/golden" --hf-ckpt "$WORK/tiny_ckpt"
+}
+
+case "$MODE" in
+  cpu) cpu_stages ;;
+  device) device_stages ;;
+  all) cpu_stages; device_stages ;;
+  *) echo "usage: $0 cpu|device|all [stage...]"; exit 2 ;;
+esac
+
+echo; echo "===== summary ($LOGDIR)"
+printf '%s\n' "${SUMMARY[@]}"
+! printf '%s\n' "${SUMMARY[@]}" | grep -q '^FAIL'
diff --git a/tests/neuron/sampling/check2.py b/tests/neuron/sampling/check2.py
new file mode 100644
index 0000000..ab4f03b
--- /dev/null
+++ b/tests/neuron/sampling/check2.py
@@ -0,0 +1,36 @@
+import os
+import torch
+from deepinfer.model_executor.layers.tilelang_ops.neuron.random_sampling import (
+    two_level_sample, quantize_weights, acc_slots_for, rng_first_two, rng_sequence)
+from deepinfer.model_executor.layers.tilelang_ops.random_sampling import random_sampling
+from deepinfer.model_executor.layers.tilelang_ops.sampler import prepare_verification_metadata, prepare_verification_metadata_ref
+
+for V in [8, 100, 1024, 4032, 8192, 129280]:
+    Q = 64
+    p = torch.rand(Q, V) ** 8; p[:, ::3] = 0; p = p / p.sum(1, keepdim=True)
+    seeds = torch.randint(-2**62, 2**62, (Q,), dtype=torch.int64)
+    pos = torch.randint(-5, 10**9, (Q,), dtype=torch.int32)
+    out = two_level_sample(p, seeds, pos, 1, V)
+    # independent: inverse CDF over columns ordered by (col % AS, col // AS)
+    AS = acc_slots_for(V)
+    fx = quantize_weights(p).to(torch.int64)
+    cols = torch.arange(V)
+    order = torch.argsort((cols % AS) * (V + 1) + cols // AS)
+    cdf = fx[:, order].cumsum(1)
+    total = cdf[:, -1]
+    lo, hi = rng_first_two(seeds, rng_sequence(pos, 1))
+    r = torch.tensor([((int(h) << 32) | int(l)) % int(t) for h, l, t in zip(hi, lo, total)])
+    exp = order[(cdf <= r[:, None]).sum(1)]
+    assert torch.equal(out.long(), exp), V
+    assert (p.gather(1, out.long()[:, None]) > 0).all()
+print("two-level == slot-ordered inverse CDF: OK")
+# strided probs rows
+p = torch.softmax(torch.randn(16, 2048), -1); big = torch.zeros(16, 4096); big[:, :2048] = p
+a = random_sampling(big[:, :2048], torch.arange(16).to(torch.uint64), torch.zeros(16, dtype=torch.int32))
+b = random_sampling(p.contiguous(), torch.arange(16).to(torch.uint64), torch.zeros(16, dtype=torch.int32))
+assert torch.equal(a, b); print("strided rows OK")
+# nextn=0 metadata path (num_steps=0) and seqlen=1 rows
+seqlen = torch.ones(37, dtype=torch.int32); ftp = torch.arange(37, dtype=torch.int32)
+for got, ref in zip(prepare_verification_metadata(ftp, seqlen, 37, 0), prepare_verification_metadata_ref(ftp, seqlen, 37)):
+    assert torch.equal(got, ref)
+print("prepare_verification_metadata num_steps=0 OK")
diff --git a/tests/neuron/sampling/kat.py b/tests/neuron/sampling/kat.py
new file mode 100644
index 0000000..2f5542e
--- /dev/null
+++ b/tests/neuron/sampling/kat.py
@@ -0,0 +1,11 @@
+import torch
+from deepinfer.model_executor.layers.tilelang_ops.neuron.random_sampling import philox4x32_10
+def t(v): return torch.tensor([v], dtype=torch.int64)
+for ctr, key, exp in [((0,0,0,0),(0,0),(0x6627e8d5,0xe169c58d,0xbc57ac4c,0x9b00dbd8)),
+    ((0xffffffff,)*4,(0xffffffff,)*2,(0x408f276d,0x41c83b0e,0xa20bc7c6,0x6d5451fd)),
+    ((0x243f6a88,0x85a308d3,0x13198a2e,0x03707344),(0xa4093822,0x299f31d0),(0xd16cfe09,0x94fdcceb,0x5001e420,0x24126ea1))]:
+    out = philox4x32_10(*map(t, ctr), *map(t, key))
+    got = tuple(int(o) for o in out)
+    print([hex(g) for g in got], got == exp)
+from deepinfer.model_executor.layers.tilelang_ops import ops
+print(ops.random_sampling_tl, ops.fused_sample_tl, ops.map_lookup_tl)
diff --git a/tests/neuron/setup_trn2.sh b/tests/neuron/setup_trn2.sh
new file mode 100755
index 0000000..a59c550
--- /dev/null
+++ b/tests/neuron/setup_trn2.sh
@@ -0,0 +1,69 @@
+#!/bin/bash
+# One-time setup of a fresh Trn2 (Neuron DLAMI, SDK 2.33+) for the DeepInfer Neuron port.
+#
+#   git clone -b neuron https://github.com/whn09/DeepInfer.git && cd DeepInfer
+#   bash tests/neuron/setup_trn2.sh            # WORK=~/di by default
+#
+# Creates $WORK/venv (layered on the DLAMI's read-only torch-neuronx venv), builds the C++ extension,
+# downloads the official reference implementation + tokenizer, generates the tiny checkpoint and the
+# CPU golden. Idempotent: finished steps are skipped. Host CPU only -- never touches a NeuronCore.
+set -euo pipefail
+
+REPO="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
+WORK="${WORK:-$HOME/di}"
+OPT_VENV="${OPT_VENV:-/opt/aws_neuronx_venv_torch_neuronx_2_14_0}"
+HF_REPO="deepseek-ai/DeepSeek-V4.1-Flash-1017"
+mkdir -p "$WORK"
+log() { echo "[setup $(date +%H:%M:%S)] $*"; }
+
+# 1. venv layered on the read-only DLAMI torch-neuronx venv
+if [ ! -x "$WORK/venv/bin/python" ]; then
+  log "venv -> $WORK/venv (layered on $OPT_VENV)"
+  PY="$(readlink -f "$OPT_VENV/bin/python")"
+  "$PY" -m venv "$WORK/venv"
+  SP="$("$WORK/venv/bin/python" -c 'import site; print(site.getsitepackages()[0])')"
+  echo "import site; site.addsitedir('$OPT_VENV/lib/python$("$PY" -c 'import sys; print(f"{sys.version_info[0]}.{sys.version_info[1]}")')/site-packages')" > "$SP/zz_opt_neuron.pth"
+  # neuronx-cc and the neuron tools live in the DLAMI venv
+  echo "export PATH=\"\$PATH:$OPT_VENV/bin\"" >> "$WORK/venv/bin/activate"
+fi
+# shellcheck disable=SC1091
+source "$WORK/venv/bin/activate"
+python -c "import torch, torch_neuronx; print('torch', torch.__version__, 'torch-neuronx', torch_neuronx.__version__)"
+which neuronx-cc >/dev/null || { echo "neuronx-cc not on PATH"; exit 1; }
+
+# 2. Python deps: requirements.txt minus torch / tilelang / tile-kernels (pure-torch shims are vendored)
+log "pip deps"
+pip install -q einops "huggingface_hub<1" ninja orjson pillow posix_ipc pybind11 pyzmq requests scipy sympy tabulate \
+  tokenizers tqdm "transformers==4.50.0" xgrammar pytest psutil safetensors
+
+# 3. C++ extension (host-only, -DWITH_NEURON)
+if ! ls "$REPO"/deepinfer/_so/neuron-*/_C*.so >/dev/null 2>&1; then
+  log "build C++ extension"
+  (cd "$REPO" && DEEPINFER_ACCELERATOR=neuron MAX_JOBS=16 bash ./scripts/build.sh build > "$WORK/build.log" 2>&1) \
+    || { tail -30 "$WORK/build.log"; exit 1; }
+fi
+
+# 4. official reference implementation + tokenizer + real config (small files only, no weights)
+REF="$WORK/ref_dsv41"
+if [ ! -f "$REF/inference/model.py" ]; then
+  log "download reference implementation from $HF_REPO"
+  mkdir -p "$REF/inference"
+  for f in inference/model.py inference/kernel.py inference/convert.py inference/engram.py inference/generate.py \
+           inference/vision.py inference/image_processor.py inference/config.json tokenizer.json tokenizer_config.json config.json; do
+    curl -fsSL "https://huggingface.co/$HF_REPO/resolve/main/$f" -o "$REF/$f"
+  done
+fi
+
+# 5. tiny checkpoint (~1.6 GB) and the CPU golden from the official reference
+if [ ! -f "$WORK/tiny_ckpt/config.json" ]; then
+  log "tiny checkpoint -> $WORK/tiny_ckpt"
+  TORCH_DEVICE_BACKEND_AUTOLOAD=0 python "$REPO/examples/neuron/make_tiny_checkpoint.py" \
+    --out "$WORK/tiny_ckpt" --tokenizer-dir "$REF" --real-config "$REF/config.json"
+fi
+if [ ! -f "$WORK/golden/long.pt" ]; then
+  log "golden (official reference on CPU) -> $WORK/golden"
+  TORCH_DEVICE_BACKEND_AUTOLOAD=0 python -I "$REPO/examples/neuron/reference_golden.py" \
+    --hf-ckpt "$WORK/tiny_ckpt" --ref-dir "$REF/inference" --converted "$WORK/tiny_ckpt_ref_mp1" \
+    --out "$WORK/golden" --threads 32
+fi
+log "done. Next: bash tests/neuron/run_all.sh"
diff --git a/tests/neuron/tile_kernels/integ.py b/tests/neuron/tile_kernels/integ.py
new file mode 100644
index 0000000..050f494
--- /dev/null
+++ b/tests/neuron/tile_kernels/integ.py
@@ -0,0 +1,20 @@
+import torch
+import deepinfer  # noqa: sets vendor path
+from deepinfer.model_executor.models.v4_modules import mhc_ops
+import tile_kernels.torch as ref
+torch.manual_seed(0)
+t, h, M = 6, 256, 4
+m3 = M * (M + 2)
+res = torch.randn(1, t, M, h, dtype=torch.bfloat16)
+fn = torch.randn(m3, M * h) * 0.02
+w = torch.rand(M * h) + 0.5
+print('n_splits', mhc_ops.default_n_splits())
+for ns in (1, 16):
+    try:
+        out = mhc_ops.mhc_pre_norm_fn(res, fn, w, 1e-6, n_splits=ns)
+        r = ref.mhc_pre_norm_fn_ref(res, fn, w, 1e-6)
+        print(ns, 'max diff', (out - r).abs().max().item())
+    except Exception as e:
+        print(ns, 'ERR', type(e).__name__, str(e)[:300])
+x = torch.randn(t, 2, 128, dtype=torch.bfloat16)
+o = mhc_ops.expand_to_mhc(x, 4); print('expand ok', o.shape)
diff --git a/tests/neuron/tile_kernels/plugin/conftest.py b/tests/neuron/tile_kernels/plugin/conftest.py
new file mode 100644
index 0000000..33dac33
--- /dev/null
+++ b/tests/neuron/tile_kernels/plugin/conftest.py
@@ -0,0 +1,15 @@
+import os, sys
+sys.path.insert(0, os.path.join(os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..", "..")), "deepinfer", "neuron_ops", "vendor"))
+_D = os.path.dirname(os.path.abspath(__file__))
+os.environ.setdefault("TK_TEST_MERGE_MODE", "overwrite")
+os.environ.setdefault("TK_BENCHMARK_TESTS_DIR", _D)
+os.environ.setdefault("TK_BENCHMARK_BASELINES_DIR", os.path.join(_D, "bb"))
+os.environ.setdefault("TK_BENCHMARK_FAIL_MODE", "none")
+os.environ.setdefault("TK_GPU_MEM_PROFILES_DIR", os.path.join(_D, "gm"))
+pytest_plugins = [
+    "tile_kernels.testing.pytest.benchmark",
+    "tile_kernels.testing.pytest.gpu_mem",
+    "tile_kernels.testing.pytest.random",
+    "tile_kernels.testing.pytest.tilelang_warning",
+    "tile_kernels.testing.pytest.xdist_failfast",
+]
diff --git a/tests/neuron/tile_kernels/plugin/test_plugin.py b/tests/neuron/tile_kernels/plugin/test_plugin.py
new file mode 100644
index 0000000..fe2d56f
--- /dev/null
+++ b/tests/neuron/tile_kernels/plugin/test_plugin.py
@@ -0,0 +1,6 @@
+import torch
+
+
+def test_timer(benchmark_timer, seed):
+    us = benchmark_timer(lambda: torch.ones(8).sum())
+    assert us > 0 and isinstance(seed, int)
diff --git a/tests/neuron/tile_kernels/test_tile_kernels.py b/tests/neuron/tile_kernels/test_tile_kernels.py
new file mode 100644
index 0000000..78c99e0
--- /dev/null
+++ b/tests/neuron/tile_kernels/test_tile_kernels.py
@@ -0,0 +1,557 @@
+"""CPU correctness tests: Neuron pure-torch tile_kernels vs the upstream tile_kernels.torch references."""
+
+import os
+import sys
+
+VENDOR = os.path.join(os.path.normpath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "..", "..")), "deepinfer", "neuron_ops", "vendor")
+sys.path.insert(0, VENDOR)
+
+import pytest
+import torch
+
+import tile_kernels
+from tile_kernels import quant
+from tile_kernels import torch as ref
+from tile_kernels.quant.common import get_packed_ue8m0_torch_dtype
+
+torch.set_num_threads(4)
+
+
+def bits(t):
+    if t.dtype in (torch.float8_e4m3fn,):
+        return t.view(torch.uint8)
+    return t
+
+
+def assert_same(a, b, check_stride=True):
+    assert a.dtype == b.dtype, (a.dtype, b.dtype)
+    assert a.shape == b.shape, (a.shape, b.shape)
+    if check_stride and a.numel() > 1:
+        assert a.stride() == b.stride(), (a.stride(), b.stride())
+    if a.dtype.is_floating_point and a.dtype != torch.float8_e4m3fn:
+        assert torch.equal(torch.nan_to_num(a), torch.nan_to_num(b)), (a - b).abs().max()
+    else:
+        assert torch.equal(bits(a), bits(b))
+
+
+# ---------------------------------------------------------------------------
+# config / surface
+# ---------------------------------------------------------------------------
+
+
+def test_surface():
+    assert tile_kernels.config.is_ascend() is False
+    assert tile_kernels.get_num_vec_cores() == 2 * tile_kernels.get_num_sms()
+    assert tile_kernels.get_max_ub_per_vector_core() > 0 and tile_kernels.config.get_max_ub_per_vector_core() > 0
+    prev = tile_kernels.get_num_sms()
+    tile_kernels.set_num_sms(3)
+    assert tile_kernels.get_num_sms() == 3 and tile_kernels.get_num_vec_cores() == 6
+    tile_kernels.set_num_sms(prev)
+    tile_kernels.set_pdl(True)
+    tile_kernels.set_pdl(False)
+    tile_kernels.use_deterministic_algorithms(True)
+    tile_kernels.set_token_alignment(64)
+    assert tile_kernels.get_token_alignment() == 64
+    tile_kernels.set_token_alignment(128)
+    assert tile_kernels.utils.align(5, 4) == 8
+    assert tile_kernels.utils.dtype_to_str(torch.bfloat16) == 'bf16'
+    assert get_packed_ue8m0_torch_dtype() == torch.int16
+    from tile_kernels.config import is_ascend  # noqa
+    from tile_kernels.mhc.expand_kernel import expand_to_mhc_fwd  # noqa
+    from tile_kernels.mhc.head_compute_mix_kernel import mhc_head_compute_mix_fwd  # noqa
+    from tile_kernels.mhc.norm_fn_kernel import mhc_fn_normw_merge_fwd, mhc_reduce_partials_and_rmsnorm_fwd  # noqa
+    from tile_kernels.mhc.post_kernel import mhc_post_fwd  # noqa
+    from tile_kernels.mhc.pre_apply_mix_kernel import mhc_pre_apply_mix_fwd  # noqa
+    from tile_kernels.mhc.pre_split_mixes_kernel import mhc_pre_split_mixes_fwd  # noqa
+    from tile_kernels.mhc.sinkhorn_kernel import mhc_sinkhorn_fwd  # noqa
+    from tile_kernels.mhc.pre_big_fuse_kernel import mhc_pre_big_fuse_fwd  # noqa
+    from tile_kernels.moe import moe_topk_gate  # noqa
+    from tile_kernels.quant import swiglu_forward_and_per_token_cast, swiglu_forward, per_token_cast, batched_transpose_weight_sf  # noqa
+    from tile_kernels.transform import apply_rotary  # noqa
+    from tile_kernels.torch import cast, cast_back  # noqa
+    from tile_kernels import get_max_ub_per_vector_core, get_num_vec_cores  # noqa
+    for name in ('engram_hash', 'engram_gate_fwd', 'fused_weight'):
+        assert callable(getattr(tile_kernels.engram, name))
+    assert callable(tile_kernels.torch.engram.make_offsets)
+    for name in ('per_token_cast', 'per_token_cast_back', 'per_token_cast_and_cast_back'):
+        assert callable(getattr(quant, name))
+    import importlib
+
+    for mod in ('benchmark', 'gpu_mem', 'random', 'tilelang_warning', 'xdist_failfast'):
+        importlib.import_module(f'tile_kernels.testing.pytest.{mod}')
+
+
+# ---------------------------------------------------------------------------
+# per_token_cast
+# ---------------------------------------------------------------------------
+
+LAYOUTS = [
+    # round_sf, packed, col_major, e4m3_sf
+    (False, False, False, False),
+    (True, False, False, False),
+    (True, True, False, False),
+    (True, True, True, False),
+    (True, False, True, False),
+    (False, False, True, False),
+    (False, False, False, True),
+]
+
+
+@pytest.mark.parametrize('layout', LAYOUTS)
+@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
+@pytest.mark.parametrize('npc', [16, 32, 64, 128])
+@pytest.mark.parametrize('m,hidden', [(1, 128), (7, 256), (33, 512), (64, 192)])
+@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
+def test_per_token_cast(layout, fmt, npc, m, hidden, dtype):
+    round_sf, packed, col, e4m3_sf = layout
+    if hidden % npc:
+        pytest.skip('hidden % npc')
+    torch.manual_seed(m * 1000 + hidden + npc)
+    x = (torch.randn(m, hidden) * torch.logspace(-3, 3, hidden)).to(dtype)
+    x[0, :npc] = 0  # all-zero group exercises the amax clamp
+    kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col, use_e4m3_sf=e4m3_sf)
+    q, sf = quant.per_token_cast(x, fmt, npc, **kw)
+    rq, rsf = ref.cast(x, fmt, (1, npc), **kw)
+    assert_same(q, rq)
+    assert_same(sf, rsf)
+    # sf_only / precomputed sf variants
+    sf2 = quant.per_token_cast_with_sf_only(x, fmt, npc, use_tma_aligned_col_major_sf=col, round_sf=round_sf, use_packed_ue8m0=packed, use_e4m3_sf=e4m3_sf)
+    assert_same(sf2, rsf)
+    if not packed and not col:
+        q3 = quant.per_token_cast_with_precomputed_sf(x, fmt, npc, rsf)
+        rq3 = ref.cast(x, fmt, (1, npc), sf=rsf.float())
+        assert_same(q3, rq3)
+    # out= in place
+    if not col:
+        out_x = torch.empty_like(q)
+        out_sf = torch.empty_like(sf)
+        rx, rs = quant.per_token_cast(x, fmt, npc, out=(out_x, out_sf), **kw)
+        assert rx.data_ptr() == out_x.data_ptr() and rs.data_ptr() == out_sf.data_ptr()
+        assert_same(out_x, rq)
+        assert_same(out_sf, rsf)
+
+
+def test_per_token_cast_ascend_layout_bytes():
+    """Packed col-major sf: byte (k % 2) of storage element [k // 2, m] is exponent of group k (Ascend get_sf_index)."""
+    torch.manual_seed(0)
+    m, hidden, npc = 5, 256, 32
+    x = torch.randn(m, hidden, dtype=torch.bfloat16) * 10
+    _, sf = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
+    assert sf.dtype == torch.int16 and sf.shape == (m, hidden // npc // 2) and sf.stride() == (1, m)
+    storage = sf.T.contiguous().view(torch.uint8)  # (P, 2m) bytes
+    _, sf_rm = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=False)
+    exp = (sf_rm.view(torch.int32) >> 23)
+    for mi in range(m):
+        for k in range(hidden // npc):
+            assert storage[k // 2, 2 * mi + k % 2].item() == exp[mi, k].item()
+    # row-major packed: byte 2p + j of row m = group 2p + j
+    _, sfp = quant.per_token_cast(x, 'e4m3', npc, round_sf=True, use_packed_ue8m0=True)
+    assert sfp.shape == (m, hidden // npc // 2) and sfp.is_contiguous()
+    assert torch.equal(sfp.view(torch.uint8).to(torch.int32), exp)
+
+
+@pytest.mark.parametrize('in_layout', ['fp32_128', 'packed_col_32', 'packed_row_32', 'fp32_1x32'])
+@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
+def test_per_token_cast_with_input_sf(in_layout, fmt):
+    torch.manual_seed(1)
+    m, hidden = 256, 256
+    w = torch.randn(m, hidden) * 3
+    if in_layout == 'fp32_128':
+        xq, xsf = ref.cast(w, 'e4m3', (128, 128), round_sf=True)
+        blk = (128, 128)
+    elif in_layout == 'fp32_1x32':
+        xq, xsf = ref.cast(w, 'e4m3', (1, 32))
+        blk = (1, 32)
+    elif in_layout == 'packed_col_32':
+        xq, xsf = ref.cast(w, 'e4m3', (1, 32), round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
+        blk = (1, 32)
+    else:
+        xq, xsf = ref.cast(w, 'e4m3', (1, 32), round_sf=True, use_packed_ue8m0=True)
+        blk = (1, 32)
+    kw = dict(round_sf=True, use_tma_aligned_col_major_sf=True, use_packed_ue8m0=True)
+    q, sf = quant.per_token_cast((xq, xsf), fmt, 32, x_block_size=blk, **kw)
+    rq, rsf = ref.cast((xq, xsf), fmt, (1, 32), x_block_size=blk, **kw)
+    assert_same(q, rq)
+    assert_same(sf, rsf)
+
+
+def test_per_token_cast_empty():
+    x = torch.empty(0, 128, dtype=torch.bfloat16)
+    q, sf = quant.per_token_cast(x, 'e4m3', 32, round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
+    assert q.shape == (0, 128) and sf.shape == (0, 2) and sf.dtype == torch.int16
+
+
+# ---------------------------------------------------------------------------
+# cast_back / cast_and_cast_back / per_block_cast
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize('layout', LAYOUTS)
+@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
+@pytest.mark.parametrize('npc', [16, 32, 128])
+@pytest.mark.parametrize('out_fmt', ['bf16', 'fp32'])
+def test_per_token_cast_back(layout, fmt, npc, out_fmt):
+    round_sf, packed, col, e4m3_sf = layout
+    torch.manual_seed(npc)
+    x = torch.randn(19, 256, dtype=torch.bfloat16) * 5
+    kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col, use_e4m3_sf=e4m3_sf)
+    q, sf = ref.cast(x, fmt, (1, npc), **kw)
+    out = quant.per_token_cast_back((q, sf), out_fmt, npc)
+    rout = ref.cast_back((q, sf), out_fmt, (1, npc))
+    assert_same(out, rout)
+
+
+@pytest.mark.parametrize('fmt', ['e4m3', 'e2m1'])
+@pytest.mark.parametrize('npc', [16, 32, 64, 128])
+@pytest.mark.parametrize('sf_mode', ['plain', 'round', 'e4m3'])
+@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
+def test_per_token_cast_and_cast_back(fmt, npc, sf_mode, dtype):
+    torch.manual_seed(npc)
+    base = torch.randn(23, 384, dtype=dtype) * 7
+    big = torch.zeros(23, 512, dtype=dtype)
+    big[:, :384] = base
+    x = big[:, :384]  # strided rows
+    round_sf, e4m3_sf = sf_mode == 'round', sf_mode == 'e4m3'
+    quant.per_token_cast_and_cast_back(x, fmt, npc, round_sf, e4m3_sf)
+    q, sf = ref.cast(base, fmt, (1, npc), round_sf=round_sf, use_e4m3_sf=e4m3_sf)
+    expected = ref.cast_back((q, sf), 'fp32', (1, npc)).to(dtype)
+    assert torch.equal(x, expected)
+    assert torch.equal(big[:, 384:], torch.zeros(23, 128, dtype=dtype))
+
+
+@pytest.mark.parametrize('block', [(128, 128), (32, 32), (1, 128)])
+@pytest.mark.parametrize('layout', LAYOUTS[:6])
+def test_per_block_cast(block, layout):
+    round_sf, packed, col, _ = layout
+    torch.manual_seed(2)
+    x = torch.randn(256, 384, dtype=torch.bfloat16)
+    kw = dict(round_sf=round_sf, use_packed_ue8m0=packed, use_tma_aligned_col_major_sf=col)
+    q, sf = quant.per_block_cast(x, 'e4m3', block, **kw)
+    rq, rsf = ref.cast(x, 'e4m3', block, **kw)
+    assert_same(q, rq)
+    assert_same(sf, rsf)
+
+
+# ---------------------------------------------------------------------------
+# swiglu
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
+@pytest.mark.parametrize('with_weights', [False, True])
+@pytest.mark.parametrize('clamp', [None, 3.0])
+@pytest.mark.parametrize('with_psum', [False, True])
+def test_swiglu_forward(dtype, with_weights, clamp, with_psum):
+    torch.manual_seed(3)
+    tile_kernels.set_token_alignment(8)
+    n, hidden = 64, 256
+    x = (torch.randn(n, hidden * 2) * 3).to(dtype)
+    w = torch.rand(n) if with_weights else None
+    rsf = 2.5 if with_weights else None
+    psum = torch.tensor([5, 8, 8, 20, 37, 44], dtype=torch.int32) if with_psum else None  # aligned-prefix layout
+    for fmt in ('fp32', 'bf16'):
+        cnt = torch.zeros(4, dtype=torch.int64) if clamp is not None else None
+        rcnt = torch.zeros(4, dtype=torch.int64) if clamp is not None else None
+        out = quant.swiglu_forward(x, fmt, psum, w, rsf, clamp, cnt)
+        rout = ref.swiglu_forward(x, psum, w, rsf, clamp, rcnt)
+        assert torch.equal(out, rout.to(out.dtype))
+        if cnt is not None:
+            assert torch.equal(cnt, rcnt), (cnt, rcnt)
+    # fused fp8 cast (deepinfer get_fp8_cast_params)
+    kw = dict(round_sf=True, use_packed_ue8m0=True, use_tma_aligned_col_major_sf=True)
+    for npc in (32, 128):
+        q, sf = quant.swiglu_forward_and_per_token_cast(x, 'e4m3', npc, psum, w, rsf, clamp_value=clamp, **kw)
+        rq, rsf_ = ref.cast(ref.swiglu_forward(x, psum, w, rsf, clamp), 'e4m3', (1, npc), **kw)
+        assert_same(q, rq)
+        assert_same(sf, rsf_)
+    tile_kernels.set_token_alignment(128)
+
+
+# ---------------------------------------------------------------------------
+# batched_transpose_weight_sf
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize('block', [32, 128])
+def test_batched_transpose_weight_sf(block):
+    torch.manual_seed(4)
+
+    def make(*shape):
+        e = torch.randint(110, 140, shape, dtype=torch.int32)
+        return (e << 23).view(torch.float32)
+
+    sfs = [make(384 // block, 640 // block), make(3, 256 // block, 896 // block), make(128 // block, 128 // block)]
+    flags = [False, True, True]
+    outs = quant.batched_transpose_weight_sf(sfs, block, flags, outs=[None] * 3)
+    refs = ref.batched_transpose_weight_sf(sfs, block, flags)
+    for (a0, a1), (r0, r1) in zip(outs, refs):
+        assert_same(a0, r0)
+        assert_same(a1, r1)
+    # in-place rewrite with new values
+    sfs2 = [make(*s.shape) for s in sfs]
+    outs2 = quant.batched_transpose_weight_sf(sfs2, block, flags, outs=outs)
+    refs2 = ref.batched_transpose_weight_sf(sfs2, block, flags)
+    for (a0, a1), (o0, o1), (r0, r1) in zip(outs2, outs, refs2):
+        assert a0.data_ptr() == o0.data_ptr() and a1.data_ptr() == o1.data_ptr()
+        assert_same(a0, r0)
+        assert_same(a1, r1)
+
+
+# ---------------------------------------------------------------------------
+# rope
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize('rot_dim', [64, 128])
+@pytest.mark.parametrize('ndim', [3, 4])
+@pytest.mark.parametrize('interleaved', [False, True])
+@pytest.mark.parametrize('conjugate', [False, True])
+@pytest.mark.parametrize('use_positions', [False, True])
+@pytest.mark.parametrize('dtype', [torch.bfloat16, torch.float32])
+def test_apply_rotary(rot_dim, ndim, interleaved, conjugate, use_positions, dtype):
+    torch.manual_seed(5)
+    seqlen_ro, s, h, hk, b = 64, 9, 4, 2, 2
+    cache = torch.randn(seqlen_ro, rot_dim)
+    cos, sin = cache[:, : rot_dim // 2], cache[:, rot_dim // 2 :]
+    lead = (s,) if ndim == 3 else (b, s)
+    full_q = torch.randn(*lead, h, rot_dim + 32).to(dtype)
+    q = full_q[..., 32:]  # tail partial rope, strided
+    k = torch.randn(*lead, hk, rot_dim).to(dtype)
+    pos = torch.randint(0, seqlen_ro - 4, lead, dtype=torch.int32) if use_positions else None
+    off = 3
+    rq = ref.rotary_embedding_ref(q.float(), cos, sin, None if pos is None else pos.long(), off, interleaved, conjugate)
+    rk = ref.rotary_embedding_ref(k.float(), cos, sin, None if pos is None else pos.long(), off, interleaved, conjugate)
+    head = full_q[..., :32].clone()
+    tile_kernels.transform.apply_rotary(q, cache, k, pos, interleaved, conjugate, off)
+    tol = dict(atol=0, rtol=0) if dtype == torch.float32 else dict(atol=1e-2, rtol=1e-2)
+    torch.testing.assert_close(q.float(), rq.to(dtype).float(), **(dict(atol=1e-5, rtol=1e-5) if dtype == torch.float32 else tol))
+    torch.testing.assert_close(k.float(), rk.to(dtype).float(), **(dict(atol=1e-5, rtol=1e-5) if dtype == torch.float32 else tol))
+    assert torch.equal(full_q[..., :32], head)
+    # bitwise vs the Ascend op order: out1 = fma(first, cos, second * -sin), out2 = fma(first, sin, second * cos)
+    def fma(a, b, c):
+        return (a.double() * b.double() + c.double()).float()
+
+    q2 = full_q[..., 32:].clone()
+    q0 = torch.randn_like(q2.float()).to(dtype)
+    q2.copy_(q0)
+    tile_kernels.transform.apply_rotary(q2, cache, None, pos, interleaved, conjugate, off)
+    idx = (pos.long() + off) if pos is not None else (torch.arange(s) + off).expand(*lead[:-1], s)
+    c_, s_ = cos[idx].unsqueeze(-2), sin[idx].unsqueeze(-2)
+    s_ = -s_ if conjugate else s_
+    xf = q0.float()
+    f, sc = (xf[..., 0::2], xf[..., 1::2]) if interleaved else (xf[..., : rot_dim // 2], xf[..., rot_dim // 2 :])
+    c_, s_ = c_.expand_as(f), s_.expand_as(f)
+    o1, o2 = fma(f, c_, sc * -s_), fma(f, s_, sc * c_)
+    exp = torch.stack([o1, o2], -1).flatten(-2) if interleaved else torch.cat([o1, o2], -1)
+    assert torch.equal(q2, exp.to(dtype))
+
+
+# ---------------------------------------------------------------------------
+# moe
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize('num_shared', [0, 1, 2])
+@pytest.mark.parametrize('with_map', [False, True])
+@pytest.mark.parametrize('with_mask', [False, True])
+@pytest.mark.parametrize('with_image', [False, True])
+@pytest.mark.parametrize('with_fix', [False, True])
+def test_moe_topk_gate(num_shared, with_map, with_mask, with_image, with_fix):
+    torch.manual_seed(6)
+    t, e, k = 37, 64, 6
+    logits = torch.randn(t, e) * 2
+    logits[3, 10] = logits[3, 11]  # tie
+    bias = torch.randn(e) * 0.1
+    kwargs = dict(num_topk=k, use_shared_as_routed=num_shared > 0, num_shared_experts=num_shared, routed_scaling_factor=2.5, ep_rank=3, scoring_func='sqrtsoftplus', bias=bias)
+    if with_mask:
+        kwargs['mask'] = torch.rand(t) > 0.3
+    if with_image:
+        kwargs['image_bias'] = torch.randn(e) * 0.1
+        kwargs['image_token_mask'] = torch.rand(t) > 0.5
+    if with_map:
+        nl = e + num_shared
+        dup = 3
+        count = torch.randint(1, dup + 1, (nl,), dtype=torch.int32)
+        kwargs['to_physical_map'] = torch.randint(0, 500, (nl, dup), dtype=torch.int32)
+        kwargs['logical_count'] = count
+    unm = unm_ref = None
+    if with_fix:
+        unm = torch.stack([torch.randperm(e)[:k] for _ in range(t)]).to(torch.int64)
+        unm_ref = unm.clone()
+        kwargs['fix_routing_mask'] = torch.rand(t) > 0.5
+    idx, w = tile_kernels.moe.moe_topk_gate(logits, unmapped_topk_idx=unm, **kwargs)
+    ridx, rw = ref.moe_topk_gate_forward(logits, unmapped_topk_idx=unm_ref, **kwargs)
+    assert torch.equal(idx, ridx)
+    torch.testing.assert_close(w, rw, atol=1e-6, rtol=1e-6)
+    if with_fix:
+        assert torch.equal(unm, unm_ref)
+
+
+def test_moe_topk_gate_out_scores_force_random():
+    torch.manual_seed(7)
+    t, e, k = 16, 32, 4
+    logits = torch.randn(t, e)
+    mask = torch.ones(t, dtype=torch.bool)
+    mask[0] = False
+    fr = torch.zeros(t, dtype=torch.bool)
+    fr[1:4] = True
+    out = (torch.empty(t, k, dtype=torch.int64), torch.empty(t, k))
+    scores = torch.empty(t, e)
+    unm = torch.empty(t, k, dtype=torch.int64)
+    idx, w = tile_kernels.moe.moe_topk_gate(logits, k, False, 0, 1.0, 0, mask=mask, force_random=fr, out=out, scores=scores, unmapped_topk_idx=unm)
+    assert idx.data_ptr() == out[0].data_ptr()
+    assert (idx[0] == -1).all() and (w[0] == 0).all() and (scores[0] == 0).all() and (unm[0] == -1).all()
+    assert (unm[1:4] == -1).all() and (scores[1:4] == 0).all()
+    for r in range(1, 4):
+        assert len(set(idx[r].tolist())) == k and idx[r].min() >= 0 and idx[r].max() < e
+    ridx, rw = ref.moe_topk_gate_forward(logits[4:], k, False, 0, 1.0, 0, 'sqrtsoftplus')
+    assert torch.equal(idx[4:], ridx)
+    torch.testing.assert_close(scores[4:], torch.nn.functional.softplus(logits[4:]).sqrt())
+
+
+# ---------------------------------------------------------------------------
+# engram
+# ---------------------------------------------------------------------------
+
+
+def test_engram_hash():
+    torch.manual_seed(8)
+    t, n, layers, tables = 29, 4, 2, 3
+    ids = torch.randint(0, 100000, (t, n), dtype=torch.int32)
+    mult = torch.randint(1, 2**40, (layers, n), dtype=torch.int64)
+    vs = torch.randint(1000, 50000, (layers, n - 1, tables), dtype=torch.int32)
+    off = tile_kernels.torch.engram.make_offsets(vs)
+    mask = torch.rand(t) > 0.7
+    for m in (None, mask):
+        out = tile_kernels.engram.engram_hash(ids, mult, vs, off, m)
+        rout = ref.engram_hash_ref(ids, mult, vs, off, m)
+        assert out.dtype == torch.int32 and torch.equal(out, rout)
+
+
+@pytest.mark.parametrize('inplace', [False, True])
+@pytest.mark.parametrize('with_mask', [False, True])
+def test_engram_gate(inplace, with_mask):
+    torch.manual_seed(9)
+    t, hc, d = 13, 4, 256
+    hs = torch.randn(t, hc, d, dtype=torch.bfloat16)
+    kv = torch.randn(t, hc + 1, d, dtype=torch.bfloat16)
+    wh = torch.randn(hc, d, dtype=torch.bfloat16)
+    we = torch.randn(hc, d, dtype=torch.bfloat16)
+    mask = (torch.rand(t) > 0.5) if with_mask else None
+    wf = tile_kernels.engram.fused_weight(wh, we)
+    assert wf.dtype == torch.float32 and torch.equal(wf, wh.float() * we.float())
+    r_out, r_dot, r_g, r_rx, r_rk = ref.engram_gate_ref(hs, kv, wh, we, 1e-4, 1e-6, save_for_backward=True, image_token_mask=mask)
+    hs_in = hs.clone()
+    out, dot, g, rx, rk = tile_kernels.engram.engram_gate_fwd(hs_in, kv, wf, 1e-6, 1e-4, save_for_backward=True, image_token_mask=mask, inplace=inplace)
+    if inplace:
+        assert out.data_ptr() == hs_in.data_ptr()
+    sel = ~mask if mask is not None else slice(None)
+    assert (out.float() - r_out.float()).abs().max() <= 0.0625
+    assert torch.equal(out[mask], hs[mask]) if mask is not None else True
+    torch.testing.assert_close(dot[sel], r_dot[sel], atol=1e-3, rtol=1e-4)
+    torch.testing.assert_close(g[sel], r_g[sel], atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(rx[sel], r_rx[sel], atol=1e-6, rtol=1e-5)
+    torch.testing.assert_close(rk[sel], r_rk[sel], atol=1e-6, rtol=1e-5)
+    out2, *rest = tile_kernels.engram.engram_gate_fwd(hs.clone(), kv, wf, 1e-6, 1e-4, save_for_backward=False, image_token_mask=mask)
+    assert all(r is None for r in rest)
+    assert torch.equal(out2, out)
+
+
+# ---------------------------------------------------------------------------
+# mhc
+# ---------------------------------------------------------------------------
+
+MHC = 4
+
+
+def test_mhc_expand_and_mix():
+    torch.manual_seed(10)
+    n, h = 11, 256
+    x = torch.randn(n, h, dtype=torch.bfloat16)
+    o = torch.empty(n, MHC, h, dtype=torch.bfloat16)
+    tile_kernels.mhc.expand_to_mhc_fwd(x, o)
+    assert torch.equal(o, ref.expand_to_mhc_ref(x, MHC))
+
+    inp = torch.randn(n, MHC)
+    scale, base = torch.randn(1), torch.randn(MHC)
+    out = torch.empty_like(inp)
+    tile_kernels.mhc.mhc_head_compute_mix_fwd(inp, scale, base, out, 1e-6)
+    torch.testing.assert_close(out, ref.mhc_head_compute_mix_ref(inp, scale, base, 1e-6), atol=0, rtol=0)
+
+    m3 = MHC * (MHC + 2)
+    mixes = torch.randn(1, n, m3)
+    scale3, base3 = torch.randn(3), torch.randn(m3)
+    pre, post, comb = torch.empty(n, MHC), torch.empty(n, MHC), torch.empty(n, MHC * MHC)
+    tile_kernels.mhc.mhc_pre_split_mixes_fwd(mixes.view(n, m3), scale3, base3, pre, post, comb, 2.0, 1e-6)
+    rpre, rpost, rcomb = ref.mhc_pre_split_mixes_ref(mixes, scale3, base3, MHC, 2.0, 1e-6)
+    torch.testing.assert_close(pre, rpre.view(n, MHC), atol=1e-7, rtol=1e-6)
+    torch.testing.assert_close(post, rpost.view(n, MHC), atol=1e-7, rtol=1e-6)
+    torch.testing.assert_close(comb, rcomb.reshape(n, MHC * MHC), atol=1e-7, rtol=1e-6)
+
+    xr = torch.randn(n, MHC, h, dtype=torch.bfloat16)
+    mix = torch.rand(n, MHC)
+    o2 = torch.empty(n, h, dtype=torch.bfloat16)
+    tile_kernels.mhc.mhc_pre_apply_mix_fwd(xr, mix, o2)
+    torch.testing.assert_close(o2.float(), ref.mhc_pre_apply_mix_ref(xr.float(), mix.unsqueeze(-1)).float(), atol=1e-2, rtol=1e-2)
+
+    cm = torch.randn(n, MHC, MHC)
+    cmo = torch.empty_like(cm)
+    tile_kernels.mhc.mhc_sinkhorn_fwd(cm, cmo, 10, 1e-6)
+    torch.testing.assert_close(cmo, ref.sinkhorn_normalize_ref(cm, 10, 1e-6), atol=1e-6, rtol=1e-5)
+
+
+def test_mhc_post():
+    torch.manual_seed(11)
+    a, b, h = 2, 5, 256
+    x = torch.randn(a, b, h, dtype=torch.bfloat16)
+    res = torch.randn(a, b, MHC, h, dtype=torch.bfloat16)
+    post = torch.rand(a, b, MHC, 1)
+    comb = torch.rand(a, b, MHC, MHC)
+    out = tile_kernels.mhc.mhc_post_fwd(x, res, post, comb)
+    rout = ref.mhc_post_ref(x, res, post, comb)
+    assert out.dtype == torch.bfloat16 and out.shape == res.shape
+    assert (out.float() - rout.float()).abs().max() <= 0.0625
+    out2 = torch.empty_like(res)
+    r = tile_kernels.mhc.mhc_post_fwd(x, res, post, comb, out=out2)
+    assert r.data_ptr() == out2.data_ptr() and torch.equal(out2, out)
+
+
+@pytest.mark.parametrize('n_splits', [1, 4, 16])
+def test_mhc_norm_fn(n_splits):
+    torch.manual_seed(12)
+    t, h = 9, 256
+    m3 = MHC * (MHC + 2)
+    residual = torch.randn(1, t, MHC, h, dtype=torch.bfloat16)
+    fn = torch.randn(m3, MHC * h) * 0.02
+    normw = torch.rand(MHC * h) + 0.5
+    merged = torch.empty_like(fn)
+    tile_kernels.mhc.mhc_fn_normw_merge_fwd(fn, normw, merged)
+    assert torch.equal(merged, fn * normw)
+    om_s, sq_s = ref.mhc_pre_norm_fn_partials_ref(residual, merged, n_splits)  # (S, T, 1, M3), (S, T, 1)
+    om, sq, out = torch.empty(t, 1, m3), torch.empty(t, 1), torch.empty(t, m3)
+    tile_kernels.mhc.mhc_reduce_partials_and_rmsnorm_fwd(om_s, sq_s, om, sq, out, MHC * h, 1e-6, n_splits)
+    rout = ref.mhc_pre_norm_fn_ref(residual, fn, normw, 1e-6).view(t, m3)
+    torch.testing.assert_close(out, rout, atol=1e-4, rtol=1e-4)
+    torch.testing.assert_close(om, om_s.sum(0))
+    torch.testing.assert_close(sq, sq_s.sum(0))
+
+
+@pytest.mark.parametrize('n_splits', [1, 16])
+def test_mhc_pre_big_fuse(n_splits):
+    torch.manual_seed(13)
+    t, h = 7, 256
+    m3 = MHC * (MHC + 2)
+    residual = torch.randn(t, MHC, h, dtype=torch.bfloat16)
+    fn = torch.randn(m3, MHC * h) * 0.02
+    om_s, sq_s = ref.mhc_pre_norm_fn_partials_ref(residual.unsqueeze(0), fn, n_splits)
+    gemm_mul, gemm_sq = om_s.view(n_splits, t, m3), sq_s.view(n_splits, t)
+    scale, base = torch.rand(3), torch.randn(m3) * 0.1
+    post, comb, layer_in = torch.empty(t, MHC), torch.empty(t, MHC * MHC), torch.empty(t, h, dtype=torch.bfloat16)
+    tile_kernels.mhc.mhc_pre_big_fuse_fwd(gemm_mul, gemm_sq, scale, base, residual, post, comb, layer_in, 1e-6, 1e-6, 1e-6, 2.0, 20, n_splits=n_splits)
+    mixes = ref.mhc_pre_norm_fn_ref(residual.unsqueeze(0), fn, None, 1e-6)  # (1, T, M3)
+    rpre, rpost, rcomb = ref.mhc_pre_split_mixes_ref(mixes, scale, base, MHC, 2.0, 1e-6)
+    rcomb = ref.sinkhorn_normalize_ref(rcomb, 20, 1e-6)
+    rli = ref.mhc_pre_apply_mix_ref(residual.unsqueeze(0).float(), rpre)
+    torch.testing.assert_close(post, rpost.view(t, MHC), atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(comb, rcomb.reshape(t, MHC * MHC), atol=1e-5, rtol=1e-4)
+    torch.testing.assert_close(layer_in.float(), rli.view(t, h).float(), atol=2e-2, rtol=1e-2)