From e55384ae95afd1a0dc5a0d6486e26b4b4d596b56 Mon Sep 17 00:00:00 2001
From: Brandon Music <brandon.m.music@gmail.com>
Date: Mon, 20 Jul 2026 10:43:34 -0400
Subject: [PATCH 01/11] cute/compiler: include device compute capability in the
 compile disk-cache key

---
 sparkinfer/_lib/compiler.py                   | 92 +++++++++++++------
 tests/_lib/test_compile_cache.py              | 40 ++++++++
 .../acceptance/corpus/ptx_capture.py          | 30 +++---
 .../acceptance/e2e/contract.py                |  3 +-
 .../evidence/kernel_resources.py              | 42 ++++-----
 5 files changed, 142 insertions(+), 65 deletions(-)

diff --git a/sparkinfer/_lib/compiler.py b/sparkinfer/_lib/compiler.py
index 622f0a88..7fa2a45a 100644
--- a/sparkinfer/_lib/compiler.py
+++ b/sparkinfer/_lib/compiler.py
@@ -1072,12 +1072,16 @@ def _compile_cache_payload_log_value(
 ) -> dict[str, Any]:
     if payload is None:
         return {}
-    if len(payload) == 10 and payload[0] == "sparkinfer_cute_compile_cache_v5_explicit_spec":
+    if (
+        len(payload) == 11
+        and payload[0] == "sparkinfer_cute_compile_cache_v6_explicit_spec"
+    ):
         (
             _version,
             target_key,
             _sparkinfer_fingerprint,
             toolchain_key,
+            _device_arch,
             spec_hash,
             spec_json,
             kwargs_hash,
@@ -1146,13 +1150,14 @@ def _compile_cache_payload_log_value(
                 summary["toolchain"] = toolchain_summary
         return summary
 
-    if len(payload) != 8:
+    if len(payload) != 9:
         return {}
     (
         _version,
         target_key,
         _sparkinfer_fingerprint,
         toolchain_key,
+        _device_arch,
         args_key,
         kwargs_key,
         options_key,
@@ -1183,8 +1188,8 @@ def _compile_cache_payload_log_value(
 def _is_explicit_spec_payload(payload: tuple[object, ...] | None) -> bool:
     return (
         payload is not None
-        and len(payload) == 10
-        and payload[0] == "sparkinfer_cute_compile_cache_v5_explicit_spec"
+        and len(payload) == 11
+        and payload[0] == "sparkinfer_cute_compile_cache_v6_explicit_spec"
     )
 
 
@@ -1443,6 +1448,33 @@ def _distribution_version(name: str) -> str:
         return ""
 
 
+_DEVICE_ARCH_KEY: tuple[object, ...] | None = None
+
+
+def _device_arch_key() -> tuple[object, ...]:
+    """Compute capability of the device this process will compile for.
+
+    Deliberately NOT lru_cached on failure: if CUDA is not initialized yet the
+    query is retried on the next call, so an early call can never memoize a
+    bogus "unknown" and let a cubin built for one architecture be served to
+    another from the shared on-disk cache.
+    """
+    global _DEVICE_ARCH_KEY
+    if _DEVICE_ARCH_KEY is not None:
+        return _DEVICE_ARCH_KEY
+    try:
+        import torch
+
+        if not torch.cuda.is_available():
+            return ("arch", "unavailable")
+        major, minor = torch.cuda.get_device_capability()
+        name = torch.cuda.get_device_name()
+    except Exception:
+        return ("arch", "unavailable")
+    _DEVICE_ARCH_KEY = ("arch", int(major), int(minor), str(name))
+    return _DEVICE_ARCH_KEY
+
+
 @lru_cache(maxsize=1)
 def _runtime_toolchain_key() -> tuple[object, ...]:
     from .runtime_patches import cutlass_runtime_patch_status
@@ -1538,6 +1570,7 @@ def _static_compile_cache_context(compile_callable: Any) -> tuple[object, ...]:
     return (
         _sparkinfer_package_fingerprint(),
         _runtime_toolchain_key(),
+        _device_arch_key(),
         _compile_options_cache_key(compile_callable),
         _compile_environment_key(),
     )
@@ -1849,16 +1882,18 @@ def _compile_disk_cache_payload(
     (
         package_fingerprint,
         runtime_toolchain,
+        device_arch,
         compile_options,
         compile_environment,
     ) = _static_compile_cache_context(compile_callable)
     if compile_spec is not None:
         kwargs_json_key, kwargs_hash_key = _compile_kwargs_json_key(kwargs)
         return (
-            "sparkinfer_cute_compile_cache_v5_explicit_spec",
+            "sparkinfer_cute_compile_cache_v6_explicit_spec",
             _explicit_spec_compile_target(func),
             package_fingerprint,
             runtime_toolchain,
+            device_arch,
             compile_spec.hash_key,
             compile_spec.json_key,
             kwargs_hash_key,
@@ -1867,10 +1902,11 @@ def _compile_disk_cache_payload(
             compile_environment,
         )
     return (
-        "sparkinfer_cute_compile_cache_v2",
+        "sparkinfer_cute_compile_cache_v3",
         _normalize_compile_target(func, set()),
         package_fingerprint,
         runtime_toolchain,
+        device_arch,
         _structural_cache_key(args),
         _structural_cache_key(kwargs),
         compile_options,
@@ -2008,25 +2044,25 @@ def _semantic_compile_manifest_payload(
         "cache_format": cache_format,
         "target": _semantic_target_key(cache_payload[1]),
     }
-    if cache_format == "sparkinfer_cute_compile_cache_v5_explicit_spec":
-        semantic["compile_spec_hash"] = cache_payload[4]
+    if cache_format == "sparkinfer_cute_compile_cache_v6_explicit_spec":
+        semantic["compile_spec_hash"] = cache_payload[5]
         try:
-            semantic["compile_spec"] = json.loads(str(cache_payload[5]))
+            semantic["compile_spec"] = json.loads(str(cache_payload[6]))
         except (TypeError, ValueError, json.JSONDecodeError):
-            semantic["compile_spec"] = str(cache_payload[5])
-        if cache_payload[6]:
-            semantic["compile_kwargs_hash"] = cache_payload[6]
+            semantic["compile_spec"] = str(cache_payload[6])
+        if cache_payload[7]:
+            semantic["compile_kwargs_hash"] = cache_payload[7]
             try:
-                semantic["compile_kwargs"] = json.loads(str(cache_payload[7]))
+                semantic["compile_kwargs"] = json.loads(str(cache_payload[8]))
             except (TypeError, ValueError, json.JSONDecodeError):
-                semantic["compile_kwargs"] = str(cache_payload[7])
-        semantic["compile_options"] = _manifest_json_value(cache_payload[8])
-        semantic["compile_environment"] = _manifest_json_value(cache_payload[9])
+                semantic["compile_kwargs"] = str(cache_payload[8])
+        semantic["compile_options"] = _manifest_json_value(cache_payload[9])
+        semantic["compile_environment"] = _manifest_json_value(cache_payload[10])
     else:
-        semantic["args"] = _semantic_structural_key(cache_payload[4])
-        semantic["kwargs"] = _semantic_structural_key(cache_payload[5])
-        semantic["compile_options"] = _manifest_json_value(cache_payload[6])
-        semantic["compile_environment"] = _manifest_json_value(cache_payload[7])
+        semantic["args"] = _semantic_structural_key(cache_payload[5])
+        semantic["kwargs"] = _semantic_structural_key(cache_payload[6])
+        semantic["compile_options"] = _manifest_json_value(cache_payload[7])
+        semantic["compile_environment"] = _manifest_json_value(cache_payload[8])
     return semantic
 
 
@@ -2260,9 +2296,9 @@ def _build_compile_manifest(
         allow_nan=False,
     )
     cache_format = str(cache_payload[0]) if cache_payload else "unknown"
-    explicit = cache_format == "sparkinfer_cute_compile_cache_v5_explicit_spec"
-    options_index = 8 if explicit else 6
-    environment_index = 9 if explicit else 7
+    explicit = cache_format == "sparkinfer_cute_compile_cache_v6_explicit_spec"
+    options_index = 9 if explicit else 7
+    environment_index = 10 if explicit else 8
     launch_metadata = (
         _extract_launch_dynamic_smem_bytes(compiled)
         if compiled is not None
@@ -2308,12 +2344,12 @@ def _build_compile_manifest(
         ).hexdigest(),
     }
     if explicit:
-        manifest["compile_spec_hash"] = str(cache_payload[4])
-        manifest["compile_spec_json"] = str(cache_payload[5])
-        manifest["compile_kwargs_hash"] = str(cache_payload[6])
-        manifest["compile_kwargs_json"] = str(cache_payload[7])
+        manifest["compile_spec_hash"] = str(cache_payload[5])
+        manifest["compile_spec_json"] = str(cache_payload[6])
+        manifest["compile_kwargs_hash"] = str(cache_payload[7])
+        manifest["compile_kwargs_json"] = str(cache_payload[8])
         try:
-            spec = json.loads(str(cache_payload[5]))
+            spec = json.loads(str(cache_payload[6]))
         except (TypeError, ValueError, json.JSONDecodeError):
             spec = None
         if isinstance(spec, dict):
diff --git a/tests/_lib/test_compile_cache.py b/tests/_lib/test_compile_cache.py
index 322f56ca..bb37c016 100644
--- a/tests/_lib/test_compile_cache.py
+++ b/tests/_lib/test_compile_cache.py
@@ -9,6 +9,8 @@
 import importlib
 from pathlib import Path
 
+import pytest
+
 compiler = importlib.import_module("sparkinfer._lib.compiler")
 
 
@@ -49,3 +51,41 @@ def test_cache_dir_resolution_order(monkeypatch):
 
     monkeypatch.setenv("SPARKINFER_COMPILE_CACHE_DIR", "/explicit")
     assert compiler._cute_compile_cache_dir() == Path("/explicit")
+
+
+def test_disk_cache_key_includes_device_arch(monkeypatch):
+    compile_callable = object()
+    monkeypatch.setattr(
+        compiler,
+        "_static_compile_cache_context",
+        lambda _callable: (
+            "package",
+            "toolchain",
+            ("arch", 12, 0, "gpu"),
+            (),
+            (),
+        ),
+    )
+
+    payload = compiler._compile_disk_cache_payload(
+        compile_callable,
+        test_disk_cache_key_includes_device_arch,
+        (),
+        {},
+    )
+
+    assert payload[0] == "sparkinfer_cute_compile_cache_v3"
+    assert payload[4] == ("arch", 12, 0, "gpu")
+
+
+def test_device_arch_key_retries_after_unavailable(monkeypatch):
+    torch = pytest.importorskip("torch")
+    monkeypatch.setattr(compiler, "_DEVICE_ARCH_KEY", None)
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+    assert compiler._device_arch_key() == ("arch", "unavailable")
+    assert compiler._DEVICE_ARCH_KEY is None
+
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (12, 1))
+    monkeypatch.setattr(torch.cuda, "get_device_name", lambda: "SM121")
+    assert compiler._device_arch_key() == ("arch", 12, 1, "SM121")
diff --git a/validation/cutlass_migration/acceptance/corpus/ptx_capture.py b/validation/cutlass_migration/acceptance/corpus/ptx_capture.py
index 8e967342..49ff4ae7 100644
--- a/validation/cutlass_migration/acceptance/corpus/ptx_capture.py
+++ b/validation/cutlass_migration/acceptance/corpus/ptx_capture.py
@@ -765,25 +765,25 @@ def _semantic_target_key(target_key: Any) -> Any:
 
 
 def _semantic_payload_from_cache_payload(cache_payload: list[Any]) -> dict[str, Any]:
-    if len(cache_payload) != 10:
+    if len(cache_payload) != 11:
         raise RuntimeError(f"explicit cache payload has {len(cache_payload)} fields")
-    if cache_payload[0] != "sparkinfer_cute_compile_cache_v5_explicit_spec":
+    if cache_payload[0] != "sparkinfer_cute_compile_cache_v6_explicit_spec":
         raise RuntimeError(f"unsupported cache format {cache_payload[0]!r}")
     semantic: dict[str, Any] = {
         "cache_format": cache_payload[0],
         "target": _semantic_target_key(cache_payload[1]),
-        "compile_spec_hash": cache_payload[4],
+        "compile_spec_hash": cache_payload[5],
     }
     semantic["compile_spec"] = _load_json_document(
-        str(cache_payload[5]), "cache-payload compile spec"
+        str(cache_payload[6]), "cache-payload compile spec"
     )
-    if cache_payload[6]:
-        semantic["compile_kwargs_hash"] = cache_payload[6]
+    if cache_payload[7]:
+        semantic["compile_kwargs_hash"] = cache_payload[7]
         semantic["compile_kwargs"] = _load_json_document(
-            str(cache_payload[7]), "cache-payload compile kwargs"
+            str(cache_payload[8]), "cache-payload compile kwargs"
         )
-    semantic["compile_options"] = _manifest_json_value(cache_payload[8])
-    semantic["compile_environment"] = _manifest_json_value(cache_payload[9])
+    semantic["compile_options"] = _manifest_json_value(cache_payload[9])
+    semantic["compile_environment"] = _manifest_json_value(cache_payload[10])
     return semantic
 
 
@@ -883,8 +883,8 @@ def _validate_compile_manifest(
         toolchain_names.append(entry[0])
     if len(toolchain_names) != len(set(toolchain_names)):
         raise RuntimeError("compile manifest toolchain repeats a component")
-    compile_spec_hash = cache_payload[4]
-    compile_spec_json = cache_payload[5]
+    compile_spec_hash = cache_payload[5]
+    compile_spec_json = cache_payload[6]
     if (
         not isinstance(compile_spec_hash, str)
         or not _SHA256_RE.fullmatch(compile_spec_hash)
@@ -912,8 +912,8 @@ def _validate_compile_manifest(
     if raw.get("compile_spec_version") != compile_spec.get("version"):
         raise RuntimeError("compile manifest version differs from compile spec")
 
-    compile_kwargs_hash = cache_payload[6]
-    compile_kwargs_json = cache_payload[7]
+    compile_kwargs_hash = cache_payload[7]
+    compile_kwargs_json = cache_payload[8]
     if (
         raw.get("compile_kwargs_hash") != compile_kwargs_hash
         or raw.get("compile_kwargs_json") != compile_kwargs_json
@@ -930,13 +930,13 @@ def _validate_compile_manifest(
         _load_json_document(compile_kwargs_json, "compile kwargs")
     elif compile_kwargs_hash:
         raise RuntimeError("compile manifest has a kwargs hash without JSON")
-    if raw.get("compile_options") != cache_payload[8]:
+    if raw.get("compile_options") != cache_payload[9]:
         raise RuntimeError("compile manifest compile options differ from payload")
     if not isinstance(raw["compile_options"], list) or any(
         not isinstance(option, str) or not option for option in raw["compile_options"]
     ):
         raise RuntimeError("compile manifest compile options are invalid")
-    if raw.get("compile_environment") != cache_payload[9]:
+    if raw.get("compile_environment") != cache_payload[10]:
         raise RuntimeError("compile manifest compile environment differs from payload")
     compile_environment = raw["compile_environment"]
     if not isinstance(compile_environment, list) or any(
diff --git a/validation/cutlass_migration/acceptance/e2e/contract.py b/validation/cutlass_migration/acceptance/e2e/contract.py
index a28751c3..cdebfbbf 100644
--- a/validation/cutlass_migration/acceptance/e2e/contract.py
+++ b/validation/cutlass_migration/acceptance/e2e/contract.py
@@ -345,7 +345,8 @@ def _validate_cache_artifact(
         f"{location}: unsupported compile manifest schema",
     )
     _require(
-        manifest.get("cache_format") == "sparkinfer_cute_compile_cache_v5_explicit_spec",
+        manifest.get("cache_format")
+        == "sparkinfer_cute_compile_cache_v6_explicit_spec",
         f"{location}: exact explicit compile specification is required",
     )
     cache_key = manifest.get("cache_key")
diff --git a/validation/cutlass_migration/evidence/kernel_resources.py b/validation/cutlass_migration/evidence/kernel_resources.py
index a26eddd5..78cedece 100644
--- a/validation/cutlass_migration/evidence/kernel_resources.py
+++ b/validation/cutlass_migration/evidence/kernel_resources.py
@@ -413,27 +413,27 @@ def _semantic_target_key(target_key: Any) -> Any:
 
 
 def _semantic_payload_from_cache_payload(cache_payload: list[Any]) -> dict[str, Any]:
-    if len(cache_payload) != 10:
+    if len(cache_payload) != 11:
         raise ValueError(f"explicit cache payload has {len(cache_payload)} fields")
-    if cache_payload[0] != "sparkinfer_cute_compile_cache_v5_explicit_spec":
+    if cache_payload[0] != "sparkinfer_cute_compile_cache_v6_explicit_spec":
         raise ValueError(f"unsupported cache format {cache_payload[0]!r}")
     semantic: dict[str, Any] = {
         "cache_format": cache_payload[0],
         "target": _semantic_target_key(cache_payload[1]),
-        "compile_spec_hash": cache_payload[4],
+        "compile_spec_hash": cache_payload[5],
     }
     try:
-        semantic["compile_spec"] = json.loads(str(cache_payload[5]))
+        semantic["compile_spec"] = json.loads(str(cache_payload[6]))
     except (TypeError, ValueError, json.JSONDecodeError):
-        semantic["compile_spec"] = str(cache_payload[5])
-    if cache_payload[6]:
-        semantic["compile_kwargs_hash"] = cache_payload[6]
+        semantic["compile_spec"] = str(cache_payload[6])
+    if cache_payload[7]:
+        semantic["compile_kwargs_hash"] = cache_payload[7]
         try:
-            semantic["compile_kwargs"] = json.loads(str(cache_payload[7]))
+            semantic["compile_kwargs"] = json.loads(str(cache_payload[8]))
         except (TypeError, ValueError, json.JSONDecodeError):
-            semantic["compile_kwargs"] = str(cache_payload[7])
-    semantic["compile_options"] = _manifest_json_value(cache_payload[8])
-    semantic["compile_environment"] = _manifest_json_value(cache_payload[9])
+            semantic["compile_kwargs"] = str(cache_payload[8])
+    semantic["compile_options"] = _manifest_json_value(cache_payload[9])
+    semantic["compile_environment"] = _manifest_json_value(cache_payload[10])
     return semantic
 
 
@@ -504,31 +504,31 @@ def _manifest_integrity_status(
         return "invalid-package-fingerprint"
     if raw.get("toolchain") != cache_payload[3]:
         return "toolchain-payload-mismatch"
-    if raw.get("compile_spec_hash") != cache_payload[4]:
+    if raw.get("compile_spec_hash") != cache_payload[5]:
         return "compile-spec-payload-mismatch"
     compile_spec_json = raw.get("compile_spec_json")
-    if compile_spec_json != cache_payload[5] or not isinstance(compile_spec_json, str):
+    if compile_spec_json != cache_payload[6] or not isinstance(compile_spec_json, str):
         return "compile-spec-json-mismatch"
     if hashlib.sha256(compile_spec_json.encode("utf-8")).hexdigest() != raw.get(
         "compile_spec_hash"
     ):
         return "compile-spec-hash-mismatch"
     if (
-        raw.get("compile_kwargs_hash", "") != cache_payload[6]
-        or raw.get("compile_kwargs_json", "") != cache_payload[7]
+        raw.get("compile_kwargs_hash", "") != cache_payload[7]
+        or raw.get("compile_kwargs_json", "") != cache_payload[8]
     ):
         return "compile-kwargs-payload-mismatch"
-    if cache_payload[7]:
+    if cache_payload[8]:
         if (
-            hashlib.sha256(str(cache_payload[7]).encode("utf-8")).hexdigest()
-            != cache_payload[6]
+            hashlib.sha256(str(cache_payload[8]).encode("utf-8")).hexdigest()
+            != cache_payload[7]
         ):
             return "compile-kwargs-hash-mismatch"
-    elif cache_payload[6]:
+    elif cache_payload[7]:
         return "compile-kwargs-hash-without-json"
-    if raw.get("compile_options") != cache_payload[8]:
+    if raw.get("compile_options") != cache_payload[9]:
         return "compile-options-payload-mismatch"
-    if raw.get("compile_environment") != cache_payload[9]:
+    if raw.get("compile_environment") != cache_payload[10]:
         return "compile-environment-payload-mismatch"
     try:
         compile_spec = json.loads(compile_spec_json)

From 8a93568a6163d6a2f02e3538dba0a20d76a9f55d Mon Sep 17 00:00:00 2001
From: Brandon Music <brandon.m.music@gmail.com>
Date: Mon, 20 Jul 2026 10:43:34 -0400
Subject: [PATCH 02/11] w4a16: EXL3 trellis (trellis3_t256) fused arm

---
 sparkinfer/_lib/intrinsics.py                 |  218 ++
 .../moe/_shared/kernels/w4a16/__init__.py     |   14 +-
 sparkinfer/moe/_shared/kernels/w4a16/host.py  |   31 +-
 .../moe/_shared/kernels/w4a16/kernel.py       | 2087 ++++++++++++++++-
 .../moe/_shared/kernels/w4a16/prepare.py      |  637 +++++
 5 files changed, 2869 insertions(+), 118 deletions(-)

diff --git a/sparkinfer/_lib/intrinsics.py b/sparkinfer/_lib/intrinsics.py
index 0183a34b..5131a517 100644
--- a/sparkinfer/_lib/intrinsics.py
+++ b/sparkinfer/_lib/intrinsics.py
@@ -6016,3 +6016,221 @@ def nf3_codebook_pools(codebook) -> tuple:
         lo[reg] |= (bf16 & 0xFF) << (8 * slot)
         hi[reg] |= ((bf16 >> 8) & 0xFF) << (8 * slot)
     return lo[0], lo[1], hi[0], hi[1]
+
+# ===== TRELLIS-3.0 (QTIP/EXL3 3INST cb=1) dequant — added for NF3->trellis swap =====
+_TRELLIS_MCG  = 0xCBAC1FED   # codebook.cuh:39  MCG multiplier (cb==1)
+_TRELLIS_MASK = 0x8FFF8FFF   # codebook.cuh:41  lop3 operand b
+_TRELLIS_OR   = 0x3B603B60   # codebook.cuh:41  lop3 operand c ; immLut 0x6a == (a & b) ^ c
+
+
+@dsl_user_op
+def packed_dequant_trellis_to_half2x4(
+    win_a, win_b, bits: int = 3, *, loc=None, ip=None
+):
+    """8 trellis windows -> 4 fp16x2 fragments for one compile-time bitrate.
+
+    win_a/win_b: two 32-bit funnel windows extracted from the staged unit by the
+    caller (_trellis_funnel32). Element/lane order matches the NF3 primitive:
+      o0=(dec w0, dec w1) o1=(dec w2, dec w3) o2=(dec w4, dec w5) o3=(dec w6, dec w7)
+    with w7=win_a, w6=win_a>>bits, ... w4=win_a>>(3*bits), and
+         w3=win_b, w2=win_b>>bits, ... w0=win_b>>(3*bits).
+
+    ``bits`` is a trace-time specialization in {3,4,5,6}, so the generated PTX
+    keeps the same immediate shifts as exl3's templated dq4/dq8 implementations.
+    The caller supplies independently valid four-weight windows for 5/6 bpw,
+    whose full eight-weight span can cross three ring words.
+    """
+    bits = int(bits)
+    if bits not in (3, 4, 5, 6):
+        raise ValueError(f"unsupported trellis bitrate {bits}; expected 3, 4, 5, or 6")
+    asm = """
+        {
+            .reg .b32 w0,w1,w2,w3,w4,w5,w6,w7, lo, hi, M;
+            mov.b32 M, 0xCBAC1FED;
+            and.b32 w7, $4, 0xffff;
+            shr.u32 w6, $4, __B1__;  and.b32 w6, w6, 0xffff;
+            shr.u32 w5, $4, __B2__;  and.b32 w5, w5, 0xffff;
+            shr.u32 w4, $4, __B3__;  and.b32 w4, w4, 0xffff;
+            and.b32 w3, $5, 0xffff;
+            shr.u32 w2, $5, __B1__;  and.b32 w2, w2, 0xffff;
+            shr.u32 w1, $5, __B2__;  and.b32 w1, w1, 0xffff;
+            shr.u32 w0, $5, __B3__;  and.b32 w0, w0, 0xffff;
+            mul.lo.u32 w0, w0, M;  lop3.b32 w0, w0, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w1, w1, M;  lop3.b32 w1, w1, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w2, w2, M;  lop3.b32 w2, w2, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w3, w3, M;  lop3.b32 w3, w3, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w4, w4, M;  lop3.b32 w4, w4, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w5, w5, M;  lop3.b32 w5, w5, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w6, w6, M;  lop3.b32 w6, w6, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w7, w7, M;  lop3.b32 w7, w7, 0x8fff8fff, 0x3b603b60, 0x6a;
+            prmt.b32 lo, w0, w1, 0x5410;  prmt.b32 hi, w0, w1, 0x7632;  add.rn.f16x2 $0, lo, hi;
+            prmt.b32 lo, w2, w3, 0x5410;  prmt.b32 hi, w2, w3, 0x7632;  add.rn.f16x2 $1, lo, hi;
+            prmt.b32 lo, w4, w5, 0x5410;  prmt.b32 hi, w4, w5, 0x7632;  add.rn.f16x2 $2, lo, hi;
+            prmt.b32 lo, w6, w7, 0x5410;  prmt.b32 hi, w6, w7, 0x7632;  add.rn.f16x2 $3, lo, hi;
+        }
+        """
+    asm = (
+        asm.replace("__B1__", str(bits))
+        .replace("__B2__", str(2 * bits))
+        .replace("__B3__", str(3 * bits))
+    )
+    result = llvm.inline_asm(
+        llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]),
+        [Uint32(win_a).ir_value(loc=loc, ip=ip), Uint32(win_b).ir_value(loc=loc, ip=ip)],
+        asm,
+        "=r,=r,=r,=r,r,r",
+        has_side_effects=False, is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT, loc=loc, ip=ip,
+    )
+    o0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
+    o1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
+    o2 = llvm.extractvalue(T.i32(), result, [2], loc=loc, ip=ip)
+    o3 = llvm.extractvalue(T.i32(), result, [3], loc=loc, ip=ip)
+    return Uint32(o0), Uint32(o1), Uint32(o2), Uint32(o3)
+
+
+@dsl_user_op
+def packed_dequant_trellis_stream_to_half2x4(
+    stream_lo, stream_hi, bits: int, *, loc=None, ip=None
+):
+    """Decode eight windows from a bitrate-strided 64-bit aligned stream.
+
+    Bit zero of ``{stream_hi,stream_lo}`` begins w7; w6..w0 begin at successive
+    ``bits`` offsets. This is required at 6 bpw, where four 16-bit windows span
+    34 bits and cannot be represented by one 32-bit base window.
+    """
+    bits = int(bits)
+    if bits not in (3, 4, 5, 6):
+        raise ValueError(f"unsupported trellis bitrate {bits}; expected 3, 4, 5, or 6")
+    asm = """
+        {
+            .reg .b32 w0,w1,w2,w3,w4,w5,w6,w7, lo, hi, M;
+            mov.b32 M, 0xCBAC1FED;
+            and.b32 w7, $4, 0xffff;
+            __I6__
+            __I5__
+            __I4__
+            __I3__
+            __I2__
+            __I1__
+            __I0__
+            mul.lo.u32 w0, w0, M;  lop3.b32 w0, w0, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w1, w1, M;  lop3.b32 w1, w1, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w2, w2, M;  lop3.b32 w2, w2, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w3, w3, M;  lop3.b32 w3, w3, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w4, w4, M;  lop3.b32 w4, w4, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w5, w5, M;  lop3.b32 w5, w5, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w6, w6, M;  lop3.b32 w6, w6, 0x8fff8fff, 0x3b603b60, 0x6a;
+            mul.lo.u32 w7, w7, M;  lop3.b32 w7, w7, 0x8fff8fff, 0x3b603b60, 0x6a;
+            prmt.b32 lo, w0, w1, 0x5410;  prmt.b32 hi, w0, w1, 0x7632;  add.rn.f16x2 $0, lo, hi;
+            prmt.b32 lo, w2, w3, 0x5410;  prmt.b32 hi, w2, w3, 0x7632;  add.rn.f16x2 $1, lo, hi;
+            prmt.b32 lo, w4, w5, 0x5410;  prmt.b32 hi, w4, w5, 0x7632;  add.rn.f16x2 $2, lo, hi;
+            prmt.b32 lo, w6, w7, 0x5410;  prmt.b32 hi, w6, w7, 0x7632;  add.rn.f16x2 $3, lo, hi;
+        }
+        """
+    for name, register, multiple in (
+        ("__I6__", "w6", 1), ("__I5__", "w5", 2),
+        ("__I4__", "w4", 3), ("__I3__", "w3", 4),
+        ("__I2__", "w2", 5), ("__I1__", "w1", 6),
+        ("__I0__", "w0", 7),
+    ):
+        shift = multiple * bits
+        if shift < 32:
+            instruction = (
+                f"shf.r.wrap.b32 {register}, $4, $5, {shift}; "
+                f"and.b32 {register}, {register}, 0xffff;"
+            )
+        else:
+            instruction = (
+                f"shr.u32 {register}, $5, {shift - 32}; "
+                f"and.b32 {register}, {register}, 0xffff;"
+            )
+        asm = asm.replace(name, instruction)
+    result = llvm.inline_asm(
+        llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]),
+        [
+            Uint32(stream_lo).ir_value(loc=loc, ip=ip),
+            Uint32(stream_hi).ir_value(loc=loc, ip=ip),
+        ],
+        asm,
+        "=r,=r,=r,=r,r,r",
+        has_side_effects=False,
+        is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT,
+        loc=loc,
+        ip=ip,
+    )
+    o0 = llvm.extractvalue(T.i32(), result, [0], loc=loc, ip=ip)
+    o1 = llvm.extractvalue(T.i32(), result, [1], loc=loc, ip=ip)
+    o2 = llvm.extractvalue(T.i32(), result, [2], loc=loc, ip=ip)
+    o3 = llvm.extractvalue(T.i32(), result, [3], loc=loc, ip=ip)
+    return Uint32(o0), Uint32(o1), Uint32(o2), Uint32(o3)
+
+
+@dsl_user_op
+def packed_dequant_trellis3_to_half2x4(win_a, win_b, *, loc=None, ip=None):
+    """Compatibility wrapper for the original 3-bpw primitive."""
+    return packed_dequant_trellis_to_half2x4(win_a, win_b, 3, loc=loc, ip=ip)
+
+
+@dsl_user_op
+def _f16x2_to_bf16x2(p, *, loc=None, ip=None):
+    """{hi_f16, lo_f16} -> {hi_bf16, lo_bf16}, RNE, lane order preserved."""
+    r = llvm.inline_asm(
+        llvm.StructType.get_literal([T.i32()]),
+        [Uint32(p).ir_value(loc=loc, ip=ip)],
+        """
+        {
+            .reg .b16 hlo, hhi;
+            .reg .f32 flo, fhi;
+            mov.b32 {hlo, hhi}, $1;
+            cvt.f32.f16 flo, hlo;
+            cvt.f32.f16 fhi, hhi;
+            cvt.rn.bf16x2.f32 $0, fhi, flo;
+        }
+        """,
+        "=r,r", has_side_effects=False, is_align_stack=False,
+        asm_dialect=llvm.AsmDialect.AD_ATT, loc=loc, ip=ip,
+    )
+    return Uint32(llvm.extractvalue(T.i32(), r, [0], loc=loc, ip=ip))
+
+
+@dsl_user_op
+def packed_dequant_trellis_to_bfloat2x4(
+    win_a, win_b, bits: int = 3, *, loc=None, ip=None
+):
+    """Diagnostic bf16 conversion of the generalized fp16 trellis decode.
+
+    Production trellis kernels should use the fp16 primitive and fp16 MMA arm:
+    that is bit-faithful to exl3 and saves 16 conversion instructions per
+    fragment.  This wrapper remains available for the D4 bf16 A/B gate.
+    """
+    h0, h1, h2, h3 = packed_dequant_trellis_to_half2x4(
+        win_a, win_b, bits, loc=loc, ip=ip
+    )
+    return (_f16x2_to_bf16x2(h0, loc=loc, ip=ip), _f16x2_to_bf16x2(h1, loc=loc, ip=ip),
+            _f16x2_to_bf16x2(h2, loc=loc, ip=ip), _f16x2_to_bf16x2(h3, loc=loc, ip=ip))
+
+
+@dsl_user_op
+def packed_dequant_trellis_stream_to_bfloat2x4(
+    stream_lo, stream_hi, bits: int, *, loc=None, ip=None
+):
+    """Diagnostic bf16 conversion of the 64-bit-stream trellis decode."""
+    h0, h1, h2, h3 = packed_dequant_trellis_stream_to_half2x4(
+        stream_lo, stream_hi, bits, loc=loc, ip=ip
+    )
+    return (
+        _f16x2_to_bf16x2(h0, loc=loc, ip=ip),
+        _f16x2_to_bf16x2(h1, loc=loc, ip=ip),
+        _f16x2_to_bf16x2(h2, loc=loc, ip=ip),
+        _f16x2_to_bf16x2(h3, loc=loc, ip=ip),
+    )
+
+
+@dsl_user_op
+def packed_dequant_trellis3_to_bfloat2x4(win_a, win_b, *, loc=None, ip=None):
+    """Compatibility wrapper for the former 3-bpw bf16 primitive."""
+    return packed_dequant_trellis_to_bfloat2x4(
+        win_a, win_b, 3, loc=loc, ip=ip
+    )
diff --git a/sparkinfer/moe/_shared/kernels/w4a16/__init__.py b/sparkinfer/moe/_shared/kernels/w4a16/__init__.py
index de331843..2c953dd6 100644
--- a/sparkinfer/moe/_shared/kernels/w4a16/__init__.py
+++ b/sparkinfer/moe/_shared/kernels/w4a16/__init__.py
@@ -1,3 +1,13 @@
-"""Internal CuTeDSL W4A16 MoE implementation package."""
+"""Public preparation and execution surfaces for CuTeDSL W4A16 kernels."""
 
-__all__: list[str] = []
+from .kernel import run_trellis256_dense
+from .prepare import (
+    PreparedTrellis256DenseWeight,
+    prepare_trellis256_dense_weight,
+)
+
+__all__ = [
+    "PreparedTrellis256DenseWeight",
+    "prepare_trellis256_dense_weight",
+    "run_trellis256_dense",
+]
diff --git a/sparkinfer/moe/_shared/kernels/w4a16/host.py b/sparkinfer/moe/_shared/kernels/w4a16/host.py
index c73f199d..856a046a 100644
--- a/sparkinfer/moe/_shared/kernels/w4a16/host.py
+++ b/sparkinfer/moe/_shared/kernels/w4a16/host.py
@@ -38,6 +38,8 @@ class W4A16PackedBuffers:
     block_expert_ids: torch.Tensor | None = None
     packed_route_count: torch.Tensor | None = None
     expert_offsets: torch.Tensor | None = None
+    rotation_a_gate: torch.Tensor | None = None
+    rotation_a_up: torch.Tensor | None = None
 
 
 @dataclass(frozen=True)
@@ -50,6 +52,7 @@ class W4A16BufferPlan:
     fc2_c_tmp_elements: int
     intermediate_cache13_elements: int
     intermediate_cache2_elements: int
+    rotation_a_elements: int = 0
 
 
 def validate_activation(activation: str) -> bool:
@@ -272,6 +275,7 @@ def plan_w4a16_buffers(
     topk: int,
     route_num_experts: int | None = None,
     sms: int,
+    full_rotation: bool = False,
 ) -> W4A16BufferPlan:
     routed_rows = int(m) * int(topk)
     route_num_experts = (
@@ -305,6 +309,7 @@ def plan_w4a16_buffers(
         ),
         intermediate_cache13_elements=routed_rows * max(fc1_cols, hidden_size),
         intermediate_cache2_elements=routed_rows * intermediate_size,
+        rotation_a_elements=(routed_rows * hidden_size if full_rotation else 0),
     )
 
 
@@ -316,6 +321,7 @@ def make_w4a16_packed_buffers(
     dtype: torch.dtype,
     device: torch.device,
     route_num_experts: int | None = None,
+    full_rotation: bool = False,
 ) -> W4A16PackedBuffers:
     route_num_experts = (
         int(prepared.num_experts)
@@ -329,6 +335,7 @@ def make_w4a16_packed_buffers(
         topk=topk,
         route_num_experts=route_num_experts,
         sms=sms,
+        full_rotation=full_rotation,
     )
     fc1_c_tmp = torch.empty(
         (plan.fc1_c_tmp_elements,),
@@ -351,7 +358,11 @@ def make_w4a16_packed_buffers(
             dtype=dtype,
             device=device,
         ),
-        output=torch.empty((m, prepared.hidden_size), dtype=dtype, device=device),
+        output=torch.empty(
+            (m, prepared.hidden_size),
+            dtype=torch.float32 if full_rotation else dtype,
+            device=device,
+        ),
         fc1_c_tmp=fc1_c_tmp,
         fc2_c_tmp=fc2_c_tmp,
         packed_route_indices=torch.empty(
@@ -364,6 +375,24 @@ def make_w4a16_packed_buffers(
         expert_offsets=torch.empty(
             (route_num_experts + 1,), dtype=torch.int32, device=device
         ),
+        rotation_a_gate=(
+            torch.empty(
+                (plan.routed_rows, int(prepared.hidden_size)),
+                dtype=torch.float16,
+                device=device,
+            )
+            if full_rotation
+            else None
+        ),
+        rotation_a_up=(
+            torch.empty(
+                (plan.routed_rows, int(prepared.hidden_size)),
+                dtype=torch.float16,
+                device=device,
+            )
+            if full_rotation
+            else None
+        ),
     )
 
 
diff --git a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
index a2bfc11e..0c7f59ff 100644
--- a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
+++ b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
@@ -14,7 +14,7 @@
 import torch
 from cutlass.base_dsl.compiler import OptLevel
 from cutlass._mlir.dialects import llvm
-from cutlass.cutlass_dsl import Int32, Int64, T, Uint32, dsl_user_op
+from cutlass.cutlass_dsl import Int32, Int64, T, Uint32, Uint64, dsl_user_op
 
 from sparkinfer._lib.compiler import (
     KernelCompileSpec,
@@ -56,6 +56,10 @@
     packed_dequant_e8m0x4_to_bfloat2x2,
     packed_dequant_e8m0x4_to_half2x2,
     packed_dequant_nf3x8_to_bfloat2x4,
+    packed_dequant_trellis_to_bfloat2x4,
+    packed_dequant_trellis_to_half2x4,
+    packed_dequant_trellis_stream_to_bfloat2x4,
+    packed_dequant_trellis_stream_to_half2x4,
     nf3_codebook_pools,
     pack_f32x2_to_bfloat2,
     pack_f32x2_to_f16x2,
@@ -111,12 +115,16 @@
 _DEVICE_MAX_REG_BYTES = 255 * 1024
 _DEFAULT_MAX_SHARED_MEM = 101_376
 _SCALAR_ACC_FRAGMENT_WIDTH = 1
-_WEIGHT_LAYOUTS = {"packed", "modelopt", "nf3_2p1"}
+_WEIGHT_LAYOUTS = {"packed", "modelopt", "nf3_2p1", "trellis3_t256"}
 _MODEL_OPT_W13_LAYOUTS = {"w13", "w31"}
+_TRELLIS256_W13_LAYOUTS = {"packed", "trellis3_t256_proj"}
 # NF3 codebook: 8 bf16 codepoints for the 3-bit ("nf3_2p1") weight layout. The
 # device dequant reads these as prmt pools (see nf3_codebook_pools); the host
 # injects the 4 pool words as compile-time constants into the GEMM.
 _NF3_CODEBOOK = (-1.0, -0.6047, -0.3563, -0.1275, 0.1275, 0.3563, 0.6047, 1.0)
+# Native EXL3 t256 tiles contain 256 tail-biting codes at one compile-time
+# bitrate. Their exact storage is [16*bits] int16 == [8*bits] uint32 per tile.
+_TRELLIS256_BITS = (3, 4, 5, 6)
 _SCALE_FORMATS = {
     "e4m3_k16": "e4m3_k16",
     "e8m0_k32": "e8m0_k32",
@@ -281,17 +289,15 @@ def _shared_memory_footprint(
     tile_k: int,
     scale_format: str = "e4m3_k16",
     weight_layout: str = "packed",
+    weight_bits: int = 4,
 ) -> int:
     cta_m = int(cta_m_blocks) * 16
     cta_n = int(tile_n)
     cta_k = int(tile_k)
     sh_block_meta_size = cta_m * 16
     sh_a_size = _STAGES * (cta_m * cta_k) * 2
-    sh_b_size = _STAGES * (cta_k * cta_n // _PACK_FACTOR) * 4
-    if weight_layout == "nf3_2p1":
-        # NF3 stages 12-byte triples per 32-code unit instead of 16-byte int4
-        # units (see b_unit_bytes) -- 3/4 of the packed B stage footprint.
-        sh_b_size = sh_b_size * 3 // 4
+    staged_weight_bits = 3 if weight_layout == "nf3_2p1" else int(weight_bits)
+    sh_b_size = _STAGES * (cta_k * cta_n * staged_weight_bits // 8)
     sh_red_size = cta_m * (cta_n + 8) * 2
     sh_bias_size = cta_n * 2
     tmp_size = min(sh_b_size, sh_red_size) + sh_bias_size
@@ -316,6 +322,7 @@ def _determine_blocks_per_sm(
     max_shared_mem: int,
     scale_format: str = "e4m3_k16",
     weight_layout: str = "packed",
+    weight_bits: int = 4,
 ) -> int:
     num_regs = _w4a16_num_regs(
         cta_threads=cta_threads,
@@ -332,6 +339,7 @@ def _determine_blocks_per_sm(
         tile_k=tile_k,
         scale_format=scale_format,
         weight_layout=weight_layout,
+        weight_bits=weight_bits,
     )
     blocks_per_sm_limit = min(
         _DEVICE_MAX_REG_BYTES // register_bytes,
@@ -370,6 +378,7 @@ def _candidate_tile_fits(
     max_shared_mem: int,
     scale_format: str = "e4m3_k16",
     weight_layout: str = "packed",
+    weight_bits: int = 4,
     allow_logical_tail: bool = False,
 ) -> bool:
     if int(tile_k) == -1 or int(tile_n) == -1 or int(cta_threads) == -1:
@@ -399,6 +408,7 @@ def _candidate_tile_fits(
         tile_k=tile_k,
         scale_format=scale_format,
         weight_layout=weight_layout,
+        weight_bits=weight_bits,
     )
     return smem_bytes <= int(max_shared_mem)
 
@@ -415,6 +425,7 @@ def _select_tile_config(
     required_cta_threads: int | None = None,
     scale_format: str = "e4m3_k16",
     weight_layout: str = "packed",
+    weight_bits: int = 4,
     allow_logical_tail: bool = False,
 ) -> tuple[int, int, int, int]:
     cta_m_blocks = _covering_count(moe_block_size, 16)
@@ -439,6 +450,7 @@ def _select_tile_config(
             max_shared_mem=int(max_shared_mem) - 512,
             scale_format=scale_format,
             weight_layout=weight_layout,
+            weight_bits=weight_bits,
             allow_logical_tail=allow_logical_tail,
         ):
             continue
@@ -460,6 +472,7 @@ def _select_tile_config(
             max_shared_mem=max_shared_mem,
             scale_format=scale_format,
             weight_layout=weight_layout,
+            weight_bits=weight_bits,
         )
         if blocks_per_sm_limit > best_blocks_per_sm:
             best_blocks_per_sm = blocks_per_sm_limit
@@ -489,6 +502,8 @@ class W4A16GemmCompileResult:
     weight_layout: str = "packed"
     scale_format: str = "e4m3_k16"
     w13_layout: str = "w13"
+    dense_route_fast_path: bool = False
+    trellis_bits: int = 3
 
 
 @dataclass(frozen=True)
@@ -508,6 +523,11 @@ class W4A16TopKSumCompileResult:
     m: int
     topk: int
     hidden_size: int
+    full_rotation: bool = False
+    num_experts: int = 0
+    route_num_experts: int = 0
+    route_ids_dtype: torch.dtype = torch.int32
+    use_expert_map: bool = False
 
 
 @dataclass(frozen=True)
@@ -539,6 +559,12 @@ class W4A16FusedMoeCompileResult:
     scale_format: str = "e4m3_k16"
     tc_decode_fused_sum: bool = False
     collect_activation_amax: bool = False
+    schedule_whole_tiles: bool = False
+    intermediate_rotation: bool = False
+    dual_a: bool = False
+    trellis_bits: int = 3
+    full_rotation: bool = False
+    rotation_input_dtype: str = "fp16"
 
 
 @dataclass(frozen=True)
@@ -662,9 +688,13 @@ def __init__(
         weight_layout: str = "packed",
         scale_format: str = "e4m3_k16",
         w13_layout: str = "w13",
+        trellis_bits: int = 3,
         source_n_rotation: int = 0,
         single_token_route_fast_path: bool = False,
         direct_topk_routes: bool = False,
+        dense_route_fast_path: bool = False,
+        dual_a: bool = False,
+        route_major_a: bool = False,
         fused_topk_sum: bool = False,
         fused_sum_topk: int = 1,
         schedule_whole_tiles: bool = False,
@@ -673,10 +703,16 @@ def __init__(
             raise ValueError(f"unsupported element_dtype {element_dtype!r}")
         if weight_layout not in _WEIGHT_LAYOUTS:
             raise ValueError(f"unsupported W4A16 weight_layout {weight_layout!r}")
+        trellis_bits = int(trellis_bits)
         scale_format = _normalize_scale_format(scale_format)
         if weight_layout == "modelopt":
             if w13_layout not in _MODEL_OPT_W13_LAYOUTS:
                 raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}")
+        elif weight_layout == "trellis3_t256":
+            if w13_layout not in _TRELLIS256_W13_LAYOUTS:
+                raise ValueError(
+                    f"unsupported trellis3_t256 w13_layout {w13_layout!r}"
+                )
         else:
             w13_layout = "packed"
             source_n_rotation = 0
@@ -687,6 +723,18 @@ def __init__(
                 )
             if element_dtype != "bf16":
                 raise ValueError("nf3_2p1 W4A16 weights are bf16-activation only (v1)")
+        if weight_layout == "trellis3_t256":
+            if trellis_bits not in _TRELLIS256_BITS:
+                raise ValueError(
+                    "trellis3_t256 bits must be one of "
+                    f"{_TRELLIS256_BITS}, got {trellis_bits}"
+                )
+            if scale_format != "e4m3_k32":
+                raise ValueError(
+                    "trellis3_t256 W4A16 weights require scale_format='e4m3_k32'"
+                )
+        elif trellis_bits != 3:
+            raise ValueError("trellis_bits is only valid for trellis3_t256 weights")
         if epilogue_activation not in (None, "relu2"):
             raise ValueError(
                 "W4A16 GEMM epilogue activation currently supports only relu2"
@@ -750,7 +798,23 @@ def __init__(
         self.is_fp16 = element_dtype == "fp16"
         self.epilogue_relu2 = epilogue_activation == "relu2"
         self.weight_layout = weight_layout
+        self.trellis_bits = trellis_bits
         self.weight_layout_nf3 = weight_layout == "nf3_2p1"
+        self.weight_layout_trellis256 = weight_layout == "trellis3_t256"
+        self.weight_layout_trellis256_proj = (
+            self.weight_layout_trellis256
+            and w13_layout == "trellis3_t256_proj"
+        )
+        if self.weight_layout_trellis256_proj and (
+            self.size_n % 2 != 0 or (self.size_n // 2) % self.tile_n != 0
+        ):
+            raise ValueError(
+                "trellis3_t256_proj requires each FC1 projection to contain "
+                "an integral number of CTA N tiles"
+            )
+        self.b_region_variable = (
+            self.weight_layout_nf3 or self.weight_layout_trellis256
+        )
         self.scale_format = scale_format
         self.scale_format_e8m0_k32 = scale_format == "e8m0_k32"
         # k32 scale cadence (two K16 rows share one scale group). Shared by the
@@ -787,6 +851,26 @@ def __init__(
         self.source_n_rotation = int(source_n_rotation)
         self.single_token_route_fast_path = bool(single_token_route_fast_path)
         self.direct_topk_routes = bool(direct_topk_routes)
+        self.dense_route_fast_path = bool(dense_route_fast_path)
+        if self.dense_route_fast_path and (
+            self.direct_topk_routes
+            or self.single_token_route_fast_path
+            or self.num_experts != 1
+            or self.top_k != 1
+            or self.mul_topk_weights
+        ):
+            raise ValueError(
+                "dense_route_fast_path requires E=1, top_k=1, no top-k weights, "
+                "and no other route fast path"
+            )
+        self.dual_a = bool(dual_a)
+        if self.dual_a and not self.weight_layout_trellis256_proj:
+            raise ValueError(
+                "dual_a is only valid for projection-major trellis3_t256 FC1"
+            )
+        self.route_major_a = bool(route_major_a)
+        if self.route_major_a and not self.dual_a:
+            raise ValueError("route_major_a requires the exact dual-A FC1 path")
         self.fused_topk_sum = bool(fused_topk_sum)
         self.fused_sum_topk = int(fused_sum_topk)
         # Whole-tile persistent scheduling: every mn-tile is computed by one
@@ -794,8 +878,14 @@ def __init__(
         # the split-K tail machinery entirely. Requires the host to bound the
         # wave count; used by the exact-geometry hybrid decode schedule.
         self.schedule_whole_tiles = bool(schedule_whole_tiles)
-        if self.schedule_whole_tiles and not self.direct_topk_routes:
-            raise ValueError("schedule_whole_tiles requires direct_topk_routes")
+        if (
+            self.schedule_whole_tiles
+            and not self.direct_topk_routes
+            and not self.weight_layout_trellis256
+        ):
+            raise ValueError(
+                "schedule_whole_tiles requires direct_topk_routes or trellis3_t256"
+            )
         if self.fused_topk_sum and not self.direct_topk_routes:
             raise ValueError("fused_topk_sum requires direct_topk_routes")
         if self.fused_topk_sum and self.fused_sum_topk < 1:
@@ -825,6 +915,11 @@ def __init__(
             max_shared_mem=max_shared_mem,
             scale_format=scale_format,
             weight_layout=weight_layout,
+            weight_bits=(
+                max(4, self.trellis_bits)
+                if self.weight_layout_trellis256
+                else 4
+            ),
         )
 
         # W4A16 shared-memory geometry, in int4 units unless noted.
@@ -842,15 +937,18 @@ def __init__(
         self.b_sh_stride_threads = self.b_sh_stride
         self.b_sh_stage = self.b_sh_stride * self.cta_k_blocks
         self.b_sh_wr_iters = self.b_sh_stage // self.cta_threads
-        # NF3 ("nf3_2p1") stages 12-byte triples per 32-code unit instead of the
-        # 16-byte int4 unit; the per-thread unit count is unchanged, only the
-        # bytes-per-unit and the 16-byte cp.async chunk count differ.
-        self.b_unit_bytes = 12 if self.weight_layout_nf3 else 16
+        # Native t256 uses 4*bits bytes per 32 codes; NF3 uses 12 bytes and
+        # packed/modelopt use 16 bytes for the same logical unit.
+        self.b_unit_bytes = (
+            4 * self.trellis_bits
+            if self.weight_layout_trellis256
+            else 12 if self.weight_layout_nf3 else 16
+        )
         self.b_sh_stage_bytes = self.b_sh_stage * self.b_unit_bytes
-        if self.weight_layout_nf3:
+        if self.b_region_variable:
             if self.b_sh_stage_bytes % 16 != 0:
                 raise ValueError(
-                    "nf3_2p1 B stage bytes must be a multiple of 16; "
+                    "nf3/trellis B stage bytes must be a multiple of 16; "
                     f"got {self.b_sh_stage_bytes}"
                 )
             self.b_sh_chunks = self.b_sh_stage_bytes // 16
@@ -892,6 +990,12 @@ def __init__(
         self.sh_a_off = self.sh_s_off + _STAGES * self.s_sh_stage
         self.shared_int4 = self.sh_a_off + _STAGES * self.a_sh_stage
         self.shared_words = self.shared_int4 * 4
+        if self.shared_words * 4 > int(max_shared_mem):
+            raise ValueError(
+                "W4A16 shared-memory footprint exceeds device opt-in limit: "
+                f"{self.shared_words * 4} > {int(max_shared_mem)} bytes "
+                f"(layout={self.weight_layout})"
+            )
 
     @property
     def __cache_key__(self) -> tuple[object, ...]:
@@ -912,11 +1016,15 @@ def __cache_key__(self) -> tuple[object, ...]:
             self.element_dtype,
             self.epilogue_relu2,
             self.weight_layout,
+            self.trellis_bits,
             self.scale_format,
             self.w13_layout,
             self.source_n_rotation,
             self.single_token_route_fast_path,
             self.direct_topk_routes,
+            self.dense_route_fast_path,
+            self.dual_a,
+            self.route_major_a,
             self.fused_topk_sum,
             self.fused_sum_topk,
             self.cta_m_blocks,
@@ -1045,9 +1153,10 @@ def _mma_rhs_fragments_as_mma_a_m16n8k16_f32(
     @cute.jit
     def __call__(
         self,
-        a_bf16_flat: cute.Tensor,
+        a_bf16_ptr: cute.Pointer,
+        a_alt_bf16_ptr: cute.Pointer,
         b_i32_flat: cute.Tensor,
-        c_bf16_flat: cute.Tensor,
+        c_bf16_ptr: cute.Pointer,
         scales_i32_flat: cute.Tensor,
         global_scale: cute.Tensor,
         packed_route_indices: cute.Tensor,
@@ -1060,9 +1169,28 @@ def __call__(
         grid_x: cutlass.Int32,
         stream: cuda.CUstream,
     ):
+        a_bf16_flat = cute.make_tensor(
+            a_bf16_ptr,
+            layout=cute.make_layout(
+                (active_m * Int32(self.size_k),), stride=(1,)
+            ),
+        )
+        a_alt_bf16_flat = cute.make_tensor(
+            a_alt_bf16_ptr,
+            layout=cute.make_layout(
+                (active_m * Int32(self.size_k),), stride=(1,)
+            ),
+        )
+        c_bf16_flat = cute.make_tensor(
+            c_bf16_ptr,
+            layout=cute.make_layout(
+                (active_m * Int32(self.top_k) * Int32(self.size_n),), stride=(1,)
+            ),
+        )
         grid = (grid_x, 1, 1)
         self.kernel(
             a_bf16_flat,
+            a_alt_bf16_flat,
             b_i32_flat,
             c_bf16_flat,
             scales_i32_flat,
@@ -1085,6 +1213,7 @@ def __call__(
     def kernel(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
         b_i32_flat: cute.Tensor,
         c_bf16_flat: cute.Tensor,
         scales_i32_flat: cute.Tensor,
@@ -1117,6 +1246,7 @@ class Storage:
         grid_x, _, _ = cute.arch.grid_dim()
         self._run_persistent_gemm(
             a_bf16_flat,
+            a_alt_bf16_flat,
             b_i32_flat,
             c_bf16_flat,
             scales_i32_flat,
@@ -1138,6 +1268,7 @@ class Storage:
     def _run_persistent_gemm(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
         b_i32_flat: cute.Tensor,
         c_bf16_flat: cute.Tensor,
         scales_i32_flat: cute.Tensor,
@@ -1157,7 +1288,11 @@ def _run_persistent_gemm(
     ):
         n_tiles = Int32(self.n_tiles)
         route_blocks = active_size_m * Int32(self.top_k)
-        if cutlass.const_expr(not self.direct_topk_routes):
+        if cutlass.const_expr(self.dense_route_fast_path):
+            route_blocks = (
+                active_size_m + Int32(self.moe_block_size) - Int32(1)
+            ) // Int32(self.moe_block_size)
+        elif cutlass.const_expr(not self.direct_topk_routes):
             route_blocks = packed_route_count[Int32(0)].to(Int32) // Int32(
                 self.moe_block_size
             )
@@ -1306,13 +1441,16 @@ def _run_persistent_gemm(
                         lock_slot,
                     )
                 else:
-                    if cutlass.const_expr(self.direct_topk_routes):
+                    if cutlass.const_expr(self.dense_route_fast_path):
+                        expert_idx = Int32(0)
+                    elif cutlass.const_expr(self.direct_topk_routes):
                         expert_idx = packed_route_indices[route_block_idx].to(Int32)
                     else:
                         expert_idx = block_expert_ids[route_block_idx].to(Int32)
                     if expert_idx >= Int32(0):
                         self._run_tile(
                             a_bf16_flat,
+                            a_alt_bf16_flat,
                             b_i32_flat,
                             c_bf16_flat,
                             scales_i32_flat,
@@ -1356,6 +1494,26 @@ def _read_moe_block_data(
         global_scale_f32: cutlass.Float32,
         active_size_m: Int32,
     ) -> Int32:
+        if cutlass.const_expr(self.dense_route_fast_path):
+            block_row = route_block_idx * Int32(self.moe_block_size)
+            valid_rows = active_size_m - block_row
+            if valid_rows > Int32(self.moe_block_size):
+                valid_rows = Int32(self.moe_block_size)
+            if valid_rows < Int32(0):
+                valid_rows = Int32(0)
+            if tid < Int32(self.moe_block_size):
+                row = block_row + tid
+                st_shared_i32(
+                    smem_base + Int32(self.sh_route_off * 16) + tid * Int32(4),
+                    row,
+                )
+                st_shared_i32(
+                    smem_base + Int32(self.sh_rd_route_off * 16) + tid * Int32(4),
+                    row,
+                )
+            cute.arch.sync_threads()
+            return valid_rows
+
         if cutlass.const_expr(self.direct_topk_routes):
             if tid == Int32(0):
                 idx = route_block_idx
@@ -1378,6 +1536,8 @@ def _read_moe_block_data(
                 ].to(Int32)
                 st_shared_i32(smem_base + Int32(self.sh_route_off * 16), idx)
                 rd_row = idx // Int32(self.top_k)
+                if cutlass.const_expr(self.route_major_a):
+                    rd_row = idx
                 st_shared_i32(smem_base + Int32(self.sh_rd_route_off * 16), rd_row)
                 if cutlass.const_expr(self.mul_topk_weights):
                     safe_idx = idx
@@ -1431,6 +1591,8 @@ def _read_moe_block_data(
                 smem_base + Int32(self.sh_route_off * 16) + tid * Int32(4)
             )
             rd_row = idx // Int32(self.top_k)
+            if cutlass.const_expr(self.route_major_a):
+                rd_row = idx
             st_shared_i32(
                 smem_base + Int32(self.sh_rd_route_off * 16) + tid * Int32(4),
                 rd_row,
@@ -1460,6 +1622,7 @@ def _read_moe_block_data(
     def _run_tile(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
         b_i32_flat: cute.Tensor,
         c_bf16_flat: cute.Tensor,
         scales_i32_flat: cute.Tensor,
@@ -1483,6 +1646,7 @@ def _run_tile(
         if cutlass.const_expr(self.uses_m_block_8):
             self._run_tile_m8(
                 a_bf16_flat,
+                a_alt_bf16_flat,
                 b_i32_flat,
                 c_bf16_flat,
                 scales_i32_flat,
@@ -1506,6 +1670,7 @@ def _run_tile(
         else:
             self._run_tile_large_m(
                 a_bf16_flat,
+                a_alt_bf16_flat,
                 b_i32_flat,
                 c_bf16_flat,
                 scales_i32_flat,
@@ -1653,6 +1818,7 @@ def _a_shared_read_offset(self, tid: Int32, lanes_per_row: cutlass.Constexpr[int
     def _run_tile_m8(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
         b_i32_flat: cute.Tensor,
         c_bf16_flat: cute.Tensor,
         scales_i32_flat: cute.Tensor,
@@ -1714,6 +1880,7 @@ def _run_tile_m8(
         k_tiles = reduce_tile_count
         self._prefetch_initial_tiles(
             a_bf16_flat,
+            a_alt_bf16_flat,
             b_i32_flat,
             scales_i32_flat,
             smem_base,
@@ -1757,6 +1924,7 @@ def _run_tile_m8(
         )
         self._run_mma_pipeline(
             a_bf16_flat,
+            a_alt_bf16_flat,
             b_i32_flat,
             scales_i32_flat,
             smem_base,
@@ -1812,6 +1980,7 @@ def _run_tile_m8(
     def _run_tile_large_m(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
         b_i32_flat: cute.Tensor,
         c_bf16_flat: cute.Tensor,
         scales_i32_flat: cute.Tensor,
@@ -1896,6 +2065,7 @@ def _run_tile_large_m(
         k_tiles = reduce_tile_count
         self._prefetch_initial_tiles(
             a_bf16_flat,
+            a_alt_bf16_flat,
             b_i32_flat,
             scales_i32_flat,
             smem_base,
@@ -1939,6 +2109,7 @@ def _run_tile_large_m(
         )
         self._run_mma_pipeline(
             a_bf16_flat,
+            a_alt_bf16_flat,
             b_i32_flat,
             scales_i32_flat,
             smem_base,
@@ -1994,6 +2165,7 @@ def _run_tile_large_m(
     def _run_mma_pipeline(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
         b_i32_flat: cute.Tensor,
         scales_i32_flat: cute.Tensor,
         smem_base: Int32,
@@ -2048,6 +2220,7 @@ def _run_mma_pipeline(
 
                         self._prefetch_pipeline_step(
                             a_bf16_flat,
+                            a_alt_bf16_flat,
                             b_i32_flat,
                             scales_i32_flat,
                             smem_base,
@@ -2072,7 +2245,13 @@ def _run_mma_pipeline(
                         )
 
                         for jj in cutlass.range_constexpr(4):
-                            if cutlass.const_expr(self.weight_layout_nf3):
+                            if cutlass.const_expr(self.weight_layout_trellis256):
+                                self._scaled_dequant_b_fragment_trellis256(
+                                    b_frag,
+                                    b_scale_cur[0, jj],
+                                    b_scale_cur[1, jj],
+                                )
+                            elif cutlass.const_expr(self.weight_layout_nf3):
                                 # NF3 triple: lo16 = half (jj % 2) of word
                                 # (jj // 2); hi8 = byte jj of the hi word; scale
                                 # register is the same per-jj lane as packed.
@@ -2618,6 +2797,10 @@ def _load_b_scale_registers(
         pipe: Int32,
         kk: Int32,
     ):
+        if cutlass.const_expr(self.weight_layout_trellis256):
+            # TRELLIS-256: no per-weight scale; the (2,4) bundle carries the 4
+            # jj-tiles' pre-funnelled windows -> regs[0,jj]=win_a, regs[1,jj]=win_b.
+            return self._load_b_registers_trellis256(smem_base, tid, pipe, kk)
         if cutlass.const_expr(self.weight_layout == "modelopt"):
             q0, q1, q2, q3 = self._load_b_registers_modelopt_shared(
                 smem_base,
@@ -2820,6 +3003,128 @@ def _scaled_dequant_b_fragment_nf3(
         frag[1, 0] = self._elem2_mul(o2, s_lane1)
         frag[1, 1] = self._elem2_mul(o3, s_lane1)
 
+    @cute.jit
+    def _trellis256_lane_geom(
+        self,
+        lane: Int32,
+        weight_offset: cutlass.Constexpr[int],
+        weight_count: cutlass.Constexpr[int],
+    ):
+        bits = Int32(self.trellis_bits)
+        ring_u32 = Int32(8 * self.trellis_bits)
+        t_offset = Int32(8) * lane + Int32(weight_offset)
+        b1 = (t_offset + Int32(257)) * bits
+        b0 = b1 - Int32(16)
+        b2 = b1 + Int32((int(weight_count) - 1) * self.trellis_bits)
+        i0 = b0 >> Int32(5)
+        i2 = (b2 - Int32(1)) >> Int32(5)
+        ia = i0 - ring_u32 * (i0 >= ring_u32).to(Int32)
+        ib = i2 - ring_u32 * (i2 >= ring_u32).to(Int32)
+        s2 = (i2 + Int32(1)) * Int32(32) - b2
+        return ia, ib, s2, i2 - i0
+
+    @cute.jit
+    def _trellis_funnel256(self, a: Uint32, b: Uint32, s2: Int32):
+        merged = (Int64(a) << Int64(32)) | Int64(b)
+        return Uint32(merged >> Int64(s2))
+
+    @cute.jit
+    def _scaled_dequant_b_fragment_trellis256(
+        self,
+        frag: cute.Tensor,
+        win_a: Uint32,
+        win_b: Uint32,
+    ):
+        if cutlass.const_expr(self.trellis_bits == 6):
+            if cutlass.const_expr(self.is_fp16):
+                o0, o1, o2, o3 = packed_dequant_trellis_stream_to_half2x4(
+                    win_a, win_b, self.trellis_bits
+                )
+            else:
+                o0, o1, o2, o3 = packed_dequant_trellis_stream_to_bfloat2x4(
+                    win_a, win_b, self.trellis_bits
+                )
+        elif cutlass.const_expr(self.is_fp16):
+            o0, o1, o2, o3 = packed_dequant_trellis_to_half2x4(
+                win_a, win_b, self.trellis_bits
+            )
+        else:
+            o0, o1, o2, o3 = packed_dequant_trellis_to_bfloat2x4(
+                win_a, win_b, self.trellis_bits
+            )
+        frag[0, 0] = o0
+        frag[0, 1] = o1
+        frag[1, 0] = o2
+        frag[1, 1] = o3
+
+    @cute.jit
+    def _load_b_registers_trellis256(
+        self,
+        smem_base: Int32,
+        tid: Int32,
+        pipe: Int32,
+        kk: Int32,
+    ):
+        lane = tid & Int32(31)
+        warp_id = tid >> Int32(5)
+        warp_row = warp_id // Int32(self.tb_n_warps)
+        w_n = warp_id % Int32(self.tb_n_warps)
+        kt_local = Int32(self.b_sh_wr_iters) * warp_row + kk
+        b_region = (
+            smem_base
+            + Int32(self.sh_b_off * 16)
+            + pipe * Int32(self.b_sh_stage_bytes)
+        )
+        tile_u32 = 8 * self.trellis_bits
+        base_u32 = (
+            kt_local * Int32(self.cta_n_blocks) + Int32(4) * w_n
+        ) * Int32(tile_u32)
+        wa = [Uint32(0), Uint32(0), Uint32(0), Uint32(0)]
+        wb = [Uint32(0), Uint32(0), Uint32(0), Uint32(0)]
+        if cutlass.const_expr(self.trellis_bits <= 4):
+            ia, ib, s2, _ = self._trellis256_lane_geom(lane, 0, 8)
+            for jj in cutlass.range_constexpr(4):
+                tbase = base_u32 + Int32(jj * tile_u32)
+                a = ld_shared_u32(b_region + (tbase + ia) * Int32(4))
+                b = ld_shared_u32(b_region + (tbase + ib) * Int32(4))
+                wa[jj] = self._trellis_funnel256(a, b, s2)
+                wb[jj] = self._trellis_funnel256(
+                    a, b, s2 + Int32(4 * self.trellis_bits)
+                )
+        elif cutlass.const_expr(self.trellis_bits == 5):
+            ib0, ib1, sb, _ = self._trellis256_lane_geom(lane, 0, 4)
+            ia0, ia1, sa, _ = self._trellis256_lane_geom(lane, 4, 4)
+            for jj in cutlass.range_constexpr(4):
+                tbase = base_u32 + Int32(jj * tile_u32)
+                b0 = ld_shared_u32(b_region + (tbase + ib0) * Int32(4))
+                b1 = ld_shared_u32(b_region + (tbase + ib1) * Int32(4))
+                a0 = ld_shared_u32(b_region + (tbase + ia0) * Int32(4))
+                a1 = ld_shared_u32(b_region + (tbase + ia1) * Int32(4))
+                wb[jj] = self._trellis_funnel256(b0, b1, sb)
+                wa[jj] = self._trellis_funnel256(a0, a1, sa)
+        else:
+            i0, i2, s2, delta = self._trellis256_lane_geom(lane, 0, 8)
+            i1 = i0 + Int32(1)
+            i1 = i1 - Int32(tile_u32) * (i1 >= Int32(tile_u32)).to(Int32)
+            for jj in cutlass.range_constexpr(4):
+                tbase = base_u32 + Int32(jj * tile_u32)
+                z0 = ld_shared_u32(b_region + (tbase + i0) * Int32(4))
+                z1 = ld_shared_u32(b_region + (tbase + i1) * Int32(4))
+                z2 = ld_shared_u32(b_region + (tbase + i2) * Int32(4))
+                stream64 = Uint64(0)
+                if delta == Int32(1):
+                    stream64 = (
+                        (Uint64(z0) << Uint64(32)) | Uint64(z2)
+                    ) >> Uint64(s2)
+                else:
+                    lower64 = (Uint64(z1) << Uint64(32)) | Uint64(z2)
+                    stream64 = (lower64 >> Uint64(s2)) | (
+                        Uint64(z0) << (Uint64(64) - Uint64(s2))
+                    )
+                wa[jj] = Uint32(stream64)
+                wb[jj] = Uint32(stream64 >> Uint64(32))
+        return wa[0], wa[1], wa[2], wa[3], wb[0], wb[1], wb[2], wb[3]
+
     @cute.jit
     def _mma_accumulate_m8(
         self,
@@ -3105,6 +3410,7 @@ def _load_b_registers_modelopt_shared(
     def _stage_k_tile_async(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
         b_i32_flat: cute.Tensor,
         scales_i32_flat: cute.Tensor,
         smem_base: Int32,
@@ -3144,22 +3450,84 @@ def _stage_k_tile_async(
                     Int32(i * self.a_sh_wr_delta) + a_sh_wr
                 ),
             )
+            a_src = get_ptr_as_int64(a_bf16_flat, a_int4 * Int32(8))
+            if cutlass.const_expr(self.dual_a):
+                # Projection-major FC1 validates that each CTA N tile is wholly
+                # gate or wholly up.  Select the matching pre-rotated A operand
+                # with the same logical projection boundary used by B staging.
+                # The selected source still lands in the one existing A SMEM
+                # stage, so this adds neither shared memory nor MMA work.
+                if output_n_tile >= Int32(self.n_tiles // 2):
+                    a_src = get_ptr_as_int64(
+                        a_alt_bf16_flat, a_int4 * Int32(8)
+                    )
             if cutlass.const_expr(self.has_k_tile_tail):
                 a_k_int4 = tile_idx * Int32(self.a_gl_rd_delta_o) + a_gl_rd_col0
                 if row < block_valid_rows and a_k_int4 < a_gl_stride:
                     cp_async4_shared_global(
                         a_dst,
-                        get_ptr_as_int64(a_bf16_flat, a_int4 * Int32(8)),
+                        a_src,
                     )
                 else:
                     st_shared_v4_u32(a_dst, Uint32(0), Uint32(0), Uint32(0), Uint32(0))
             else:
                 cp_async4_shared_global_pred(
                     a_dst,
-                    get_ptr_as_int64(a_bf16_flat, a_int4 * Int32(8)),
+                    a_src,
                     (row < block_valid_rows).to(Int32),
                 )
 
+        if cutlass.const_expr(self.weight_layout_trellis256):
+            t256_n16 = self.size_n // 16
+            t256_tile_u32 = 8 * self.trellis_bits
+            t256_expert_u32 = (self.size_k // 16) * t256_n16 * t256_tile_u32
+            t256_chunks_per_kt = self.cta_n_blocks * (2 * self.trellis_bits)
+            t256_total_chunks = self.cta_k_blocks * t256_chunks_per_kt
+            for i in cutlass.range_constexpr(self.b_sh_wr_iters_nf3):
+                t256_chunk = Int32(i * self.cta_threads) + tid
+                t256_kt = t256_chunk // Int32(t256_chunks_per_kt)
+                t256_in_kt = t256_chunk - t256_kt * Int32(t256_chunks_per_kt)
+                b_dst = (
+                    smem_base
+                    + Int32(self.sh_b_off * 16)
+                    + pipe * Int32(self.b_sh_stage_bytes)
+                    + t256_chunk * Int32(16)
+                )
+                if cutlass.const_expr(self.weight_layout_trellis256_proj):
+                    t256_half_n16 = t256_n16 // 2
+                    t256_out_n16 = output_n_tile * Int32(self.cta_n_blocks)
+                    t256_proj = (
+                        t256_out_n16 >= Int32(t256_half_n16)
+                    ).to(Int32)
+                    t256_local_n16 = (
+                        t256_out_n16 - t256_proj * Int32(t256_half_n16)
+                    )
+                    t256_proj_expert_u32 = (
+                        (self.size_k // 16) * t256_half_n16 * t256_tile_u32
+                    )
+                    t256_plane_u32 = self.num_experts * t256_proj_expert_u32
+                    b_src_i32 = (
+                        t256_proj * Int32(t256_plane_u32)
+                        + expert_idx * Int32(t256_proj_expert_u32)
+                        + (tile_idx * Int32(self.cta_k_blocks) + t256_kt)
+                        * Int32(t256_half_n16 * t256_tile_u32)
+                        + t256_local_n16 * Int32(t256_tile_u32)
+                        + t256_in_kt * Int32(4)
+                    )
+                else:
+                    b_src_i32 = (
+                        Int32(t256_expert_u32) * expert_idx
+                        + (tile_idx * Int32(self.cta_k_blocks) + t256_kt)
+                        * Int32(t256_n16 * t256_tile_u32)
+                        + output_n_tile * Int32(self.cta_n_blocks * t256_tile_u32)
+                        + t256_in_kt * Int32(4)
+                    )
+                cp_async4_shared_global_pred(
+                    b_dst,
+                    get_ptr_as_int64(b_i32_flat, b_src_i32),
+                    (t256_chunk < Int32(t256_total_chunks)).to(Int32),
+                )
+
         if cutlass.const_expr(self.weight_layout_nf3):
             # NF3 flat-span staging: the packer lays the per-(expert,
             # output_n_tile, k-stage) B block out as ONE contiguous 12-byte-triple
@@ -3192,7 +3560,7 @@ def _stage_k_tile_async(
                 )
 
         for i in cutlass.range_constexpr(
-            0 if self.weight_layout_nf3 else self.b_sh_wr_iters
+            0 if self.b_region_variable else self.b_sh_wr_iters
         ):
             b_src_int4 = (
                 b_gl_rd_base
@@ -3221,40 +3589,47 @@ def _stage_k_tile_async(
                     Int32(i * self.cta_threads) + tid,
                 )
 
-        if tid < Int32(self.s_sh_stage):
-            s_k_group = tile_idx * Int32(self.s_tb_groups) + tid // Int32(
-                self.s_sh_stride
-            )
-            s_n_group = Int32(self.s_sh_stride) * output_n_tile + (
-                tid % Int32(self.s_sh_stride)
-            )
-            s_src_int4 = scales_expert_off + s_gl_stride * s_k_group + s_n_group
-            s_dst = self._int4_addr(
-                smem_base,
-                Int32(self.sh_s_off) + pipe * Int32(self.s_sh_stage) + tid,
-            )
-            if cutlass.const_expr(self.has_logical_tail):
-                if s_k_group < Int32(self.scale_k_groups) and s_n_group < Int32(
-                    self.scale_n_groups
-                ):
+        # trellis3_t256 has no per-weight scale and its register-load arm returns
+        # before touching scale SMEM.  Const-expr-elide the otherwise dead HBM
+        # reads so every layer can share a four-byte aligned dummy scale tensor
+        # instead of retaining 54 MiB of packed ones.
+        if cutlass.const_expr(not self.weight_layout_trellis256):
+            if tid < Int32(self.s_sh_stage):
+                s_k_group = tile_idx * Int32(self.s_tb_groups) + tid // Int32(
+                    self.s_sh_stride
+                )
+                s_n_group = Int32(self.s_sh_stride) * output_n_tile + (
+                    tid % Int32(self.s_sh_stride)
+                )
+                s_src_int4 = scales_expert_off + s_gl_stride * s_k_group + s_n_group
+                s_dst = self._int4_addr(
+                    smem_base,
+                    Int32(self.sh_s_off) + pipe * Int32(self.s_sh_stage) + tid,
+                )
+                if cutlass.const_expr(self.has_logical_tail):
+                    if s_k_group < Int32(self.scale_k_groups) and s_n_group < Int32(
+                        self.scale_n_groups
+                    ):
+                        cp_async4_shared_global(
+                            s_dst,
+                            get_ptr_as_int64(scales_i32_flat, s_src_int4 * Int32(4)),
+                        )
+                    else:
+                        st_shared_v4_u32(
+                            s_dst, Uint32(0), Uint32(0), Uint32(0), Uint32(0)
+                        )
+                else:
                     cp_async4_shared_global(
                         s_dst,
                         get_ptr_as_int64(scales_i32_flat, s_src_int4 * Int32(4)),
                     )
-                else:
-                    st_shared_v4_u32(s_dst, Uint32(0), Uint32(0), Uint32(0), Uint32(0))
-            else:
-                cp_async4_shared_global(
-                    s_dst,
-                    get_ptr_as_int64(scales_i32_flat, s_src_int4 * Int32(4)),
-                )
-
         cute.arch.cp_async_commit_group()
 
     @cute.jit
     def _prefetch_pipeline_step(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
         b_i32_flat: cute.Tensor,
         scales_i32_flat: cute.Tensor,
         smem_base: Int32,
@@ -3280,6 +3655,7 @@ def _prefetch_pipeline_step(
         if cutlass.const_expr(kk == self.b_sh_wr_iters - 2):
             self._prefetch_lookahead_tile(
                 a_bf16_flat,
+                a_alt_bf16_flat,
                 b_i32_flat,
                 scales_i32_flat,
                 smem_base,
@@ -3306,6 +3682,7 @@ def _prefetch_pipeline_step(
     def _prefetch_initial_tiles(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
         b_i32_flat: cute.Tensor,
         scales_i32_flat: cute.Tensor,
         smem_base: Int32,
@@ -3329,6 +3706,7 @@ def _prefetch_initial_tiles(
             if Int32(pipe) < k_tiles:
                 self._stage_k_tile_async(
                     a_bf16_flat,
+                    a_alt_bf16_flat,
                     b_i32_flat,
                     scales_i32_flat,
                     smem_base,
@@ -3357,6 +3735,7 @@ def _prefetch_initial_tiles(
     def _prefetch_lookahead_tile(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
         b_i32_flat: cute.Tensor,
         scales_i32_flat: cute.Tensor,
         smem_base: Int32,
@@ -3382,6 +3761,7 @@ def _prefetch_lookahead_tile(
         if fetch_tile < k_tiles:
             self._stage_k_tile_async(
                 a_bf16_flat,
+                a_alt_bf16_flat,
                 b_i32_flat,
                 scales_i32_flat,
                 smem_base,
@@ -4223,11 +4603,15 @@ def __init__(
         weight_layout: str = "packed",
         scale_format: str = "e4m3_k16",
         w13_layout: str = "w13",
+        trellis_bits: int = 3,
         direct_topk_routes: bool = False,
         tc_decode_fused_sum: bool = False,
         tc_zero_output: bool = True,
         collect_activation_amax: bool = False,
         schedule_whole_tiles: bool = False,
+        intermediate_rotation: bool = False,
+        full_rotation: bool = False,
+        rotation_input_dtype: str = "fp16",
     ):
         activation = normalize_moe_activation(activation)
         is_gated = validate_activation(activation)
@@ -4243,6 +4627,11 @@ def __init__(
         if weight_layout == "modelopt":
             if w13_layout not in _MODEL_OPT_W13_LAYOUTS:
                 raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}")
+        elif weight_layout == "trellis3_t256":
+            if w13_layout not in _TRELLIS256_W13_LAYOUTS:
+                raise ValueError(
+                    f"unsupported trellis3_t256 w13_layout {w13_layout!r}"
+                )
         else:
             w13_layout = "packed"
         self.tc_decode_fused_sum = bool(tc_decode_fused_sum)
@@ -4278,6 +4667,7 @@ def __init__(
         self.swiglu_alpha = float(swiglu_alpha)
         self.swiglu_beta = float(swiglu_beta)
         self.weight_layout = weight_layout
+        self.trellis_bits = int(trellis_bits)
         self.scale_format = scale_format
         self.w13_layout = w13_layout
         self.apply_router_weight_on_input = bool(apply_router_weight_on_input)
@@ -4286,7 +4676,51 @@ def __init__(
         self.is_fp16 = element_dtype == "fp16"
         self.fast_math = bool(fast_math)
         self.direct_topk_routes = bool(direct_topk_routes)
-        self.schedule_whole_tiles = bool(schedule_whole_tiles)
+        self.schedule_whole_tiles = bool(
+            schedule_whole_tiles or weight_layout == "trellis3_t256"
+        )
+        self.intermediate_rotation = bool(intermediate_rotation)
+        if self.intermediate_rotation:
+            if weight_layout != "trellis3_t256":
+                raise ValueError(
+                    "intermediate_rotation is only supported for trellis3_t256"
+                )
+            if not is_gated or self.activation_is_swigluoai or self.has_swiglu_limit:
+                raise ValueError(
+                    "intermediate_rotation requires plain gated silu "
+                    "(no swiglu limit/oai)"
+                )
+            if int(intermediate_size) % 128 != 0:
+                raise ValueError(
+                    "intermediate_rotation requires intermediate_size % 128 == 0"
+                )
+        self.full_rotation = bool(full_rotation)
+        self.rotation_input_dtype = str(rotation_input_dtype)
+        self.rotation_input_is_fp16 = self.rotation_input_dtype == "fp16"
+        if self.full_rotation:
+            if not self.intermediate_rotation:
+                raise ValueError("full_rotation requires intermediate_rotation")
+            if element_dtype != "fp16":
+                raise ValueError("full_rotation requires fp16 GEMM operands")
+            if self.rotation_input_dtype not in {"bf16", "fp16"}:
+                raise ValueError(
+                    "full_rotation input dtype must be 'bf16' or 'fp16'"
+                )
+            if self.direct_topk_routes or self.tc_decode_fused_sum:
+                raise ValueError(
+                    "full_rotation requires the route-packed non-TC-decode path"
+                )
+            if self.apply_router_weight_on_input:
+                raise ValueError(
+                    "full_rotation applies router weights only in the fp32 top-k sum"
+                )
+            if int(hidden_size) % 128 != 0:
+                raise ValueError("full_rotation requires hidden_size % 128 == 0")
+        self.dual_a = bool(
+            self.intermediate_rotation
+            and weight_layout == "trellis3_t256"
+            and w13_layout == "trellis3_t256_proj"
+        )
         fc1_source_n_rotation = (
             int(intermediate_size)
             if (weight_layout == "modelopt" and w13_layout == "w13" and is_gated)
@@ -4308,9 +4742,12 @@ def __init__(
             weight_layout=weight_layout,
             scale_format=scale_format,
             w13_layout=w13_layout,
+            trellis_bits=self.trellis_bits,
             source_n_rotation=fc1_source_n_rotation,
             single_token_route_fast_path=size_m == 1 and not self.direct_topk_routes,
             direct_topk_routes=self.direct_topk_routes,
+            dual_a=self.dual_a,
+            route_major_a=self.full_rotation,
             schedule_whole_tiles=self.schedule_whole_tiles,
         )
         self.fc2 = W4A16GemmKernel(
@@ -4319,7 +4756,9 @@ def __init__(
             size_k=intermediate_size,
             num_experts=num_experts,
             top_k=1,
-            mul_topk_weights=not bool(apply_router_weight_on_input),
+            mul_topk_weights=(
+                not bool(apply_router_weight_on_input) and not self.full_rotation
+            ),
             tile_n=fc2_tile_n,
             tile_k=fc2_tile_k,
             moe_block_size=moe_block_size,
@@ -4327,7 +4766,10 @@ def __init__(
             element_dtype=element_dtype,
             weight_layout=weight_layout,
             scale_format=scale_format,
-            w13_layout=w13_layout,
+            w13_layout=(
+                "packed" if weight_layout == "trellis3_t256" else w13_layout
+            ),
+            trellis_bits=self.trellis_bits,
             single_token_route_fast_path=size_m == 1 and not self.direct_topk_routes,
             direct_topk_routes=self.direct_topk_routes,
             fused_topk_sum=self.tc_decode_fused_sum,
@@ -4361,6 +4803,7 @@ def __cache_key__(self) -> tuple[object, ...]:
             self.swiglu_alpha,
             self.swiglu_beta,
             self.weight_layout,
+            self.trellis_bits,
             self.scale_format,
             self.apply_router_weight_on_input,
             self.zero_fc2_output,
@@ -4369,6 +4812,10 @@ def __cache_key__(self) -> tuple[object, ...]:
             self.direct_topk_routes,
             self.tc_zero_output,
             self.collect_activation_amax,
+            self.intermediate_rotation,
+            self.dual_a,
+            self.full_rotation,
+            self.rotation_input_dtype,
             self.fc1.__cache_key__,
             self.fc2.__cache_key__,
             self.cta_threads,
@@ -4404,6 +4851,8 @@ def _clamp_swiglu_inputs(
     def __call__(
         self,
         a_bf16_ptr: cute.Pointer,
+        a_alt_bf16_ptr: cute.Pointer,
+        rotation_input_ptr: cute.Pointer,
         w13_i32_flat: cute.Tensor,
         w2_i32_flat: cute.Tensor,
         fc1_bf16_flat: cute.Tensor,
@@ -4422,13 +4871,33 @@ def __call__(
         fc1_c_tmp_f32_flat: cute.Tensor,
         fc2_c_tmp_f32_flat: cute.Tensor,
         locks_i32_flat: cute.Tensor,
+        rot_scales_flat: cute.Tensor,
+        suh_gate_flat: cute.Tensor,
+        suh_up_flat: cute.Tensor,
         active_m: cutlass.Int32,
         grid_x: cutlass.Int32,
         stream: cuda.CUstream,
     ):
+        a_rows = active_m
+        if cutlass.const_expr(self.full_rotation):
+            a_rows = active_m * Int32(self.top_k)
         a_bf16_flat = cute.make_tensor(
             a_bf16_ptr,
-            layout=cute.make_layout((active_m * Int32(self.hidden_size),), stride=(1,)),
+            layout=cute.make_layout(
+                (a_rows * Int32(self.hidden_size),), stride=(1,)
+            ),
+        )
+        a_alt_bf16_flat = cute.make_tensor(
+            a_alt_bf16_ptr,
+            layout=cute.make_layout(
+                (a_rows * Int32(self.hidden_size),), stride=(1,)
+            ),
+        )
+        rotation_input_flat = cute.make_tensor(
+            rotation_input_ptr,
+            layout=cute.make_layout(
+                (active_m * Int32(self.hidden_size),), stride=(1,)
+            ),
         )
         topk_weights_flat = cute.make_tensor(
             topk_weights_ptr,
@@ -4437,6 +4906,8 @@ def __call__(
         grid = (grid_x, 1, 1)
         self.kernel(
             a_bf16_flat,
+            a_alt_bf16_flat,
+            rotation_input_flat,
             w13_i32_flat,
             w2_i32_flat,
             fc1_bf16_flat,
@@ -4455,6 +4926,9 @@ def __call__(
             fc1_c_tmp_f32_flat,
             fc2_c_tmp_f32_flat,
             locks_i32_flat,
+            rot_scales_flat,
+            suh_gate_flat,
+            suh_up_flat,
             active_m,
         ).launch(
             grid=grid,
@@ -4472,6 +4946,8 @@ def __call__(
     def kernel(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
+        rotation_input_flat: cute.Tensor,
         w13_i32_flat: cute.Tensor,
         w2_i32_flat: cute.Tensor,
         fc1_bf16_flat: cute.Tensor,
@@ -4490,6 +4966,9 @@ def kernel(
         fc1_c_tmp_f32_flat: cute.Tensor,
         fc2_c_tmp_f32_flat: cute.Tensor,
         locks_i32_flat: cute.Tensor,
+        rot_scales_flat: cute.Tensor,
+        suh_gate_flat: cute.Tensor,
+        suh_up_flat: cute.Tensor,
         active_m: cutlass.Int32,
     ):
         tidx, _, _ = cute.arch.thread_idx()
@@ -4513,6 +4992,8 @@ class Storage:
 
         self._moe_body(
             a_bf16_flat,
+            a_alt_bf16_flat,
+            rotation_input_flat,
             w13_i32_flat,
             w2_i32_flat,
             fc1_bf16_flat,
@@ -4531,6 +5012,9 @@ class Storage:
             fc1_c_tmp_f32_flat,
             fc2_c_tmp_f32_flat,
             locks_i32_flat,
+            rot_scales_flat,
+            suh_gate_flat,
+            suh_up_flat,
             smem_base,
             tid,
             cta,
@@ -4542,6 +5026,8 @@ class Storage:
     def _moe_body(
         self,
         a_bf16_flat: cute.Tensor,
+        a_alt_bf16_flat: cute.Tensor,
+        rotation_input_flat: cute.Tensor,
         w13_i32_flat: cute.Tensor,
         w2_i32_flat: cute.Tensor,
         fc1_bf16_flat: cute.Tensor,
@@ -4560,6 +5046,9 @@ def _moe_body(
         fc1_c_tmp_f32_flat: cute.Tensor,
         fc2_c_tmp_f32_flat: cute.Tensor,
         locks_i32_flat: cute.Tensor,
+        rot_scales_flat: cute.Tensor,
+        suh_gate_flat: cute.Tensor,
+        suh_up_flat: cute.Tensor,
         smem_base: Int32,
         tid: Int32,
         cta: Int32,
@@ -4573,6 +5062,23 @@ def _moe_body(
         # barrier, FC2. The emit hooks delegate per-tile expert resolution and
         # dispatch (used by the hybrid route map); None keeps the single-tier
         # resolution inside _run_persistent_gemm.
+        if cutlass.const_expr(self.full_rotation):
+            self._run_input_rotation(
+                rotation_input_flat,
+                a_bf16_flat,
+                a_alt_bf16_flat,
+                suh_gate_flat,
+                suh_up_flat,
+                packed_route_indices,
+                block_expert_ids,
+                packed_route_count,
+                tid,
+                cta,
+                grid_x,
+                active_m,
+            )
+            self._grid_barrier(locks_i32_flat, tid, grid_x)
+
         if cutlass.const_expr(self.tc_decode_fused_sum):
             # The TC-decode FC2 epilogue atomically accumulates per-route
             # partials directly into the per-token output, so the output must be
@@ -4605,6 +5111,7 @@ def _moe_body(
         if cutlass.const_expr(self.activation_is_gated):
             self.fc1._run_persistent_gemm(
                 a_bf16_flat,
+                a_alt_bf16_flat,
                 w13_i32_flat,
                 fc1_bf16_flat,
                 w13_scales_i32_flat,
@@ -4623,17 +5130,33 @@ def _moe_body(
                 fc1_emit_tile,
             )
             self._grid_barrier(locks_i32_flat, tid, grid_x)
-            self._run_activation(
-                fc1_bf16_flat,
-                activated_bf16_flat,
-                tid,
-                cta,
-                grid_x,
-                active_m,
-            )
+            if cutlass.const_expr(self.full_rotation):
+                self._run_activation_compact(
+                    fc1_bf16_flat,
+                    activated_bf16_flat,
+                    rot_scales_flat,
+                    packed_route_indices,
+                    block_expert_ids,
+                    packed_route_count,
+                    tid,
+                    cta,
+                    grid_x,
+                    active_m,
+                )
+            else:
+                self._run_activation(
+                    fc1_bf16_flat,
+                    activated_bf16_flat,
+                    rot_scales_flat,
+                    tid,
+                    cta,
+                    grid_x,
+                    active_m,
+                )
         else:
             self.fc1._run_persistent_gemm(
                 a_bf16_flat,
+                a_alt_bf16_flat,
                 w13_i32_flat,
                 activated_bf16_flat,
                 w13_scales_i32_flat,
@@ -4671,6 +5194,7 @@ def _moe_body(
             self._zero_fc2_output(fc2_bf16_flat, tid, cta, grid_x, active_m)
             self._grid_barrier(locks_i32_flat, tid, grid_x)
         self.fc2._run_persistent_gemm(
+            activated_bf16_flat,
             activated_bf16_flat,
             w2_i32_flat,
             fc2_bf16_flat,
@@ -4713,6 +5237,181 @@ def _grid_barrier(
                     sense = ld_global_acquire_i32(sense_addr)
         cute.arch.sync_threads()
 
+    @cute.jit
+    def _run_input_rotation(
+        self,
+        x_input_flat: cute.Tensor,
+        a_gate_flat: cute.Tensor,
+        a_up_flat: cute.Tensor,
+        suh_gate_flat: cute.Tensor,
+        suh_up_flat: cute.Tensor,
+        packed_route_indices: cute.Tensor,
+        block_expert_ids: cute.Tensor,
+        packed_route_count: cute.Tensor,
+        tid: Int32,
+        cta: Int32,
+        grid_x: Int32,
+        active_m: cutlass.Int32,
+    ):
+        # One warp owns one (packed route, H128 block).  The packed route id is
+        # the logical flattened top-k route, so route//top_k is the source token
+        # and route is the row in each projection-specific A scratch.  Padding
+        # slots carry the live-route sentinel and are skipped.  Both scale
+        # multiplies are rounded to fp16 before the existing fp32-register
+        # H128 butterfly, and each H128 result is rounded to fp16 on store.
+        lane = tid & Int32(31)
+        warp_in_cta = tid >> Int32(5)
+        warps_per_cta = Int32(self.cta_threads // 32)
+        nblk = Int32(self.hidden_size // 128)
+        live_routes = active_m * Int32(self.top_k)
+        route_count = packed_route_count[Int32(0)].to(Int32)
+        gwarp = cta * warps_per_cta + warp_in_cta
+        gw_stride = grid_x * warps_per_cta
+        total_units = route_count * nblk
+        elem = lane * Int32(4)
+        unit = gwarp
+        while unit < total_units:
+            route_pos = unit // nblk
+            blk = unit - route_pos * nblk
+            route = packed_route_indices[route_pos].to(Int32)
+            expert = block_expert_ids[
+                route_pos // Int32(self.moe_block_size)
+            ].to(Int32)
+            if route >= Int32(0) and route < live_routes and expert >= Int32(0):
+                token = route // Int32(self.top_k)
+                col0 = blk * Int32(128) + elem
+                x_base = token * Int32(self.hidden_size) + col0
+                s_base = expert * Int32(self.hidden_size) + col0
+                out_base = route * Int32(self.hidden_size) + col0
+
+                x0 = cutlass.Float16(
+                    x_input_flat[x_base + Int32(0)].to(cutlass.Float32)
+                ).to(cutlass.Float32)
+                x1 = cutlass.Float16(
+                    x_input_flat[x_base + Int32(1)].to(cutlass.Float32)
+                ).to(cutlass.Float32)
+                x2 = cutlass.Float16(
+                    x_input_flat[x_base + Int32(2)].to(cutlass.Float32)
+                ).to(cutlass.Float32)
+                x3 = cutlass.Float16(
+                    x_input_flat[x_base + Int32(3)].to(cutlass.Float32)
+                ).to(cutlass.Float32)
+                sg0 = suh_gate_flat[s_base + Int32(0)].to(cutlass.Float32)
+                sg1 = suh_gate_flat[s_base + Int32(1)].to(cutlass.Float32)
+                sg2 = suh_gate_flat[s_base + Int32(2)].to(cutlass.Float32)
+                sg3 = suh_gate_flat[s_base + Int32(3)].to(cutlass.Float32)
+                su0 = suh_up_flat[s_base + Int32(0)].to(cutlass.Float32)
+                su1 = suh_up_flat[s_base + Int32(1)].to(cutlass.Float32)
+                su2 = suh_up_flat[s_base + Int32(2)].to(cutlass.Float32)
+                su3 = suh_up_flat[s_base + Int32(3)].to(cutlass.Float32)
+
+                g0 = cutlass.Float16(x0 * sg0).to(cutlass.Float32)
+                g1 = cutlass.Float16(x1 * sg1).to(cutlass.Float32)
+                g2 = cutlass.Float16(x2 * sg2).to(cutlass.Float32)
+                g3 = cutlass.Float16(x3 * sg3).to(cutlass.Float32)
+                u0 = cutlass.Float16(x0 * su0).to(cutlass.Float32)
+                u1 = cutlass.Float16(x1 * su1).to(cutlass.Float32)
+                u2 = cutlass.Float16(x2 * su2).to(cutlass.Float32)
+                u3 = cutlass.Float16(x3 * su3).to(cutlass.Float32)
+                gh0, gh1, gh2, gh3 = self._had128_quad(g0, g1, g2, g3, lane)
+                uh0, uh1, uh2, uh3 = self._had128_quad(u0, u1, u2, u3, lane)
+                a_gate_flat[out_base + Int32(0)] = cutlass.Float16(gh0)
+                a_gate_flat[out_base + Int32(1)] = cutlass.Float16(gh1)
+                a_gate_flat[out_base + Int32(2)] = cutlass.Float16(gh2)
+                a_gate_flat[out_base + Int32(3)] = cutlass.Float16(gh3)
+                a_up_flat[out_base + Int32(0)] = cutlass.Float16(uh0)
+                a_up_flat[out_base + Int32(1)] = cutlass.Float16(uh1)
+                a_up_flat[out_base + Int32(2)] = cutlass.Float16(uh2)
+                a_up_flat[out_base + Int32(3)] = cutlass.Float16(uh3)
+            unit += gw_stride
+
+    @cute.jit
+    def _run_activation_compact(
+        self,
+        fc1_bf16_flat: cute.Tensor,
+        activated_bf16_flat: cute.Tensor,
+        rot_scales_flat: cute.Tensor,
+        packed_route_indices: cute.Tensor,
+        block_expert_ids: cute.Tensor,
+        packed_route_count: cute.Tensor,
+        tid: Int32,
+        cta: Int32,
+        grid_x: Int32,
+        active_m: cutlass.Int32,
+    ):
+        # A9: traverse the already-packed expert blocks and index the persistent
+        # [E,3I] table directly.  FC1/activation rows stay at their logical
+        # flattened route ids, so packing order cannot change arithmetic or the
+        # per-row output location.
+        lane = tid & Int32(31)
+        warp_in_cta = tid >> Int32(5)
+        warps_per_cta = Int32(self.cta_threads // 32)
+        nblk = Int32(self.intermediate_size // 128)
+        live_routes = active_m * Int32(self.top_k)
+        route_count = packed_route_count[Int32(0)].to(Int32)
+        fc1_cols = Int32(self.fc1_cols)
+        isz = Int32(self.intermediate_size)
+        rot_row = Int32(3 * self.intermediate_size)
+        gwarp = cta * warps_per_cta + warp_in_cta
+        gw_stride = grid_x * warps_per_cta
+        total_units = route_count * nblk
+        elem = lane * Int32(4)
+        unit = gwarp
+        while unit < total_units:
+            route_pos = unit // nblk
+            blk = unit - route_pos * nblk
+            row = packed_route_indices[route_pos].to(Int32)
+            expert = block_expert_ids[
+                route_pos // Int32(self.moe_block_size)
+            ].to(Int32)
+            if row >= Int32(0) and row < live_routes and expert >= Int32(0):
+                col0 = blk * Int32(128) + elem
+                g_base = row * fc1_cols + col0
+                u_base = g_base + isz
+                s_base = expert * rot_row + col0
+                g0 = fc1_bf16_flat[g_base + Int32(0)].to(cutlass.Float32)
+                g1 = fc1_bf16_flat[g_base + Int32(1)].to(cutlass.Float32)
+                g2 = fc1_bf16_flat[g_base + Int32(2)].to(cutlass.Float32)
+                g3 = fc1_bf16_flat[g_base + Int32(3)].to(cutlass.Float32)
+                u0 = fc1_bf16_flat[u_base + Int32(0)].to(cutlass.Float32)
+                u1 = fc1_bf16_flat[u_base + Int32(1)].to(cutlass.Float32)
+                u2 = fc1_bf16_flat[u_base + Int32(2)].to(cutlass.Float32)
+                u3 = fc1_bf16_flat[u_base + Int32(3)].to(cutlass.Float32)
+                gh0, gh1, gh2, gh3 = self._had128_quad(g0, g1, g2, g3, lane)
+                uh0, uh1, uh2, uh3 = self._had128_quad(u0, u1, u2, u3, lane)
+                svg0 = rot_scales_flat[s_base + Int32(0)].to(cutlass.Float32)
+                svg1 = rot_scales_flat[s_base + Int32(1)].to(cutlass.Float32)
+                svg2 = rot_scales_flat[s_base + Int32(2)].to(cutlass.Float32)
+                svg3 = rot_scales_flat[s_base + Int32(3)].to(cutlass.Float32)
+                svu0 = rot_scales_flat[s_base + isz + Int32(0)].to(cutlass.Float32)
+                svu1 = rot_scales_flat[s_base + isz + Int32(1)].to(cutlass.Float32)
+                svu2 = rot_scales_flat[s_base + isz + Int32(2)].to(cutlass.Float32)
+                svu3 = rot_scales_flat[s_base + isz + Int32(3)].to(cutlass.Float32)
+                ig0 = gh0 * svg0
+                ig1 = gh1 * svg1
+                ig2 = gh2 * svg2
+                ig3 = gh3 * svg3
+                iu0 = uh0 * svu0
+                iu1 = uh1 * svu1
+                iu2 = uh2 * svu2
+                iu3 = uh3 * svu3
+                down = isz + isz
+                sd0 = rot_scales_flat[s_base + down + Int32(0)].to(cutlass.Float32)
+                sd1 = rot_scales_flat[s_base + down + Int32(1)].to(cutlass.Float32)
+                sd2 = rot_scales_flat[s_base + down + Int32(2)].to(cutlass.Float32)
+                sd3 = rot_scales_flat[s_base + down + Int32(3)].to(cutlass.Float32)
+                a0 = self._silu_f32(ig0) * iu0 * sd0
+                a1 = self._silu_f32(ig1) * iu1 * sd1
+                a2 = self._silu_f32(ig2) * iu2 * sd2
+                a3 = self._silu_f32(ig3) * iu3 * sd3
+                o0, o1, o2, o3 = self._had128_quad(a0, a1, a2, a3, lane)
+                out_base = row * isz + col0
+                activated_bf16_flat[out_base + Int32(0)] = self._cast_elem(o0)
+                activated_bf16_flat[out_base + Int32(1)] = self._cast_elem(o1)
+                activated_bf16_flat[out_base + Int32(2)] = self._cast_elem(o2)
+                activated_bf16_flat[out_base + Int32(3)] = self._cast_elem(o3)
+            unit += gw_stride
+
     @cute.jit
     def _zero_fc2_output(
         self,
@@ -4841,16 +5540,142 @@ def _collect_activation_amax_epilogue(
                 )
             expert_idx += grid_x
 
+    @cute.jit
+    def _had128_quad(
+        self,
+        v0: cutlass.Float32,
+        v1: cutlass.Float32,
+        v2: cutlass.Float32,
+        v3: cutlass.Float32,
+        lane: Int32,
+    ):
+        # Blockwise-128 Walsh-Hadamard (Sylvester/natural order) across ONE warp,
+        # bit-identical in structure to exllamav3 had_hf_r_128_inner: lane `t`
+        # owns the 4 consecutive elements [4t..4t+3], so element index within the
+        # 128-block is 4*lane + reg.  H_128 = H_4(reg, in-register) (x) H_32(lane,
+        # warp butterfly), then * r_scale = 1/sqrt(128).
+        s0 = v0 + v1
+        d0 = v0 - v1
+        s1 = v2 + v3
+        d1 = v2 - v3
+        h0 = s0 + s1
+        h1 = d0 + d1
+        h2 = s0 - s1
+        h3 = d0 - d1
+        for i in cutlass.range_constexpr(5):
+            st = 1 << i
+            p0 = cute.arch.shuffle_sync_bfly(h0, offset=st)
+            p1 = cute.arch.shuffle_sync_bfly(h1, offset=st)
+            p2 = cute.arch.shuffle_sync_bfly(h2, offset=st)
+            p3 = cute.arch.shuffle_sync_bfly(h3, offset=st)
+            if (lane & Int32(st)) != Int32(0):
+                h0 = p0 - h0
+                h1 = p1 - h1
+                h2 = p2 - h2
+                h3 = p3 - h3
+            else:
+                h0 = p0 + h0
+                h1 = p1 + h1
+                h2 = p2 + h2
+                h3 = p3 + h3
+        rs = cutlass.Float32(0.088388347648)  # scale / sqrt(128), scale = 1
+        return h0 * rs, h1 * rs, h2 * rs, h3 * rs
+
+    @cute.jit
+    def _silu_f32(self, x: cutlass.Float32) -> cutlass.Float32:
+        if cutlass.const_expr(self.fast_math):
+            e = cute.math.exp(-x, fastmath=True)
+        else:
+            e = cute.math.exp(-x, fastmath=False)
+        return x / (cutlass.Float32(1.0) + e)
+
     @cute.jit
     def _run_activation(
         self,
         fc1_bf16_flat: cute.Tensor,
         activated_bf16_flat: cute.Tensor,
+        rot_scales_flat: cute.Tensor,
         tid: Int32,
         cta: Int32,
         grid_x: Int32,
         active_m: cutlass.Int32,
     ):
+        if cutlass.const_expr(self.intermediate_rotation):
+            # Rotation-aware epilogue (trellis3_t256 tail).  Warp-cooperative over
+            # (routed-row, 128-block) units; each warp owns one 128-wide block of
+            # a row's intermediate.  Per row r the FC1 output is [gate(I) | up(I)]
+            # (fc1_cols = 2I) and rot_scales_flat[r] = [svh_gate(I)|svh_up(I)|
+            # suh_down(I)] (3I fp16).  Applies:
+            #   ig2 = had128(gate) * svh_gate ;  iu2 = had128(up) * svh_up
+            #   ia  = silu(ig2) * iu2
+            #   ia2 = had128(ia * suh_down)           -> activated[r, block]
+            lane = tid & Int32(31)
+            warp_in_cta = tid >> Int32(5)
+            warps_per_cta = Int32(self.cta_threads // 32)
+            nblk = Int32(self.intermediate_size // 128)
+            fc1_cols = Int32(self.fc1_cols)
+            isz = Int32(self.intermediate_size)
+            rot_row = Int32(3 * self.intermediate_size)
+            gwarp = cta * warps_per_cta + warp_in_cta
+            gw_stride = grid_x * warps_per_cta
+            total_units = active_m * Int32(self.top_k) * nblk
+            e = lane * Int32(4)
+            unit = gwarp
+            while unit < total_units:
+                row = unit // nblk
+                blk = unit - row * nblk
+                col0 = blk * Int32(128) + e
+                g_base = row * fc1_cols + col0
+                u_base = g_base + isz
+                s_base = row * rot_row + col0
+                # load 4 gate + 4 up
+                g0 = fc1_bf16_flat[g_base + Int32(0)].to(cutlass.Float32)
+                g1 = fc1_bf16_flat[g_base + Int32(1)].to(cutlass.Float32)
+                g2 = fc1_bf16_flat[g_base + Int32(2)].to(cutlass.Float32)
+                g3 = fc1_bf16_flat[g_base + Int32(3)].to(cutlass.Float32)
+                u0 = fc1_bf16_flat[u_base + Int32(0)].to(cutlass.Float32)
+                u1 = fc1_bf16_flat[u_base + Int32(1)].to(cutlass.Float32)
+                u2 = fc1_bf16_flat[u_base + Int32(2)].to(cutlass.Float32)
+                u3 = fc1_bf16_flat[u_base + Int32(3)].to(cutlass.Float32)
+                # blockwise-128 Hadamard on gate and up (pre-silu, output side)
+                gh0, gh1, gh2, gh3 = self._had128_quad(g0, g1, g2, g3, lane)
+                uh0, uh1, uh2, uh3 = self._had128_quad(u0, u1, u2, u3, lane)
+                # post-scale svh_gate / svh_up
+                svg0 = rot_scales_flat[s_base + Int32(0)].to(cutlass.Float32)
+                svg1 = rot_scales_flat[s_base + Int32(1)].to(cutlass.Float32)
+                svg2 = rot_scales_flat[s_base + Int32(2)].to(cutlass.Float32)
+                svg3 = rot_scales_flat[s_base + Int32(3)].to(cutlass.Float32)
+                svu0 = rot_scales_flat[s_base + isz + Int32(0)].to(cutlass.Float32)
+                svu1 = rot_scales_flat[s_base + isz + Int32(1)].to(cutlass.Float32)
+                svu2 = rot_scales_flat[s_base + isz + Int32(2)].to(cutlass.Float32)
+                svu3 = rot_scales_flat[s_base + isz + Int32(3)].to(cutlass.Float32)
+                ig2_0 = gh0 * svg0
+                ig2_1 = gh1 * svg1
+                ig2_2 = gh2 * svg2
+                ig2_3 = gh3 * svg3
+                iu2_0 = uh0 * svu0
+                iu2_1 = uh1 * svu1
+                iu2_2 = uh2 * svu2
+                iu2_3 = uh3 * svu3
+                # silu(gate) * up, then pre-scale suh_down
+                d2 = isz + isz
+                sd0 = rot_scales_flat[s_base + d2 + Int32(0)].to(cutlass.Float32)
+                sd1 = rot_scales_flat[s_base + d2 + Int32(1)].to(cutlass.Float32)
+                sd2 = rot_scales_flat[s_base + d2 + Int32(2)].to(cutlass.Float32)
+                sd3 = rot_scales_flat[s_base + d2 + Int32(3)].to(cutlass.Float32)
+                a0 = self._silu_f32(ig2_0) * iu2_0 * sd0
+                a1 = self._silu_f32(ig2_1) * iu2_1 * sd1
+                a2 = self._silu_f32(ig2_2) * iu2_2 * sd2
+                a3 = self._silu_f32(ig2_3) * iu2_3 * sd3
+                # blockwise-128 Hadamard on the activation (post-silu, input side)
+                o0, o1, o2, o3 = self._had128_quad(a0, a1, a2, a3, lane)
+                out_base = row * isz + col0
+                activated_bf16_flat[out_base + Int32(0)] = self._cast_elem(o0)
+                activated_bf16_flat[out_base + Int32(1)] = self._cast_elem(o1)
+                activated_bf16_flat[out_base + Int32(2)] = self._cast_elem(o2)
+                activated_bf16_flat[out_base + Int32(3)] = self._cast_elem(o3)
+                unit += gw_stride
+            return
         idx = cta * Int32(self.cta_threads) + tid
         stride = grid_x * Int32(self.cta_threads)
         total = active_m * Int32(self.top_k * self.intermediate_size)
@@ -4926,6 +5751,10 @@ def __init__(
                 raise ValueError(
                     f"hybrid W4A16 {name} is incompatible with activation amax"
                 )
+            if moe.intermediate_rotation or moe.dual_a or moe.full_rotation:
+                raise ValueError(
+                    f"hybrid W4A16 {name} forbids trellis rotation features"
+                )
             if moe.zero_fc2_output:
                 raise ValueError(f"hybrid W4A16 {name} forbids zero_fc2_output")
             if not moe.tc_zero_output:
@@ -5057,6 +5886,7 @@ def _emit_route_map_tile(
                     if local_expert < Int32(self.tier0.num_experts):
                         if cutlass.const_expr(is_fc1):
                             self.tier0.fc1._run_tile(
+                                a_bf16_flat,
                                 a_bf16_flat,
                                 t0_b_i32_flat,
                                 c_bf16_flat,
@@ -5080,6 +5910,7 @@ def _emit_route_map_tile(
                             )
                         else:
                             self.tier0.fc2._run_tile(
+                                a_bf16_flat,
                                 a_bf16_flat,
                                 t0_b_i32_flat,
                                 c_bf16_flat,
@@ -5106,6 +5937,7 @@ def _emit_route_map_tile(
                         if local_expert < Int32(self.tier1.num_experts):
                             if cutlass.const_expr(is_fc1):
                                 self.tier1.fc1._run_tile(
+                                    a_bf16_flat,
                                     a_bf16_flat,
                                     t1_b_i32_flat,
                                     c_bf16_flat,
@@ -5129,6 +5961,7 @@ def _emit_route_map_tile(
                                 )
                             else:
                                 self.tier1.fc2._run_tile(
+                                    a_bf16_flat,
                                     a_bf16_flat,
                                     t1_b_i32_flat,
                                     c_bf16_flat,
@@ -5313,6 +6146,8 @@ class Storage:
         # route/amax parameters receive placeholder tensors that the direct
         # route + no-amax const_expr configuration never reads.
         self.tier0._moe_body(
+            a_bf16_flat,
+            a_bf16_flat,
             a_bf16_flat,
             t0_w13_i32_flat,
             t0_w2_i32_flat,
@@ -5332,6 +6167,9 @@ class Storage:
             fc1_c_tmp_f32_flat,
             fc2_c_tmp_f32_flat,
             locks_i32_flat,
+            global_topk_ids_i32_flat,
+            global_topk_ids_i32_flat,
+            global_topk_ids_i32_flat,
             smem_base,
             tid,
             cta,
@@ -5481,7 +6319,17 @@ def kernel(
 
 
 class W4A16TopKSumKernel:
-    def __init__(self, *, topk: int, hidden_size: int, element_dtype: str = "bf16"):
+    def __init__(
+        self,
+        *,
+        topk: int,
+        hidden_size: int,
+        element_dtype: str = "bf16",
+        full_rotation: bool = False,
+        num_experts: int = 0,
+        route_num_experts: int = 0,
+        use_expert_map: bool = False,
+    ):
         if element_dtype not in {"bf16", "fp16"}:
             raise ValueError(f"unsupported element_dtype {element_dtype!r}")
         if topk <= 0 or hidden_size <= 0:
@@ -5490,6 +6338,19 @@ def __init__(self, *, topk: int, hidden_size: int, element_dtype: str = "bf16"):
         self.hidden_size = int(hidden_size)
         self.element_dtype = element_dtype
         self.is_fp16 = element_dtype == "fp16"
+        self.full_rotation = bool(full_rotation)
+        self.num_experts = int(num_experts)
+        self.route_num_experts = int(route_num_experts)
+        self.use_expert_map = bool(use_expert_map)
+        if self.full_rotation:
+            if self.element_dtype != "fp16":
+                raise ValueError("full-rotation top-k sum requires fp16 route values")
+            if self.hidden_size % 128 != 0:
+                raise ValueError("full-rotation top-k sum requires hidden_size % 128 == 0")
+            if self.num_experts <= 0:
+                raise ValueError("full-rotation top-k sum requires num_experts > 0")
+            if self.use_expert_map and self.route_num_experts <= 0:
+                raise ValueError("expert-map top-k sum requires route_num_experts > 0")
         self.cta_threads = 256
 
     @cute.jit
@@ -5503,6 +6364,10 @@ def __call__(
         self,
         fc2_ptr: cute.Pointer,
         output_ptr: cute.Pointer,
+        topk_weights_ptr: cute.Pointer,
+        route_expert_ids_ptr: cute.Pointer,
+        expert_map_ptr: cute.Pointer,
+        svh_ptr: cute.Pointer,
         active_m: cutlass.Int32,
         stream: cuda.CUstream,
     ):
@@ -5516,9 +6381,39 @@ def __call__(
             output_ptr,
             layout=cute.make_layout((active_m * Int32(self.hidden_size),), stride=(1,)),
         )
-        total = active_m * Int32(self.hidden_size)
-        grid = (_covering_count(total, self.cta_threads), 1, 1)
-        self.kernel(fc2_flat, output_flat, active_m).launch(
+        topk_weights_flat = cute.make_tensor(
+            topk_weights_ptr,
+            layout=cute.make_layout((active_m * Int32(self.topk),), stride=(1,)),
+        )
+        route_expert_ids_flat = cute.make_tensor(
+            route_expert_ids_ptr,
+            layout=cute.make_layout((active_m * Int32(self.topk),), stride=(1,)),
+        )
+        expert_map_flat = cute.make_tensor(
+            expert_map_ptr,
+            layout=cute.make_layout((max(self.route_num_experts, 1),), stride=(1,)),
+        )
+        svh_flat = cute.make_tensor(
+            svh_ptr,
+            layout=cute.make_layout(
+                (max(self.num_experts * self.hidden_size, 1),), stride=(1,)
+            ),
+        )
+        if cutlass.const_expr(self.full_rotation):
+            total = active_m * Int32(self.hidden_size // 128)
+            grid = (_covering_count(total, self.cta_threads // 32), 1, 1)
+        else:
+            total = active_m * Int32(self.hidden_size)
+            grid = (_covering_count(total, self.cta_threads), 1, 1)
+        self.kernel(
+            fc2_flat,
+            output_flat,
+            topk_weights_flat,
+            route_expert_ids_flat,
+            expert_map_flat,
+            svh_flat,
+            active_m,
+        ).launch(
             grid=grid,
             block=[self.cta_threads, 1, 1],
             stream=stream,
@@ -5529,10 +6424,67 @@ def kernel(
         self,
         fc2_flat: cute.Tensor,
         output_flat: cute.Tensor,
+        topk_weights_flat: cute.Tensor,
+        route_expert_ids_flat: cute.Tensor,
+        expert_map_flat: cute.Tensor,
+        svh_flat: cute.Tensor,
         active_m: cutlass.Int32,
     ):
         tidx, _, _ = cute.arch.thread_idx()
         bidx, _, _ = cute.arch.block_idx()
+        if cutlass.const_expr(self.full_rotation):
+            tid = Int32(tidx)
+            lane = tid & Int32(31)
+            warp = tid >> Int32(5)
+            unit = Int32(bidx) * Int32(self.cta_threads // 32) + warp
+            nblk = Int32(self.hidden_size // 128)
+            total_units = active_m * nblk
+            if unit < total_units:
+                token = unit // nblk
+                blk = unit - token * nblk
+                col0 = blk * Int32(128) + lane * Int32(4)
+                acc0 = cutlass.Float32(0.0)
+                acc1 = cutlass.Float32(0.0)
+                acc2 = cutlass.Float32(0.0)
+                acc3 = cutlass.Float32(0.0)
+                for route in cutlass.range_constexpr(self.topk):
+                    row = token * Int32(self.topk) + Int32(route)
+                    raw_expert = route_expert_ids_flat[row].to(Int32)
+                    expert = raw_expert
+                    if cutlass.const_expr(self.use_expert_map):
+                        expert = Int32(-1)
+                        if (
+                            raw_expert >= Int32(0)
+                            and raw_expert < Int32(self.route_num_experts)
+                        ):
+                            expert = expert_map_flat[raw_expert].to(Int32)
+                    if expert >= Int32(0) and expert < Int32(self.num_experts):
+                        base = row * Int32(self.hidden_size) + col0
+                        v0 = fc2_flat[base + Int32(0)].to(cutlass.Float32)
+                        v1 = fc2_flat[base + Int32(1)].to(cutlass.Float32)
+                        v2 = fc2_flat[base + Int32(2)].to(cutlass.Float32)
+                        v3 = fc2_flat[base + Int32(3)].to(cutlass.Float32)
+                        h0, h1, h2, h3 = self._had128_quad(
+                            v0, v1, v2, v3, lane
+                        )
+                        sbase = expert * Int32(self.hidden_size) + col0
+                        s0 = svh_flat[sbase + Int32(0)].to(cutlass.Float32)
+                        s1 = svh_flat[sbase + Int32(1)].to(cutlass.Float32)
+                        s2 = svh_flat[sbase + Int32(2)].to(cutlass.Float32)
+                        s3 = svh_flat[sbase + Int32(3)].to(cutlass.Float32)
+                        weight = topk_weights_flat[row].to(cutlass.Float32)
+                        # Deliberate fp32-C ordering: raw route H128, fp32(svh),
+                        # then router weight, with no intermediate fp16 round.
+                        acc0 += h0 * s0 * weight
+                        acc1 += h1 * s1 * weight
+                        acc2 += h2 * s2 * weight
+                        acc3 += h3 * s3 * weight
+                out_base = token * Int32(self.hidden_size) + col0
+                output_flat[out_base + Int32(0)] = acc0
+                output_flat[out_base + Int32(1)] = acc1
+                output_flat[out_base + Int32(2)] = acc2
+                output_flat[out_base + Int32(3)] = acc3
+            return
         idx = Int32(bidx) * Int32(self.cta_threads) + Int32(tidx)
         total = active_m * Int32(self.hidden_size)
         if idx < total:
@@ -5547,6 +6499,42 @@ def kernel(
                 acc += _materialize_w4a16_topk_route_f32(route_value)
             output_flat[idx] = self._cast_elem(acc)
 
+    @cute.jit
+    def _had128_quad(
+        self,
+        v0: cutlass.Float32,
+        v1: cutlass.Float32,
+        v2: cutlass.Float32,
+        v3: cutlass.Float32,
+        lane: Int32,
+    ):
+        s0 = v0 + v1
+        d0 = v0 - v1
+        s1 = v2 + v3
+        d1 = v2 - v3
+        h0 = s0 + s1
+        h1 = d0 + d1
+        h2 = s0 - s1
+        h3 = d0 - d1
+        for i in cutlass.range_constexpr(5):
+            st = 1 << i
+            p0 = cute.arch.shuffle_sync_bfly(h0, offset=st)
+            p1 = cute.arch.shuffle_sync_bfly(h1, offset=st)
+            p2 = cute.arch.shuffle_sync_bfly(h2, offset=st)
+            p3 = cute.arch.shuffle_sync_bfly(h3, offset=st)
+            if (lane & Int32(st)) != Int32(0):
+                h0 = p0 - h0
+                h1 = p1 - h1
+                h2 = p2 - h2
+                h3 = p3 - h3
+            else:
+                h0 = p0 + h0
+                h1 = p1 + h1
+                h2 = p2 + h2
+                h3 = p3 + h3
+        rs = cutlass.Float32(0.088388347648)
+        return h0 * rs, h1 * rs, h2 * rs, h3 * rs
+
 
 _CACHE: dict[tuple, W4A16GemmCompileResult] = {}
 _FUSED_CACHE: dict[tuple, W4A16FusedMoeCompileResult] = {}
@@ -5813,7 +6801,11 @@ def compile_w4a16_gemm(
     moe_block_size: int,
     max_m_blocks: int,
     element_dtype: str = "bf16",
+    weight_layout: str = "packed",
     scale_format: str = "e4m3_k16",
+    w13_layout: str = "packed",
+    trellis_bits: int = 3,
+    dense_route_fast_path: bool = False,
 ) -> W4A16GemmCompileResult:
     scale_format = _normalize_scale_format(scale_format)
     cutlass_dtype = _cutlass_element_dtype(element_dtype)
@@ -5833,7 +6825,12 @@ def compile_w4a16_gemm(
         moe_block_size=moe_block_size,
         max_m_blocks=max_m_blocks,
         element_dtype=element_dtype,
+        weight_layout=weight_layout,
         scale_format=scale_format,
+        w13_layout=w13_layout,
+        trellis_bits=trellis_bits,
+        dense_route_fast_path=bool(dense_route_fast_path),
+        schedule_whole_tiles=weight_layout == "trellis3_t256",
     )
     cache_key = (
         "w4a16_gemm",
@@ -5851,21 +6848,19 @@ def compile_w4a16_gemm(
     compile_size_m = _fake_m_for_specialization(size_m)
     compile_route_blocks = 1
     compile_route_slots = compile_route_blocks * int(moe_block_size)
-    a_fake = cute.runtime.make_fake_compact_tensor(
-        cutlass_dtype,
-        (compile_size_m * size_k,),
-        assumed_align=16,
-    )
-    b_fake = cute.runtime.make_fake_compact_tensor(
+    a_fake = make_ptr(cutlass_dtype, 16, cute.AddressSpace.gmem, assumed_align=16)
+    if weight_layout == "trellis3_t256":
+        b_fake_elements = (
+            num_experts * (size_k // 16) * (size_n // 16) * (8 * int(trellis_bits))
+        )
+    else:
+        b_fake_elements = num_experts * (size_k // 16) * (size_n // 16 * 32)
+    b_fake = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,
-        (num_experts * (size_k // 16) * (size_n // 16 * 32),),
-        assumed_align=16,
-    )
-    c_fake = cute.runtime.make_fake_compact_tensor(
-        cutlass_dtype,
-        (compile_size_m * top_k * size_n,),
+        (b_fake_elements,),
         assumed_align=16,
     )
+    c_fake = make_ptr(cutlass_dtype, 16, cute.AddressSpace.gmem, assumed_align=16)
     scales_fake = cute.runtime.make_fake_compact_tensor(
         cutlass.Int32,
         (
@@ -5925,6 +6920,7 @@ def compile_w4a16_gemm(
     compiled = sparkinfer_compile(
         kernel,
         a_fake,
+        a_fake,
         b_fake,
         c_fake,
         scales_fake,
@@ -5940,7 +6936,7 @@ def compile_w4a16_gemm(
         current_cuda_stream(),
         compile_spec=KernelCompileSpec.from_key(
             "moe.w4a16.gemm",
-            1,
+            3,
             cache_key,
         ),
     )
@@ -5951,7 +6947,11 @@ def compile_w4a16_gemm(
         moe_block_size=moe_block_size,
         max_m_blocks=max_m_blocks,
         blocks_per_sm=kernel.blocks_per_sm,
+        weight_layout=weight_layout,
         scale_format=scale_format,
+        w13_layout=w13_layout,
+        dense_route_fast_path=bool(dense_route_fast_path),
+        trellis_bits=int(trellis_bits),
     )
     _CACHE[cache_key] = result
     return result
@@ -5979,12 +6979,21 @@ def compile_w4a16_fused_moe(
     weight_layout: str = "packed",
     scale_format: str = "e4m3_k16",
     w13_layout: str = "w13",
+    trellis_bits: int = 3,
     direct_topk_routes: bool = False,
     tc_decode_fused_sum: bool = False,
     collect_activation_amax: bool = False,
     force_tile_config: tuple[int, int, int, int] | None = None,
+    intermediate_rotation: bool = False,
+    full_rotation: bool = False,
+    rotation_input_dtype: str | None = None,
 ) -> W4A16FusedMoeCompileResult:
     scale_format = _normalize_scale_format(scale_format)
+    intermediate_rotation = bool(intermediate_rotation)
+    full_rotation = bool(full_rotation)
+    rotation_input_dtype = (
+        element_dtype if rotation_input_dtype is None else str(rotation_input_dtype)
+    )
     cutlass_dtype = _cutlass_element_dtype(element_dtype)
     device = int(torch.cuda.current_device()) if torch.cuda.is_available() else None
     activation = normalize_moe_activation(activation)
@@ -5997,9 +7006,31 @@ def compile_w4a16_fused_moe(
     )
     if weight_layout not in _WEIGHT_LAYOUTS:
         raise ValueError(f"unsupported W4A16 weight_layout {weight_layout!r}")
+    trellis_bits = int(trellis_bits)
+    if weight_layout == "trellis3_t256":
+        if trellis_bits not in _TRELLIS256_BITS:
+            raise ValueError(
+                f"trellis3_t256 bits must be one of {_TRELLIS256_BITS}, got {trellis_bits}"
+            )
+    elif trellis_bits != 3:
+        raise ValueError("trellis_bits is only valid for trellis3_t256 weights")
+    # Existing 3-bpw scheduling was conservatively planned as 4 bpw. Keep that
+    # grid contract stable for D6; widen only the 5/6-bpw specializations.
+    weight_bits = max(4, trellis_bits) if weight_layout == "trellis3_t256" else 4
+    # GATE 5: the PRODUCTION 256-weight-tile fused-megakernel B-staging is now
+    # wired (per-warp native [K/16,N/16,8*bits u32] tile staging + the per-lane
+    # bitrate-specialized read) and ADMITTED at 3 bpw against a full-GEMM
+    # numeric oracle (megakernel FC1->silu->FC2 vs a torch reference over tiles;
+    # tests/test_trellis256_fullgemm_oracle.py, max rel-err <= 2^-6).  The GATE-4
+    # NotImplementedError gate is therefore lifted.  See GATE5_RESULT.md.
     if weight_layout == "modelopt":
         if w13_layout not in _MODEL_OPT_W13_LAYOUTS:
             raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}")
+    elif weight_layout == "trellis3_t256":
+        if w13_layout not in _TRELLIS256_W13_LAYOUTS:
+            raise ValueError(
+                f"unsupported trellis3_t256 w13_layout {w13_layout!r}"
+            )
     else:
         w13_layout = "packed"
     direct_topk_routes = bool(direct_topk_routes)
@@ -6009,6 +7040,34 @@ def compile_w4a16_fused_moe(
         raise ValueError(
             "W4A16 activation amax collection requires the route-packed fused path"
         )
+    if collect_activation_amax and intermediate_rotation:
+        raise ValueError(
+            "W4A16 activation amax collection is incompatible with intermediate rotation"
+        )
+    if full_rotation:
+        if not intermediate_rotation:
+            raise ValueError("full_rotation requires intermediate_rotation")
+        if weight_layout != "trellis3_t256":
+            raise ValueError("full_rotation is only supported for trellis3_t256")
+        if element_dtype != "fp16":
+            raise ValueError("full_rotation requires element_dtype='fp16'")
+        if rotation_input_dtype not in {"bf16", "fp16"}:
+            raise ValueError(
+                "rotation_input_dtype must be 'bf16' or 'fp16' for full_rotation"
+            )
+        if direct_topk_routes or tc_decode_fused_sum:
+            raise ValueError(
+                "full_rotation requires route packing and cannot use TC decode"
+            )
+        if apply_router_weight_on_input:
+            raise ValueError(
+                "full_rotation requires apply_router_weight_on_input=False"
+            )
+    if collect_activation_amax and weight_layout == "trellis3_t256":
+        raise NotImplementedError(
+            "trellis3_t256 activation-amax collection is not exposed through the "
+            "registered launch ABI; refusing to compile a bitrate-ambiguous kernel"
+        )
     # The TC-decode path validates M in {1,2,4,8} itself and uses direct-topk
     # routing for the whole {1,2,4,8} range, so it lifts the default decode cap.
     direct_topk_m_cap = (
@@ -6044,6 +7103,7 @@ def compile_w4a16_fused_moe(
         max_shared_mem=max_shared_mem,
         scale_format=scale_format,
         weight_layout=weight_layout,
+        weight_bits=weight_bits,
         allow_logical_tail=allow_native_logical_tail,
     )
     fc2_tile_k, fc2_tile_n, fc2_cta_threads, _ = _select_tile_config(
@@ -6056,6 +7116,7 @@ def compile_w4a16_fused_moe(
         max_shared_mem=max_shared_mem,
         scale_format=scale_format,
         weight_layout=weight_layout,
+        weight_bits=weight_bits,
         allow_logical_tail=allow_native_logical_tail,
     )
     if fc1_cta_threads != fc2_cta_threads:
@@ -6071,6 +7132,7 @@ def compile_w4a16_fused_moe(
             required_cta_threads=common_cta_threads,
             scale_format=scale_format,
             weight_layout=weight_layout,
+            weight_bits=weight_bits,
             allow_logical_tail=allow_native_logical_tail,
         )
         fc2_tile_k, fc2_tile_n, fc2_cta_threads, _ = _select_tile_config(
@@ -6084,6 +7146,7 @@ def compile_w4a16_fused_moe(
             required_cta_threads=common_cta_threads,
             scale_format=scale_format,
             weight_layout=weight_layout,
+            weight_bits=weight_bits,
             allow_logical_tail=allow_native_logical_tail,
         )
         if fc1_cta_threads != fc2_cta_threads:
@@ -6133,6 +7196,7 @@ def compile_w4a16_fused_moe(
             max_shared_mem=int(max_shared_mem) - 512,
             scale_format=scale_format,
             weight_layout=weight_layout,
+            weight_bits=weight_bits,
         ):
             fc1_tile_n = 256
             fc1_tile_k = wide_fc1_tile_k
@@ -6164,6 +7228,7 @@ def compile_w4a16_fused_moe(
             max_shared_mem=int(max_shared_mem) - 512,
             scale_format=scale_format,
             weight_layout=weight_layout,
+            weight_bits=weight_bits,
         ):
             fc2_tile_n = 256
             fc2_tile_k = wide_fc2_tile_k
@@ -6209,6 +7274,7 @@ def compile_w4a16_fused_moe(
             tile_k=ultra_fc2_tile_k,
             scale_format=scale_format,
             weight_layout=weight_layout,
+            weight_bits=weight_bits,
         )
         if (
             int(intermediate_size) % ultra_fc2_tile_k == 0
@@ -6248,6 +7314,7 @@ def compile_w4a16_fused_moe(
                 max_shared_mem=int(max_shared_mem) - 512,
                 scale_format=scale_format,
                 weight_layout=weight_layout,
+                weight_bits=weight_bits,
                 allow_logical_tail=allow_native_logical_tail,
             ):
                 raise ValueError(
@@ -6279,9 +7346,13 @@ def compile_w4a16_fused_moe(
         weight_layout=weight_layout,
         scale_format=scale_format,
         w13_layout=w13_layout,
+        trellis_bits=trellis_bits,
         direct_topk_routes=direct_topk_routes,
         tc_decode_fused_sum=tc_decode_fused_sum,
         collect_activation_amax=collect_activation_amax,
+        intermediate_rotation=intermediate_rotation,
+        full_rotation=full_rotation,
+        rotation_input_dtype=rotation_input_dtype,
     )
     cache_key = (
         "w4a16_fused_moe",
@@ -6339,6 +7410,12 @@ def compile_w4a16_fused_moe(
         compile_routed_rows if direct_topk_routes else compile_route_slots
     )
     a_fake = make_ptr(cutlass_dtype, 16, cute.AddressSpace.gmem, assumed_align=16)
+    rotation_input_fake = make_ptr(
+        _cutlass_element_dtype(rotation_input_dtype),
+        16,
+        cute.AddressSpace.gmem,
+        assumed_align=16,
+    )
     if weight_layout == "modelopt":
         w13_fake = cute.runtime.make_fake_compact_tensor(
             cutlass.Uint8,
@@ -6350,6 +7427,27 @@ def compile_w4a16_fused_moe(
             (num_experts * hidden_size * (intermediate_size // 2),),
             assumed_align=16,
         )
+    elif weight_layout == "trellis3_t256":
+        w13_fake = cute.runtime.make_fake_compact_tensor(
+            cutlass.Int32,
+            (
+                num_experts
+                * (hidden_size // 16)
+                * (fc1_cols // 16)
+                * (8 * trellis_bits),
+            ),
+            assumed_align=16,
+        )
+        w2_fake = cute.runtime.make_fake_compact_tensor(
+            cutlass.Int32,
+            (
+                num_experts
+                * (intermediate_size // 16)
+                * (hidden_size // 16)
+                * (8 * trellis_bits),
+            ),
+            assumed_align=16,
+        )
     elif weight_layout == "nf3_2p1":
         # NF3: int32, 3 words per 32-code unit; (size_n // 2) units per K16 row.
         w13_fake = cute.runtime.make_fake_compact_tensor(
@@ -6470,6 +7568,24 @@ def compile_w4a16_fused_moe(
         (4 * 256 + 2,),
         assumed_align=16,
     )
+    # Legacy mode supplies per-route [R,3I]; full rotation indexes the compact
+    # persistent [E,3I] table by the packed block's expert id.
+    rot_rows = num_experts if full_rotation else compile_routed_rows
+    rot_scales_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Float16,
+        (max(rot_rows * 3 * intermediate_size, 1),),
+        assumed_align=16,
+    )
+    suh_gate_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Float16,
+        (max(num_experts * hidden_size, 1),),
+        assumed_align=16,
+    )
+    suh_up_fake = cute.runtime.make_fake_compact_tensor(
+        cutlass.Float16,
+        (max(num_experts * hidden_size, 1),),
+        assumed_align=16,
+    )
 
     raise_if_kernel_resolution_frozen(
         "cute.compile", target=kernel, cache_key=cache_key
@@ -6477,6 +7593,8 @@ def compile_w4a16_fused_moe(
     compiled = sparkinfer_compile(
         kernel,
         a_fake,
+        a_fake,
+        rotation_input_fake,
         w13_fake,
         w2_fake,
         fc1_fake,
@@ -6495,12 +7613,15 @@ def compile_w4a16_fused_moe(
         fc1_c_tmp_fake,
         fc2_c_tmp_fake,
         locks_fake,
+        rot_scales_fake,
+        suh_gate_fake,
+        suh_up_fake,
         1,
         1,
         current_cuda_stream(),
         compile_spec=KernelCompileSpec.from_key(
             "moe.w4a16.fused_moe",
-            1,
+            4,
             cache_key,
         ),
         dsl_compile_options=OptLevel(2),
@@ -6533,6 +7654,12 @@ def compile_w4a16_fused_moe(
         scale_format=scale_format,
         tc_decode_fused_sum=bool(tc_decode_fused_sum),
         collect_activation_amax=collect_activation_amax,
+        schedule_whole_tiles=kernel.schedule_whole_tiles,
+        intermediate_rotation=intermediate_rotation,
+        dual_a=kernel.dual_a,
+        trellis_bits=trellis_bits,
+        full_rotation=full_rotation,
+        rotation_input_dtype=rotation_input_dtype,
     )
     _FUSED_CACHE[cache_key] = result
     return result
@@ -6981,19 +8108,60 @@ def compile_w4a16_topk_sum(
     topk: int,
     hidden_size: int,
     element_dtype: str = "bf16",
+    full_rotation: bool = False,
+    num_experts: int = 0,
+    route_num_experts: int = 0,
+    route_ids_dtype: torch.dtype = torch.int32,
+    use_expert_map: bool = False,
 ) -> W4A16TopKSumCompileResult:
     cutlass_dtype = _cutlass_element_dtype(element_dtype)
-    cache_key = ("w4a16_topk_sum", element_dtype, topk, hidden_size)
+    if route_ids_dtype not in (torch.int32, torch.int64):
+        raise TypeError("top-k route expert ids must be int32 or int64")
+    route_cutlass_dtype = (
+        cutlass.Int32 if route_ids_dtype == torch.int32 else cutlass.Int64
+    )
+    cache_key = (
+        "w4a16_topk_sum",
+        element_dtype,
+        topk,
+        hidden_size,
+        bool(full_rotation),
+        int(num_experts),
+        int(route_num_experts),
+        str(route_ids_dtype),
+        bool(use_expert_map),
+    )
     cached = _SUM_CACHE.get(cache_key)
     if cached is not None:
         return cached
 
     fc2_fake = make_ptr(cutlass_dtype, 16, cute.AddressSpace.gmem, assumed_align=16)
-    output_fake = make_ptr(cutlass_dtype, 16, cute.AddressSpace.gmem, assumed_align=16)
+    output_dtype = cutlass.Float32 if full_rotation else cutlass_dtype
+    output_fake = make_ptr(output_dtype, 16, cute.AddressSpace.gmem, assumed_align=16)
+    topk_weights_fake = make_ptr(
+        cutlass.Float32, 4, cute.AddressSpace.gmem, assumed_align=4
+    )
+    route_ids_align = 4 if route_ids_dtype == torch.int32 else 8
+    route_ids_fake = make_ptr(
+        route_cutlass_dtype,
+        route_ids_align,
+        cute.AddressSpace.gmem,
+        assumed_align=route_ids_align,
+    )
+    expert_map_fake = make_ptr(
+        cutlass.Int32, 4, cute.AddressSpace.gmem, assumed_align=4
+    )
+    svh_fake = make_ptr(
+        cutlass.Float16, 16, cute.AddressSpace.gmem, assumed_align=16
+    )
     kernel = W4A16TopKSumKernel(
         topk=topk,
         hidden_size=hidden_size,
         element_dtype=element_dtype,
+        full_rotation=full_rotation,
+        num_experts=num_experts,
+        route_num_experts=route_num_experts,
+        use_expert_map=use_expert_map,
     )
     raise_if_kernel_resolution_frozen(
         "cute.compile", target=kernel, cache_key=cache_key
@@ -7002,11 +8170,15 @@ def compile_w4a16_topk_sum(
         kernel,
         fc2_fake,
         output_fake,
+        topk_weights_fake,
+        route_ids_fake,
+        expert_map_fake,
+        svh_fake,
         1,
         current_cuda_stream(),
         compile_spec=KernelCompileSpec.from_key(
             "moe.w4a16.topk_sum",
-            1,
+            2,
             cache_key,
         ),
     )
@@ -7015,6 +8187,11 @@ def compile_w4a16_topk_sum(
         m=0,
         topk=topk,
         hidden_size=hidden_size,
+        full_rotation=bool(full_rotation),
+        num_experts=int(num_experts),
+        route_num_experts=int(route_num_experts),
+        route_ids_dtype=route_ids_dtype,
+        use_expert_map=bool(use_expert_map),
     )
     _SUM_CACHE[cache_key] = result
     return result
@@ -7190,6 +8367,20 @@ def _w4a16_small_m_direct_launch_fake(
     return None
 
 
+_ROT_SCALES_DUMMY: dict[object, torch.Tensor] = {}
+
+
+def _rot_scales_dummy(device: torch.device) -> torch.Tensor:
+    """Tiny fp16 placeholder for the rot_scales kernel slot when the
+    intermediate-rotation epilogue is disabled (never dereferenced by the
+    const_expr-gated kernel; keeps one compiled signature across all layouts)."""
+    t = _ROT_SCALES_DUMMY.get(device)
+    if t is None:
+        t = torch.zeros(1, dtype=torch.float16, device=device)
+        _ROT_SCALES_DUMMY[device] = t
+    return t
+
+
 def _w4a16_fused_moe_launch_flat(
     a_input: torch.Tensor,
     w13_arg: torch.Tensor,
@@ -7240,14 +8431,52 @@ def _w4a16_fused_moe_launch_flat(
     tc_decode_fused_sum: bool,
     collect_activation_amax: bool,
     stream_int: int,
+    rot_scales: torch.Tensor | None = None,
+    intermediate_rotation: bool = False,
+    a_input_up: torch.Tensor | None = None,
+    trellis_bits: int = 3,
+    full_rotation: bool = False,
+    rotation_input: torch.Tensor | None = None,
+    suh_gate_table: torch.Tensor | None = None,
+    suh_up_table: torch.Tensor | None = None,
 ) -> None:
     swiglu_limit = float(swiglu_limit_value) if has_swiglu_limit else None
     collect_activation_amax = bool(collect_activation_amax)
+    intermediate_rotation = bool(intermediate_rotation)
+    full_rotation = bool(full_rotation)
     if collect_activation_amax and activation_amax is None:
         raise ValueError("activation_amax is required for calibrated W4A16 launch")
     activation_amax_arg = (
         activation_amax.view(-1) if activation_amax is not None else w13_global_scale
     )
+    if intermediate_rotation:
+        if rot_scales is None:
+            raise ValueError("intermediate_rotation launch requires rot_scales")
+        rot_scales_arg = rot_scales.view(-1)
+    else:
+        rot_scales_arg = _rot_scales_dummy(w13_global_scale.device)
+    if full_rotation and a_input_up is None:
+        raise ValueError("full_rotation launch requires a distinct up A scratch")
+    if a_input_up is None:
+        a_input_up = a_input
+    if full_rotation and a_input_up.data_ptr() == a_input.data_ptr():
+        raise ValueError("full_rotation gate/up A scratches must not alias")
+    if full_rotation and rotation_input is None:
+        raise ValueError("full_rotation launch requires the raw rotation input")
+    if rotation_input is None:
+        rotation_input = a_input
+    if full_rotation:
+        if suh_gate_table is None or suh_up_table is None:
+            raise ValueError(
+                "full_rotation launch requires suh_gate_table and suh_up_table"
+            )
+        suh_gate_arg = suh_gate_table.reshape(-1)
+        suh_up_arg = suh_up_table.reshape(-1)
+        rotation_input_dtype = _normalize_element_dtype(rotation_input.dtype)
+    else:
+        suh_gate_arg = _rot_scales_dummy(w13_global_scale.device)
+        suh_up_arg = suh_gate_arg
+        rotation_input_dtype = element_dtype
     fused = compile_w4a16_fused_moe(
         size_m=size_m,
         hidden_size=hidden_size,
@@ -7269,6 +8498,7 @@ def _w4a16_fused_moe_launch_flat(
         weight_layout=weight_layout,
         scale_format=scale_format,
         w13_layout=w13_layout,
+        trellis_bits=trellis_bits,
         direct_topk_routes=bool(direct_topk_routes),
         tc_decode_fused_sum=bool(tc_decode_fused_sum),
         collect_activation_amax=collect_activation_amax,
@@ -7276,6 +8506,9 @@ def _w4a16_fused_moe_launch_flat(
         # its selected geometry so tile-specific packs (notably NF3) resolve the
         # identical cache entry instead of silently recompiling with auto tiles.
         force_tile_config=(fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n),
+        intermediate_rotation=intermediate_rotation,
+        full_rotation=full_rotation,
+        rotation_input_dtype=rotation_input_dtype,
     )
     fused.compiled(
         make_ptr(
@@ -7284,6 +8517,18 @@ def _w4a16_fused_moe_launch_flat(
             cute.AddressSpace.gmem,
             assumed_align=16,
         ),
+        make_ptr(
+            _cutlass_element_dtype(element_dtype),
+            a_input_up.data_ptr(),
+            cute.AddressSpace.gmem,
+            assumed_align=16,
+        ),
+        make_ptr(
+            _cutlass_element_dtype(rotation_input_dtype),
+            rotation_input.data_ptr(),
+            cute.AddressSpace.gmem,
+            assumed_align=16,
+        ),
         w13_arg,
         w2_arg,
         fc1_out,
@@ -7307,6 +8552,9 @@ def _w4a16_fused_moe_launch_flat(
         fc1_scratch,
         fc2_scratch,
         workspace,
+        rot_scales_arg,
+        suh_gate_arg,
+        suh_up_arg,
         m,
         _w4a16_fused_persistent_grid_x(
             fused=fused,
@@ -7711,13 +8959,41 @@ def _w4a16_topk_sum_launch_flat(
     hidden_size: int,
     element_dtype: str,
     stream_int: int,
+    *,
+    full_rotation: bool = False,
+    num_experts: int = 0,
+    topk_weights: torch.Tensor | None = None,
+    route_expert_ids: torch.Tensor | None = None,
+    expert_map: torch.Tensor | None = None,
+    svh_table: torch.Tensor | None = None,
 ) -> None:
+    full_rotation = bool(full_rotation)
+    route_ids_dtype = (
+        torch.int32 if route_expert_ids is None else route_expert_ids.dtype
+    )
+    route_num_experts = 0 if expert_map is None else int(expert_map.numel())
     sum_kernel = compile_w4a16_topk_sum(
         m=m,
         topk=topk,
         hidden_size=hidden_size,
         element_dtype=element_dtype,
+        full_rotation=full_rotation,
+        num_experts=num_experts,
+        route_num_experts=route_num_experts,
+        route_ids_dtype=route_ids_dtype,
+        use_expert_map=expert_map is not None,
+    )
+    dummy_addr = output.data_ptr()
+    weights_addr = dummy_addr if topk_weights is None else topk_weights.data_ptr()
+    route_ids_addr = (
+        dummy_addr if route_expert_ids is None else route_expert_ids.data_ptr()
     )
+    expert_map_addr = dummy_addr if expert_map is None else expert_map.data_ptr()
+    svh_addr = dummy_addr if svh_table is None else svh_table.data_ptr()
+    route_cutlass_dtype = (
+        cutlass.Int32 if route_ids_dtype == torch.int32 else cutlass.Int64
+    )
+    route_align = 4 if route_ids_dtype == torch.int32 else 8
     sum_kernel.compiled(
         make_ptr(
             _cutlass_element_dtype(element_dtype),
@@ -7726,11 +9002,35 @@ def _w4a16_topk_sum_launch_flat(
             assumed_align=16,
         ),
         make_ptr(
-            _cutlass_element_dtype(element_dtype),
+            cutlass.Float32 if full_rotation else _cutlass_element_dtype(element_dtype),
             output.data_ptr(),
             cute.AddressSpace.gmem,
             assumed_align=16,
         ),
+        make_ptr(
+            cutlass.Float32,
+            weights_addr,
+            cute.AddressSpace.gmem,
+            assumed_align=4,
+        ),
+        make_ptr(
+            route_cutlass_dtype,
+            route_ids_addr,
+            cute.AddressSpace.gmem,
+            assumed_align=route_align,
+        ),
+        make_ptr(
+            cutlass.Int32,
+            expert_map_addr,
+            cute.AddressSpace.gmem,
+            assumed_align=4,
+        ),
+        make_ptr(
+            cutlass.Float16,
+            svh_addr,
+            cute.AddressSpace.gmem,
+            assumed_align=16,
+        ),
         m,
         cuda.CUstream(stream_int),
     )
@@ -7884,23 +9184,54 @@ def _compile_w4a16_gemm_launch(
     moe_block_size: int,
     max_m_blocks: int,
     element_dtype: str,
-    packed_route_indices: torch.Tensor,
+    packed_route_indices: torch.Tensor | None,
     sms: int,
     max_shared_mem: int,
     device: torch.device,
     c_tmp: torch.Tensor | None = None,
+    weight_layout: str = "packed",
     scale_format: str = "e4m3_k16",
+    w13_layout: str = "packed",
+    trellis_bits: int = 3,
+    dense_route_fast_path: bool = False,
+    route_slots: int | None = None,
+    force_tile_config: tuple[int, int] | None = None,
 ) -> _W4A16GemmLaunch:
-    tile_k, tile_n, _, _ = _select_tile_config(
-        problem_m=size_m,
-        problem_n=size_n,
-        problem_k=size_k,
-        top_k=top_k,
-        moe_block_size=moe_block_size,
-        sms=sms,
-        max_shared_mem=max_shared_mem,
-        scale_format=scale_format,
+    planner_weight_bits = (
+        max(4, int(trellis_bits)) if weight_layout == "trellis3_t256" else 4
     )
+    if force_tile_config is None:
+        tile_k, tile_n, _, _ = _select_tile_config(
+            problem_m=size_m,
+            problem_n=size_n,
+            problem_k=size_k,
+            top_k=top_k,
+            moe_block_size=moe_block_size,
+            sms=sms,
+            max_shared_mem=max_shared_mem,
+            scale_format=scale_format,
+            weight_layout=weight_layout,
+            weight_bits=planner_weight_bits,
+        )
+    else:
+        tile_k, tile_n = (int(v) for v in force_tile_config)
+        cta_threads = tile_n * tile_k // 64
+        if not _candidate_tile_fits(
+            problem_n=size_n,
+            problem_k=size_k,
+            cta_m_blocks=_covering_count(moe_block_size, 16),
+            tile_n=tile_n,
+            tile_k=tile_k,
+            cta_threads=cta_threads,
+            max_shared_mem=max_shared_mem,
+            scale_format=scale_format,
+            weight_layout=weight_layout,
+            weight_bits=planner_weight_bits,
+        ):
+            raise ValueError(
+                "forced standalone W4A16 tile does not fit: "
+                f"tile_k/tile_n={tile_k}/{tile_n}, N/K={size_n}/{size_k}"
+            )
     kernel = compile_w4a16_gemm(
         size_m=size_m,
         size_n=size_n,
@@ -7913,12 +9244,22 @@ def _compile_w4a16_gemm_launch(
         moe_block_size=moe_block_size,
         max_m_blocks=max_m_blocks,
         element_dtype=element_dtype,
+        weight_layout=weight_layout,
         scale_format=scale_format,
+        w13_layout=w13_layout,
+        trellis_bits=trellis_bits,
+        dense_route_fast_path=bool(dense_route_fast_path),
     )
+    if route_slots is None:
+        if packed_route_indices is None:
+            raise ValueError(
+                "standalone W4A16 launch requires route_slots or packed_route_indices"
+            )
+        route_slots = int(packed_route_indices.numel())
     c_tmp = _get_c_tmp(
         packed_gemm_scratch_elements(
             size_n=size_n,
-            route_slots=int(packed_route_indices.numel()),
+            route_slots=int(route_slots),
             moe_block_size=moe_block_size,
             sms=sms,
         ),
@@ -7960,6 +9301,259 @@ def pack_topk_routes_by_expert(
     )
 
 
+def _trellis256_dense_tile_config(size_k: int, size_n: int) -> tuple[int, int]:
+    """Return an m-invariant t256 tile so dense row bits cannot drift with m."""
+    if int(size_k) % 64 != 0:
+        raise ValueError(f"trellis3_t256 dense K must be divisible by 64, got {size_k}")
+    if int(size_n) % 256 == 0:
+        return (64, 256)
+    if int(size_n) % 128 == 0:
+        # With dense M=64, (tile_k=64,tile_n=128) maps to the measured
+        # (cta_threads=128, cta_m/n/k_blocks=4/8/4) register specialization.
+        # The superficially symmetric 128x128 tile would need a missing
+        # 256-thread 4/8/8 register entry and fail before compilation.
+        return (64, 128)
+    raise ValueError(
+        "trellis3_t256 dense N must be divisible by 128 (or 256 for the wide tile), "
+        f"got N={size_n}"
+    )
+
+
+def _resolve_exl3_hadamard_128(hadamard_128):
+    if hadamard_128 is None:
+        try:
+            import exllamav3_ext
+        except ImportError as exc:
+            raise RuntimeError(
+                "run_trellis256_dense requires exllamav3_ext.had_r_128 for the "
+                "EXL3 input/output rotations"
+            ) from exc
+        hadamard_128 = exllamav3_ext.had_r_128
+    elif not callable(hadamard_128):
+        hadamard_128 = getattr(hadamard_128, "had_r_128", None)
+    if not callable(hadamard_128):
+        raise TypeError("hadamard_128 must be callable or expose had_r_128")
+    return hadamard_128
+
+
+def _run_trellis256_dense_current_device(
+    x: torch.Tensor,
+    prepared_dense,
+    *,
+    output: torch.Tensor | None = None,
+    gemm_output: torch.Tensor | None = None,
+    c_tmp: torch.Tensor | None = None,
+    hadamard_128=None,
+    stream: cuda.CUstream | None = None,
+) -> torch.Tensor:
+    """Run one native EXL3 linear on the already-selected CUDA device.
+
+    This is a true dense entry: E=1 and contiguous row identities are synthesized
+    inside the kernel, so no top-k tensors or route-packing kernels are created.
+    Outer rotations follow EXL3 order exactly: fp16 ``suh`` multiply before the
+    input H128, and fp16 ``svh`` multiply after the output H128.  The current
+    device decode is MCG-only and rejects other codebook metadata fail-closed.
+    """
+    if getattr(prepared_dense, "weight_layout", None) != "trellis3_t256":
+        raise ValueError("run_trellis256_dense requires prepared trellis3_t256 weights")
+    if int(getattr(prepared_dense, "num_experts", 0)) != 1:
+        raise ValueError("run_trellis256_dense requires an honest E=1 prepared weight")
+    if getattr(prepared_dense, "trellis_codebook", None) != "mcg":
+        raise NotImplementedError(
+            "run_trellis256_dense currently decodes only the MCG codebook; "
+            f"got {getattr(prepared_dense, 'trellis_codebook', None)!r}"
+        )
+    trellis_bits = int(getattr(prepared_dense, "trellis_bits", 0))
+    if trellis_bits not in _TRELLIS256_BITS:
+        raise ValueError(
+            f"prepared dense weight has invalid trellis_bits={trellis_bits}"
+        )
+    compute_dtype = getattr(prepared_dense, "params_dtype", None)
+    if compute_dtype not in (torch.float16, torch.bfloat16):
+        raise TypeError(
+            "prepared dense weight must select fp16 or bf16 compute, got "
+            f"{compute_dtype}"
+        )
+    element_dtype = "fp16" if compute_dtype == torch.float16 else "bf16"
+    cutlass_dtype = (
+        cutlass.Float16 if compute_dtype == torch.float16 else cutlass.BFloat16
+    )
+    if x.ndim != 2 or int(x.shape[0]) <= 0:
+        raise ValueError(f"x must be a non-empty rank-2 tensor, got {tuple(x.shape)}")
+    if x.dtype not in (torch.float16, torch.bfloat16):
+        raise TypeError(f"x must be fp16 or bf16, got {x.dtype}")
+    if not x.is_cuda or not x.is_contiguous():
+        raise ValueError("x must be a contiguous CUDA tensor")
+    m, size_k = (int(v) for v in x.shape)
+    size_n = int(prepared_dense.out_features)
+    if size_k != int(prepared_dense.in_features):
+        raise ValueError(
+            f"x has K={size_k}, prepared dense weight expects {prepared_dense.in_features}"
+        )
+    if prepared_dense.trellis.device != x.device:
+        raise ValueError("x and prepared dense weight must be on the same CUDA device")
+    if output is None:
+        output = torch.empty((m, size_n), dtype=x.dtype, device=x.device)
+    elif (
+        tuple(output.shape) != (m, size_n)
+        or output.dtype != x.dtype
+        or output.device != x.device
+        or not output.is_contiguous()
+    ):
+        raise ValueError(
+            f"output must be contiguous {x.dtype} with shape {(m, size_n)} on {x.device}"
+        )
+    if gemm_output is None:
+        gemm_output = torch.empty(
+            (m, size_n), dtype=compute_dtype, device=x.device
+        )
+    elif (
+        tuple(gemm_output.shape) != (m, size_n)
+        or gemm_output.dtype != compute_dtype
+        or gemm_output.device != x.device
+        or not gemm_output.is_contiguous()
+        or int(gemm_output.data_ptr()) % 16 != 0
+    ):
+        raise ValueError(
+            f"gemm_output must be contiguous, 16-byte-aligned {compute_dtype} with shape "
+            f"{(m, size_n)} on {x.device}"
+        )
+    if c_tmp is not None and int(c_tmp.data_ptr()) % 16 != 0:
+        raise ValueError("c_tmp must be at least 16-byte aligned")
+
+    hadamard_128 = _resolve_exl3_hadamard_128(hadamard_128)
+    x_f16 = x if x.dtype == torch.float16 else x.to(torch.float16)
+    rotated_f16 = torch.empty_like(x_f16)
+    hadamard_128(x_f16, rotated_f16, prepared_dense.suh, None, 1.0)
+    rotated_compute = (
+        rotated_f16
+        if compute_dtype == torch.float16
+        else rotated_f16.to(torch.bfloat16)
+    )
+
+    props = torch.cuda.get_device_properties(x.device)
+    sms = int(props.multi_processor_count)
+    max_shared_mem = int(
+        getattr(props, "shared_memory_per_block_optin", _DEFAULT_MAX_SHARED_MEM)
+    )
+    moe_block_size = 64
+    route_blocks = (m + moe_block_size - 1) // moe_block_size
+    route_slots = route_blocks * moe_block_size
+    tile_k, tile_n = _trellis256_dense_tile_config(size_k, size_n)
+    launch = _compile_w4a16_gemm_launch(
+        size_m=m,
+        size_n=size_n,
+        size_k=size_k,
+        num_experts=1,
+        top_k=1,
+        mul_topk_weights=False,
+        moe_block_size=moe_block_size,
+        max_m_blocks=route_blocks,
+        element_dtype=element_dtype,
+        packed_route_indices=None,
+        sms=sms,
+        max_shared_mem=max_shared_mem,
+        device=x.device,
+        c_tmp=c_tmp,
+        weight_layout="trellis3_t256",
+        scale_format="e4m3_k32",
+        w13_layout="packed",
+        trellis_bits=trellis_bits,
+        dense_route_fast_path=True,
+        route_slots=route_slots,
+        force_tile_config=(tile_k, tile_n),
+    )
+    dummy_i32 = prepared_dense.workspace[:1]
+    stream = current_cuda_stream() if stream is None else stream
+    n_tiles = size_n // tile_n
+    grid_x = min(
+        sms * int(launch.kernel.blocks_per_sm),
+        max(route_blocks * n_tiles, 1),
+    )
+    launch.kernel.compiled(
+        make_ptr(
+            cutlass_dtype,
+            rotated_compute.data_ptr(),
+            cute.AddressSpace.gmem,
+            assumed_align=16,
+        ),
+        make_ptr(
+            cutlass_dtype,
+            rotated_compute.data_ptr(),
+            cute.AddressSpace.gmem,
+            assumed_align=16,
+        ),
+        prepared_dense.trellis,
+        make_ptr(
+            cutlass_dtype,
+            gemm_output.data_ptr(),
+            cute.AddressSpace.gmem,
+            assumed_align=16,
+        ),
+        prepared_dense.scale.view(torch.uint8).view(torch.int32).view(-1),
+        prepared_dense.global_scale,
+        dummy_i32,
+        dummy_i32,
+        dummy_i32,
+        prepared_dense.global_scale,
+        launch.c_tmp,
+        prepared_dense.workspace,
+        m,
+        grid_x,
+        stream,
+    )
+
+    gemm_f16 = (
+        gemm_output
+        if compute_dtype == torch.float16
+        else gemm_output.to(torch.float16)
+    )
+    if output.dtype == torch.float16:
+        hadamard_128(gemm_f16, output, None, prepared_dense.svh, 1.0)
+    else:
+        output_f16 = torch.empty_like(gemm_f16)
+        hadamard_128(gemm_f16, output_f16, None, prepared_dense.svh, 1.0)
+        output.copy_(output_f16)
+    return output
+
+
+def run_trellis256_dense(
+    x: torch.Tensor,
+    prepared_dense,
+    *,
+    output: torch.Tensor | None = None,
+    gemm_output: torch.Tensor | None = None,
+    c_tmp: torch.Tensor | None = None,
+    hadamard_128=None,
+    stream: cuda.CUstream | None = None,
+) -> torch.Tensor:
+    """Run one native 3/4/5/6-bpw EXL3 linear through the fused t256 GEMM.
+
+    The whole input-rotation -> GEMM -> output-rotation chain is device-guarded
+    and runs on the current Torch stream for ``x.device``.  To use a non-default
+    stream, enter ``torch.cuda.stream(...)`` around this call; accepting a raw
+    driver stream here would leave the surrounding Torch/EXL3 rotations on a
+    different stream and create an unsynchronized race.
+    """
+    if not isinstance(x, torch.Tensor) or not x.is_cuda:
+        raise ValueError("x must be a CUDA tensor")
+    if stream is not None:
+        raise ValueError(
+            "run_trellis256_dense does not accept a raw CUDA stream; select a "
+            "Torch stream with torch.cuda.stream(...) around the call"
+        )
+    with torch.cuda.device(x.device):
+        return _run_trellis256_dense_current_device(
+            x,
+            prepared_dense,
+            output=output,
+            gemm_output=gemm_output,
+            c_tmp=c_tmp,
+            hadamard_128=hadamard_128,
+            stream=None,
+        )
+
+
 def run_w4a16_moe(
     a_input: torch.Tensor,
     prepared,
@@ -7977,6 +9571,7 @@ def run_w4a16_moe(
     packed_route_count: torch.Tensor | None = None,
     expert_offsets: torch.Tensor | None = None,
     expert_map: torch.Tensor | None = None,
+    output_expert_map: torch.Tensor | None = None,
     activation_amax: torch.Tensor | None = None,
     layer_idx: int | None = None,
     apply_router_weight_on_input: bool = False,
@@ -7986,6 +9581,14 @@ def run_w4a16_moe(
     swiglu_beta: float | None = None,
     fused_launch: W4A16FusedMoeCompileResult | None = None,
     topk_sum_launch: W4A16TopKSumCompileResult | None = None,
+    intermediate_rotation_scales: torch.Tensor | None = None,
+    a_input_up: torch.Tensor | None = None,
+    full_rotation: bool = False,
+    suh_gate_table: torch.Tensor | None = None,
+    suh_up_table: torch.Tensor | None = None,
+    svh_table: torch.Tensor | None = None,
+    rotation_a_gate: torch.Tensor | None = None,
+    rotation_a_up: torch.Tensor | None = None,
     stream: cuda.CUstream | None = None,
 ) -> torch.Tensor:
     activation = normalize_moe_activation(activation)
@@ -7996,17 +9599,46 @@ def run_w4a16_moe(
         swiglu_alpha,
         swiglu_beta,
     )
-    element_dtype = _normalize_element_dtype(a_input.dtype)
-    if output.dtype != a_input.dtype:
-        raise TypeError(f"output must have dtype {a_input.dtype}, got {output.dtype}")
+    full_rotation = bool(full_rotation)
+    rotation_input_dtype = _normalize_element_dtype(a_input.dtype)
     prepared_dtype = getattr(prepared, "params_dtype", a_input.dtype)
-    if prepared_dtype != a_input.dtype:
-        raise TypeError(
-            f"prepared weights were built for {prepared_dtype}, but a_input has dtype {a_input.dtype}"
-        )
+    if full_rotation:
+        element_dtype = _normalize_element_dtype(prepared_dtype)
+        if element_dtype != "fp16":
+            raise TypeError("full_rotation requires fp16 prepared weights/scratch")
+        if output.dtype != torch.float32:
+            raise TypeError(
+                f"full_rotation output must be torch.float32, got {output.dtype}"
+            )
+    else:
+        element_dtype = rotation_input_dtype
+        if output.dtype != a_input.dtype:
+            raise TypeError(f"output must have dtype {a_input.dtype}, got {output.dtype}")
+        if prepared_dtype != a_input.dtype:
+            raise TypeError(
+                f"prepared weights were built for {prepared_dtype}, but a_input has dtype {a_input.dtype}"
+            )
     weight_layout = getattr(prepared, "weight_layout", "packed")
     if weight_layout not in _WEIGHT_LAYOUTS:
         raise ValueError(f"unsupported W4A16 weight_layout {weight_layout!r}")
+    trellis_bits = int(getattr(prepared, "trellis_bits", 3))
+    if weight_layout == "trellis3_t256":
+        if trellis_bits not in _TRELLIS256_BITS:
+            raise ValueError(
+                f"prepared trellis3_t256 bitrate must be in {_TRELLIS256_BITS}, "
+                f"got {trellis_bits}"
+            )
+        if getattr(prepared, "trellis_codebook", None) != "mcg":
+            raise NotImplementedError(
+                "trellis3_t256 execution currently requires the MCG codebook"
+            )
+        if activation_amax is not None:
+            raise NotImplementedError(
+                "trellis3_t256 activation-amax collection is not exposed through "
+                "the registered launch ABI"
+            )
+    elif trellis_bits != 3:
+        raise ValueError("trellis_bits is only valid for trellis3_t256 weights")
     scale_format = _normalize_scale_format(
         getattr(prepared, "scale_format", None)
         or (
@@ -8023,8 +9655,43 @@ def run_w4a16_moe(
     if weight_layout == "modelopt":
         if w13_layout not in _MODEL_OPT_W13_LAYOUTS:
             raise ValueError(f"unsupported W4A16 w13_layout {w13_layout!r}")
+    elif weight_layout == "trellis3_t256":
+        if w13_layout not in _TRELLIS256_W13_LAYOUTS:
+            raise ValueError(
+                f"unsupported trellis3_t256 w13_layout {w13_layout!r}"
+            )
     else:
         w13_layout = "packed"
+    dual_a_required = bool(
+        intermediate_rotation_scales is not None
+        and weight_layout == "trellis3_t256"
+        and w13_layout == "trellis3_t256_proj"
+    )
+    if full_rotation and not dual_a_required:
+        raise ValueError(
+            "full_rotation requires projection-major trellis intermediate rotation"
+        )
+    if dual_a_required and a_input_up is None and not full_rotation:
+        raise ValueError(
+            "exact projection-major trellis3_t256 rotation requires a_input_up"
+        )
+    if a_input_up is not None and (not dual_a_required or full_rotation):
+        raise ValueError(
+            "a_input_up is only valid for exact projection-major trellis3_t256 rotation"
+        )
+    if a_input_up is not None:
+        if (
+            a_input_up.shape != a_input.shape
+            or a_input_up.dtype != a_input.dtype
+            or a_input_up.device != a_input.device
+        ):
+            raise ValueError(
+                "a_input_up must match a_input shape, dtype, and device; got "
+                f"{tuple(a_input_up.shape)}/{a_input_up.dtype}/{a_input_up.device} "
+                f"vs {tuple(a_input.shape)}/{a_input.dtype}/{a_input.device}"
+            )
+        if not a_input_up.is_contiguous():
+            raise ValueError("a_input_up must be contiguous")
     if topk_weights.dtype != torch.float32:
         raise TypeError("topk_weights must be torch.float32")
     _validate_topk_ids(topk_ids, require_cuda=False, require_contiguous=False)
@@ -8035,6 +9702,9 @@ def run_w4a16_moe(
     ):
         raise ValueError("a_input, topk_weights, and topk_ids must be contiguous")
     _validate_expert_map(expert_map, device=a_input.device)
+    _validate_expert_map(output_expert_map, device=a_input.device)
+    if output_expert_map is not None and not full_rotation:
+        raise ValueError("output_expert_map is only valid with full_rotation")
 
     m, hidden_size = a_input.shape
     topk = int(topk_ids.shape[1])
@@ -8048,6 +9718,72 @@ def run_w4a16_moe(
         raise ValueError(f"output must have shape {(m, hidden_size)}")
     if expert_map is not None and int(expert_map.numel()) < int(prepared.num_experts):
         raise ValueError("expert_map cannot be shorter than the local expert count")
+    if (
+        output_expert_map is not None
+        and int(output_expert_map.numel()) < int(prepared.num_experts)
+    ):
+        raise ValueError("output_expert_map cannot be shorter than the local expert count")
+    if full_rotation:
+        if apply_router_weight_on_input:
+            raise ValueError(
+                "full_rotation applies router weights only in the fp32 top-k sum"
+            )
+        num_local_experts = int(prepared.num_experts)
+        intermediate_size_full = int(prepared.intermediate_size)
+        for name, table, shape in (
+            ("suh_gate_table", suh_gate_table, (num_local_experts, hidden_size)),
+            ("suh_up_table", suh_up_table, (num_local_experts, hidden_size)),
+            ("svh_table", svh_table, (num_local_experts, hidden_size)),
+            (
+                "intermediate_rotation_scales",
+                intermediate_rotation_scales,
+                (num_local_experts, 3 * intermediate_size_full),
+            ),
+        ):
+            if table is None:
+                raise ValueError(f"full_rotation requires {name}")
+            if (
+                table.dtype != torch.float16
+                or table.device != a_input.device
+                or tuple(table.shape) != shape
+                or not table.is_contiguous()
+            ):
+                raise ValueError(
+                    f"{name} must be contiguous fp16 {shape} on {a_input.device}; "
+                    f"got {tuple(table.shape)}/{table.dtype}/{table.device}/"
+                    f"contiguous={table.is_contiguous()}"
+                )
+        required_a = m * topk * hidden_size
+        for name, scratch in (
+            ("rotation_a_gate", rotation_a_gate),
+            ("rotation_a_up", rotation_a_up),
+        ):
+            if scratch is None:
+                raise ValueError(f"full_rotation requires preallocated {name}")
+            if (
+                scratch.dtype != torch.float16
+                or scratch.device != a_input.device
+                or not scratch.is_contiguous()
+                or int(scratch.numel()) < required_a
+            ):
+                raise ValueError(
+                    f"{name} must be contiguous fp16 on {a_input.device} with at least "
+                    f"{required_a} elements"
+                )
+        assert rotation_a_gate is not None and rotation_a_up is not None
+        if rotation_a_gate.data_ptr() == rotation_a_up.data_ptr():
+            raise ValueError("full_rotation gate/up A scratches must not alias")
+    elif any(
+        value is not None
+        for value in (
+            suh_gate_table,
+            suh_up_table,
+            svh_table,
+            rotation_a_gate,
+            rotation_a_up,
+        )
+    ):
+        raise ValueError("full-rotation tables/scratch require full_rotation=True")
     layer_idx_int = _validate_activation_amax(
         activation_amax,
         layer_idx=layer_idx,
@@ -8285,11 +10021,12 @@ def run_w4a16_moe(
             f"intermediate_cache2 has {intermediate_cache2.numel()} elements; "
             f"need at least {buffer_plan.intermediate_cache2_elements}"
         )
+    cache_dtype = prepared_dtype if full_rotation else a_input.dtype
     if (
-        intermediate_cache13.dtype != a_input.dtype
-        or intermediate_cache2.dtype != a_input.dtype
+        intermediate_cache13.dtype != cache_dtype
+        or intermediate_cache2.dtype != cache_dtype
     ):
-        raise TypeError(f"intermediate caches must be {a_input.dtype}")
+        raise TypeError(f"intermediate caches must be {cache_dtype}")
     if (
         not intermediate_cache13.is_contiguous()
         or not intermediate_cache2.is_contiguous()
@@ -8311,7 +10048,7 @@ def run_w4a16_moe(
             top_k=topk,
             activation=activation,
             apply_router_weight_on_input=bool(apply_router_weight_on_input),
-            zero_fc2_output=expert_map is not None,
+            zero_fc2_output=expert_map is not None and not full_rotation,
             moe_block_size=block_size_m,
             max_m_blocks=int(required_m_blocks),
             element_dtype=element_dtype,
@@ -8324,9 +10061,13 @@ def run_w4a16_moe(
             weight_layout=weight_layout,
             scale_format=scale_format,
             w13_layout=w13_layout,
+            trellis_bits=trellis_bits,
             direct_topk_routes=use_direct_topk_routes,
             tc_decode_fused_sum=use_tc_decode,
             collect_activation_amax=collect_activation_amax,
+            intermediate_rotation=intermediate_rotation_scales is not None,
+            full_rotation=full_rotation,
+            rotation_input_dtype=rotation_input_dtype,
         )
     else:
         if int(fused_launch.size_m) < m:
@@ -8341,7 +10082,7 @@ def run_w4a16_moe(
             topk,
             activation,
             bool(apply_router_weight_on_input),
-            expert_map is not None,
+            expert_map is not None and not full_rotation,
             element_dtype,
             bool(fast_math),
             swiglu_limit,
@@ -8350,9 +10091,14 @@ def run_w4a16_moe(
             weight_layout,
             scale_format,
             w13_layout,
+            trellis_bits,
             bool(use_direct_topk_routes),
             bool(collect_activation_amax),
             block_size_m,
+            bool(intermediate_rotation_scales is not None),
+            dual_a_required,
+            full_rotation,
+            rotation_input_dtype,
         )
         actual_fused = (
             int(fused_launch.hidden_size),
@@ -8376,9 +10122,14 @@ def run_w4a16_moe(
                 if getattr(fused_launch, "weight_layout", "packed") == "modelopt"
                 else "packed",
             ),
+            int(getattr(fused_launch, "trellis_bits", 3)),
             bool(getattr(fused_launch, "direct_topk_routes", False)),
             bool(getattr(fused_launch, "collect_activation_amax", False)),
             int(fused_launch.moe_block_size),
+            bool(getattr(fused_launch, "intermediate_rotation", False)),
+            bool(getattr(fused_launch, "dual_a", False)),
+            bool(getattr(fused_launch, "full_rotation", False)),
+            getattr(fused_launch, "rotation_input_dtype", fused_launch.element_dtype),
         )
         if actual_fused != expected_fused or int(fused_launch.max_m_blocks) < int(
             required_m_blocks
@@ -8483,7 +10234,7 @@ def run_w4a16_moe(
         topk,
         activation,
         bool(apply_router_weight_on_input),
-        expert_map is not None,
+        expert_map is not None and not full_rotation,
         block_size_m,
         int(fused.max_m_blocks),
         element_dtype,
@@ -8502,6 +10253,34 @@ def run_w4a16_moe(
         int(fused.fc2_tile_k),
         int(fused.fc2_tile_n),
     )
+    _intermediate_rotation = intermediate_rotation_scales is not None
+    if _intermediate_rotation and (
+        collect_activation_amax
+        or use_tc_decode
+        or use_direct_topk_routes
+        or weight_layout != "trellis3_t256"
+    ):
+        raise ValueError(
+            "intermediate_rotation_scales requires the packed trellis3_t256 fused "
+            "path (no calibration / tc-decode / direct-topk)"
+        )
+    if _intermediate_rotation:
+        need = (
+            int(prepared.num_experts) * 3 * intermediate_size
+            if full_rotation
+            else int(m) * topk * 3 * intermediate_size
+        )
+        rot_arg = intermediate_rotation_scales.reshape(-1)
+        if int(rot_arg.numel()) < need or rot_arg.dtype != torch.float16:
+            raise ValueError(
+                "intermediate_rotation_scales must be fp16 with >= "
+                f"{need} elements ({'experts' if full_rotation else 'routes'}"
+                "*3*intermediate); got "
+                f"numel={int(rot_arg.numel())} dtype={rot_arg.dtype}"
+            )
+        rot_arg = rot_arg[:need]
+    else:
+        rot_arg = None
     if collect_activation_amax:
         assert activation_amax is not None
         assert layer_idx_int is not None
@@ -8512,6 +10291,51 @@ def run_w4a16_moe(
             *launch_tail,
             int(stream),
         )
+    elif _intermediate_rotation or weight_layout == "trellis3_t256":
+        # Native t256 bypasses the registered torch op so its shape-derived
+        # bitrate reaches compilation without widening the stable public op ABI.
+        # Other production layouts keep the registered path byte-identical.
+        (
+            _rc_a, _rc_w13, _rc_w2, _rc_fc1, _rc_act, _rc_fc2, _rc_w13s, _rc_w2s,
+            _rc_w13g, _rc_w2g, _rc_pri, _rc_bei, _rc_prc,
+        ) = launch_common
+        (
+            _lt_tw, _lt_fc1s, _lt_fc2s, _lt_ws, _lt_m, _lt_capm, _lt_h, _lt_i,
+            _lt_ne, _lt_tk, _lt_act, _lt_arwi, _lt_zfo, _lt_bsm, _lt_mmb,
+            _lt_edt, _lt_fm, _lt_sms, _lt_msm, _lt_hsl, _lt_slv, _lt_sa, _lt_sb,
+            _lt_wl, _lt_sf, _lt_w13l, _lt_fc1tk, _lt_fc1tn, _lt_fc2tk,
+            _lt_fc2tn,
+        ) = launch_tail
+        launch_a = rotation_a_gate if full_rotation else _rc_a
+        launch_a_up = rotation_a_up if full_rotation else a_input_up
+        assert launch_a is not None
+        _w4a16_fused_moe_launch_flat(
+            a_input=launch_a, w13_arg=_rc_w13, w2_arg=_rc_w2, fc1_out=_rc_fc1,
+            activated=_rc_act, fc2_out=_rc_fc2, w13_scale_i32=_rc_w13s,
+            w2_scale_i32=_rc_w2s, w13_global_scale=_rc_w13g, w2_global_scale=_rc_w2g,
+            packed_route_indices=_rc_pri, block_expert_ids=_rc_bei,
+            packed_route_count=_rc_prc, activation_amax=None, layer_idx=0,
+            topk_weights=_lt_tw, fc1_scratch=_lt_fc1s, fc2_scratch=_lt_fc2s,
+            workspace=_lt_ws, m=_lt_m, size_m=_lt_capm, hidden_size=_lt_h,
+            intermediate_size=_lt_i, num_experts=_lt_ne, topk=_lt_tk,
+            activation=_lt_act, apply_router_weight_on_input=_lt_arwi,
+            zero_fc2_output=_lt_zfo, moe_block_size=_lt_bsm, max_m_blocks=_lt_mmb,
+            element_dtype=_lt_edt, fast_math=_lt_fm, sms=_lt_sms, max_shared_mem=_lt_msm,
+            has_swiglu_limit=_lt_hsl, swiglu_limit_value=_lt_slv, swiglu_alpha=_lt_sa,
+            swiglu_beta=_lt_sb, weight_layout=_lt_wl, scale_format=_lt_sf,
+            w13_layout=_lt_w13l, fc1_tile_k=_lt_fc1tk, fc1_tile_n=_lt_fc1tn,
+            fc2_tile_k=_lt_fc2tk, fc2_tile_n=_lt_fc2tn,
+            direct_topk_routes=bool(use_direct_topk_routes),
+            tc_decode_fused_sum=bool(use_tc_decode), collect_activation_amax=False,
+            stream_int=int(stream), rot_scales=rot_arg,
+            intermediate_rotation=_intermediate_rotation,
+            a_input_up=launch_a_up,
+            trellis_bits=trellis_bits,
+            full_rotation=full_rotation,
+            rotation_input=_rc_a,
+            suh_gate_table=suh_gate_table,
+            suh_up_table=suh_up_table,
+        )
     else:
         torch.ops.sparkinfer.w4a16_fused_moe_launch(
             *launch_common,
@@ -8525,26 +10349,58 @@ def run_w4a16_moe(
         # FC2 already wrote the top-k-summed result into `output`.
         return output
 
+    sum_expert_map = output_expert_map if output_expert_map is not None else expert_map
     if topk_sum_launch is not None:
-        expected_sum = (topk, hidden_size)
+        expected_sum = (
+            topk,
+            hidden_size,
+            full_rotation,
+            int(prepared.num_experts) if full_rotation else 0,
+            0 if sum_expert_map is None else int(sum_expert_map.numel()),
+            topk_ids.dtype if full_rotation else torch.int32,
+            bool(sum_expert_map is not None) if full_rotation else False,
+        )
         actual_sum = (
             int(topk_sum_launch.topk),
             int(topk_sum_launch.hidden_size),
+            bool(getattr(topk_sum_launch, "full_rotation", False)),
+            int(getattr(topk_sum_launch, "num_experts", 0)),
+            int(getattr(topk_sum_launch, "route_num_experts", 0)),
+            getattr(topk_sum_launch, "route_ids_dtype", torch.int32),
+            bool(getattr(topk_sum_launch, "use_expert_map", False)),
         )
         if actual_sum != expected_sum:
             raise RuntimeError(
                 "preplanned W4A16 top-k sum launch does not match requested contract: "
                 f"requested={expected_sum}, planned={actual_sum}"
             )
-    torch.ops.sparkinfer.w4a16_topk_sum_launch(
-        fc2_out,
-        output,
-        m,
-        topk,
-        hidden_size,
-        element_dtype,
-        int(stream),
-    )
+    if full_rotation:
+        assert svh_table is not None
+        _w4a16_topk_sum_launch_flat(
+            fc2_out,
+            output,
+            m,
+            topk,
+            hidden_size,
+            element_dtype,
+            int(stream),
+            full_rotation=True,
+            num_experts=int(prepared.num_experts),
+            topk_weights=topk_weights,
+            route_expert_ids=topk_ids,
+            expert_map=sum_expert_map,
+            svh_table=svh_table,
+        )
+    else:
+        torch.ops.sparkinfer.w4a16_topk_sum_launch(
+            fc2_out,
+            output,
+            m,
+            topk,
+            hidden_size,
+            element_dtype,
+            int(stream),
+        )
     return output
 
 
@@ -9294,6 +11150,7 @@ def _weight_args(prepared) -> tuple[torch.Tensor, torch.Tensor]:
     "compile_w4a16_gemm",
     "compile_w4a16_topk_sum",
     "pack_topk_routes_by_expert",
+    "run_trellis256_dense",
     "run_w4a16_moe",
     "run_w4a16_moe_hybrid",
 ]
diff --git a/sparkinfer/moe/_shared/kernels/w4a16/prepare.py b/sparkinfer/moe/_shared/kernels/w4a16/prepare.py
index 6b123b90..a650320b 100644
--- a/sparkinfer/moe/_shared/kernels/w4a16/prepare.py
+++ b/sparkinfer/moe/_shared/kernels/w4a16/prepare.py
@@ -11,6 +11,7 @@
     W4A16PackedBuffers,
     make_w4a16_packed_buffers as _make_w4a16_packed_buffers,
     unswizzle_expert_scales,
+    validate_activation,
     validate_nf3_moe_inputs,
     validate_w4a16_packed_inputs,
 )
@@ -127,6 +128,48 @@ class PreparedNF3MoeWeights:
     w13_layout: str = "packed"
     weight_layout: str = "nf3_2p1"
     scale_format: str = "e4m3_k32"
+    # Native trellis tiles are codebook-agnostic storage.  The current
+    # trellis3_t256 decoder is still MCG-only, but retaining the checkpoint
+    # codebook here lets callers fail closed instead of silently mis-decoding a
+    # MUL1 (or default-codebook) tensor.
+    trellis_codebook: str | None = None
+    trellis_bits: int = 3
+    # Projection-specific EXL3 input incoherence scales.  They remain optional
+    # for non-trellis and synthetic oracle preparation, but the coherent
+    # projection-major runtime binds both so gate/up can stage distinct rotated
+    # A operands without copying either table.
+    gate_suh: torch.Tensor | None = None
+    up_suh: torch.Tensor | None = None
+
+
+@dataclass(frozen=True)
+class PreparedTrellis256DenseWeight:
+    """Zero-copy native EXL3 tensor plus its outer rotations.
+
+    ``trellis`` is the flattened int32 view of one native
+    ``[K/16,N/16,16*bits]i16`` payload. ``suh`` and ``svh`` retain the
+    checkpoint tensors by reference. The current compute decoder is MCG-only;
+    MUL1 is preserved as metadata so execution can reject it rather than
+    misdecode it.
+    """
+
+    trellis: torch.Tensor
+    suh: torch.Tensor
+    svh: torch.Tensor
+    scale: torch.Tensor
+    global_scale: torch.Tensor
+    workspace: torch.Tensor
+    in_features: int
+    out_features: int
+    params_dtype: torch.dtype
+    trellis_bits: int
+    trellis_codebook: str
+    mcg: torch.Tensor | None = None
+    mul1: torch.Tensor | None = None
+    num_experts: int = 1
+    weight_layout: str = "trellis3_t256"
+    scale_format: str = "e4m3_k32"
+    w13_layout: str = "packed"
 
 
 def _make_workspace(
@@ -1428,6 +1471,597 @@ def prepare_nf3_moe_weights(
     )
 
 
+_TRELLIS256_W13_LAYOUTS = {"packed", "trellis3_t256_proj"}
+_TRELLIS256_CODEBOOKS = {
+    "default": "default",
+    "cb0": "default",
+    "mcg": "mcg",
+    "cb1": "mcg",
+    "mul1": "mul1",
+    "cb2": "mul1",
+}
+_TRELLIS256_CODEBOOK_SENTINELS = {
+    0: "default",
+    0xCBAC1FED: "mcg",
+    0x83DCD12D: "mul1",
+}
+
+
+def _normalize_trellis256_codebook(codebook: str | int) -> str:
+    if isinstance(codebook, int):
+        normalized = _TRELLIS256_CODEBOOK_SENTINELS.get(
+            int(codebook) & 0xFFFFFFFF
+        )
+        if normalized is None:
+            raise ValueError(
+                "unsupported trellis256 codebook sentinel "
+                f"{int(codebook) & 0xFFFFFFFF:#010x}; expected default "
+                "0x00000000, MCG 0xcbac1fed, or MUL1 0x83dcd12d"
+            )
+        return normalized
+    normalized = _TRELLIS256_CODEBOOKS.get(str(codebook).strip().lower())
+    if normalized is None:
+        raise ValueError(
+            f"unsupported trellis256 codebook {codebook!r}; expected "
+            "'default', 'mcg', or 'mul1'"
+        )
+    return normalized
+
+
+def _trellis256_random_native_tensor(
+    shape: tuple[int, ...],
+    *,
+    device: torch.device,
+    generator: torch.Generator,
+) -> torch.Tensor:
+    """Generate all 32-bit native tile words without changing their bit pattern."""
+    raw = torch.randint(
+        0,
+        1 << 32,
+        shape,
+        dtype=torch.int64,
+        device=device,
+        generator=generator,
+    )
+    raw = torch.where(raw >= 2**31, raw - 2**32, raw)
+    return raw.to(torch.int32).contiguous()
+
+
+def _trellis256_bits_from_native_tensor(tensor: torch.Tensor, *, name: str) -> int:
+    """Recover the EXL3 bitrate from the native tile's final dimension."""
+    if tensor.ndim < 1:
+        raise ValueError(f"trellis3_t256 {name} must have at least one dimension")
+    words_per_bit = 16 if tensor.dtype == torch.int16 else 8
+    if tensor.dtype not in (torch.int16, torch.int32):
+        raise TypeError(
+            f"trellis3_t256 {name} must use native int16 or int32 storage, "
+            f"got {tensor.dtype}"
+        )
+    last = int(tensor.shape[-1])
+    if last % words_per_bit != 0:
+        raise ValueError(
+            f"trellis3_t256 {name} final dimension {last} is not a native "
+            f"{tensor.dtype} tile width"
+        )
+    bits = last // words_per_bit
+    if bits not in (3, 4, 5, 6):
+        raise ValueError(
+            f"trellis3_t256 {name} encodes unsupported {bits}-bpw storage; "
+            "expected 3, 4, 5, or 6"
+        )
+    return bits
+
+
+def _trellis256_flat_native_view(
+    tensor: torch.Tensor,
+    *,
+    name: str,
+    expected_prefix_shape: tuple[int, ...],
+    trellis_bits: int,
+    device: torch.device,
+) -> torch.Tensor:
+    trellis_bits = int(trellis_bits)
+    expected_i16_shape = expected_prefix_shape + (16 * trellis_bits,)
+    expected_i32_shape = expected_prefix_shape + (8 * trellis_bits,)
+    if tensor.device != device:
+        raise ValueError(
+            f"trellis3_t256 {name} must be on {device}, got {tensor.device}"
+        )
+    if tensor.dtype == torch.int16:
+        expected_shape = expected_i16_shape
+    elif tensor.dtype == torch.int32:
+        expected_shape = expected_i32_shape
+    else:
+        raise TypeError(
+            f"trellis3_t256 {name} must be native int16 tiles ({16 * trellis_bits} words) "
+            f"or the identical int32 view ({8 * trellis_bits} words), got {tensor.dtype}"
+        )
+    if tuple(tensor.shape) != expected_shape:
+        raise ValueError(
+            f"trellis3_t256 {name} requires native {trellis_bits}-bit EXL3 shape "
+            f"{expected_shape} for dtype {tensor.dtype}, got {tuple(tensor.shape)}"
+        )
+    if not tensor.is_contiguous():
+        raise ValueError(f"trellis3_t256 {name} must be contiguous")
+    if int(tensor.data_ptr()) % 16 != 0:
+        raise ValueError(
+            f"trellis3_t256 {name} must be at least 16-byte aligned for cp.async"
+        )
+    return tensor.view(torch.int32).reshape(-1)
+
+
+def prepare_trellis256_moe_weights(
+    w13: torch.Tensor | None = None,
+    w2: torch.Tensor | None = None,
+    *,
+    hidden_size: int,
+    intermediate_size: int,
+    num_experts: int,
+    activation: str,
+    fc1_tile_n: int,
+    fc2_tile_n: int,
+    device: torch.device | str | int | None = None,
+    seed: int = 0,
+    params_dtype: torch.dtype = torch.bfloat16,
+    w13_layout: str = "packed",
+    trellis_bits: int | None = None,
+    dummy_scale: torch.Tensor | None = None,
+    codebook: str | int = "mcg",
+    gate_suh: torch.Tensor | None = None,
+    up_suh: torch.Tensor | None = None,
+) -> PreparedNF3MoeWeights:
+    """Wrap or synthesize native EXL3 tiles for ``trellis3_t256``.
+
+    Supplying both ``w13`` and ``w2`` is the production path: no bytes are
+    copied or permuted; each tensor is only viewed as contiguous int32 words and
+    flattened.  Omitting both tensors is the deterministic full-GEMM-oracle
+    path selected by ``device`` and ``seed``.  ``gate_suh`` and ``up_suh`` are
+    optional zero-copy bindings for the two projection-specific EXL3 input
+    rotations; when supplied, both must be contiguous fp16 ``[E,H]`` tensors on
+    the weight device.
+
+    Plain FC1 storage is expert-major
+    ``[E,H/16,FC1_N/16,16*bits]i16``. ``trellis3_t256_proj`` instead requires one
+    projection-major backing ``[2,E,H/16,I/16,16*bits]i16`` so gate/up fallback views
+    can continue to alias the same live storage.  FC2 is always the plain native
+    ``[E,I/16,H/16,16*bits]i16`` layout.
+    """
+    hidden_size = int(hidden_size)
+    intermediate_size = int(intermediate_size)
+    num_experts = int(num_experts)
+    fc1_tile_n = int(fc1_tile_n)
+    fc2_tile_n = int(fc2_tile_n)
+    if params_dtype not in (torch.bfloat16, torch.float16):
+        raise ValueError(
+            "trellis3_t256 W4A16 weights require fp16 or bf16 activations"
+        )
+    requested_trellis_bits = None if trellis_bits is None else int(trellis_bits)
+    if requested_trellis_bits is not None and requested_trellis_bits not in (3, 4, 5, 6):
+        raise ValueError(
+            "trellis3_t256 bits must be one of 3, 4, 5, or 6, "
+            f"got {requested_trellis_bits}"
+        )
+    if num_experts <= 0:
+        raise ValueError(
+            f"trellis3_t256 requires num_experts > 0, got {num_experts}"
+        )
+    if hidden_size <= 0 or intermediate_size <= 0:
+        raise ValueError(
+            "trellis3_t256 requires positive hidden_size and intermediate_size, "
+            f"got H={hidden_size} I={intermediate_size}"
+        )
+    if hidden_size % 16 != 0 or intermediate_size % 16 != 0:
+        raise ValueError(
+            "native EXL3 tiles require hidden_size and intermediate_size to be "
+            f"multiples of 16, got H={hidden_size} I={intermediate_size}"
+        )
+    if hidden_size % 32 != 0 or intermediate_size % 32 != 0:
+        raise ValueError(
+            "trellis3_t256 uses E4M3 K/32 kernel plumbing and therefore requires "
+            "hidden_size and intermediate_size to be multiples of 32; "
+            f"got H={hidden_size} I={intermediate_size}"
+        )
+    for name, tile_n in (("fc1_tile_n", fc1_tile_n), ("fc2_tile_n", fc2_tile_n)):
+        if tile_n < 64 or tile_n % 16 != 0:
+            raise ValueError(
+                f"trellis3_t256 {name} must be a multiple of 16 and at least "
+                f"64 for the current W4A16 kernel, got {tile_n}"
+            )
+
+    is_gated = validate_activation(activation)
+    w13_rows = intermediate_size * (2 if is_gated else 1)
+    if w13_layout not in _TRELLIS256_W13_LAYOUTS:
+        raise ValueError(
+            f"unsupported trellis3_t256 w13_layout {w13_layout!r}; expected "
+            "'packed' or 'trellis3_t256_proj'"
+        )
+    if w13_layout == "trellis3_t256_proj":
+        if not is_gated:
+            raise ValueError(
+                "trellis3_t256_proj requires a gated activation with separate "
+                "gate/up FC1 projections"
+            )
+        if intermediate_size % fc1_tile_n != 0:
+            raise ValueError(
+                "trellis3_t256_proj requires each FC1 projection to contain an "
+                f"integral number of CTA N tiles: I={intermediate_size}, "
+                f"fc1_tile_n={fc1_tile_n}"
+            )
+    elif w13_rows % fc1_tile_n != 0:
+        raise ValueError(
+            "trellis3_t256 has no FC1 logical-tail path: "
+            f"FC1_N={w13_rows} must be divisible by fc1_tile_n={fc1_tile_n}"
+        )
+    if hidden_size % fc2_tile_n != 0:
+        raise ValueError(
+            "trellis3_t256 has no FC2 logical-tail path: "
+            f"FC2_N={hidden_size} must be divisible by fc2_tile_n={fc2_tile_n}"
+        )
+    normalized_codebook = _normalize_trellis256_codebook(codebook)
+
+    have_w13 = w13 is not None
+    have_w2 = w2 is not None
+    if have_w13 != have_w2:
+        raise ValueError("trellis3_t256 requires both w13 and w2, or neither")
+    synthetic = not have_w13
+    if synthetic:
+        resolved_trellis_bits = requested_trellis_bits or 3
+        if device is None:
+            raise ValueError(
+                "device is required when synthesizing trellis3_t256 oracle weights"
+            )
+        if isinstance(device, int):
+            resolved_device = torch.device("cuda", int(device))
+        else:
+            resolved_device = torch.device(device)
+        if resolved_device.type != "cuda":
+            raise ValueError(
+                "trellis3_t256 W4A16 weights require a CUDA device, got "
+                f"{resolved_device}"
+            )
+        generator = torch.Generator(device=resolved_device)
+        generator.manual_seed(int(seed))
+        if w13_layout == "trellis3_t256_proj":
+            w13_i32_shape = (
+                2,
+                num_experts,
+                hidden_size // 16,
+                intermediate_size // 16,
+                8 * resolved_trellis_bits,
+            )
+        else:
+            w13_i32_shape = (
+                num_experts,
+                hidden_size // 16,
+                w13_rows // 16,
+                8 * resolved_trellis_bits,
+            )
+        w2_i32_shape = (
+            num_experts,
+            intermediate_size // 16,
+            hidden_size // 16,
+            8 * resolved_trellis_bits,
+        )
+        packed_w13 = _trellis256_random_native_tensor(
+            w13_i32_shape, device=resolved_device, generator=generator
+        ).reshape(-1)
+        packed_w2 = _trellis256_random_native_tensor(
+            w2_i32_shape, device=resolved_device, generator=generator
+        ).reshape(-1)
+    else:
+        assert w13 is not None and w2 is not None
+        w13_bits = _trellis256_bits_from_native_tensor(w13, name="w13")
+        w2_bits = _trellis256_bits_from_native_tensor(w2, name="w2")
+        if w13_bits != w2_bits:
+            raise ValueError(
+                f"trellis3_t256 w13/w2 bitrate mismatch: {w13_bits} vs {w2_bits}"
+            )
+        resolved_trellis_bits = w13_bits
+        if (
+            requested_trellis_bits is not None
+            and requested_trellis_bits != resolved_trellis_bits
+        ):
+            raise ValueError(
+                "explicit trellis_bits disagrees with native tensor shape: "
+                f"requested={requested_trellis_bits}, inferred={resolved_trellis_bits}"
+            )
+        resolved_device = w13.device
+        if resolved_device.type != "cuda":
+            raise ValueError(
+                "trellis3_t256 W4A16 weights require CUDA storage, got "
+                f"{resolved_device}"
+            )
+        if device is not None:
+            requested_device = (
+                torch.device("cuda", int(device))
+                if isinstance(device, int)
+                else torch.device(device)
+            )
+            device_matches = (
+                requested_device.type == resolved_device.type
+                and (
+                    requested_device.index is None
+                    or requested_device.index == resolved_device.index
+                )
+            )
+            if not device_matches:
+                raise ValueError(
+                    "explicit trellis3_t256 device does not match supplied weights: "
+                    f"device={requested_device}, weights={resolved_device}"
+                )
+        if w13_layout == "trellis3_t256_proj":
+            expected_w13_prefix = (
+                2,
+                num_experts,
+                hidden_size // 16,
+                intermediate_size // 16,
+            )
+        else:
+            expected_w13_prefix = (
+                num_experts,
+                hidden_size // 16,
+                w13_rows // 16,
+            )
+        expected_w2_prefix = (
+            num_experts,
+            intermediate_size // 16,
+            hidden_size // 16,
+        )
+        packed_w13 = _trellis256_flat_native_view(
+            w13,
+            name="w13",
+            expected_prefix_shape=expected_w13_prefix,
+            trellis_bits=resolved_trellis_bits,
+            device=resolved_device,
+        )
+        packed_w2 = _trellis256_flat_native_view(
+            w2,
+            name="w2",
+            expected_prefix_shape=expected_w2_prefix,
+            trellis_bits=resolved_trellis_bits,
+            device=resolved_device,
+        )
+
+    have_gate_suh = gate_suh is not None
+    have_up_suh = up_suh is not None
+    if have_gate_suh != have_up_suh:
+        raise ValueError(
+            "trellis3_t256 projection input scales require both gate_suh and up_suh"
+        )
+    if have_gate_suh:
+        if w13_layout != "trellis3_t256_proj":
+            raise ValueError(
+                "trellis3_t256 gate_suh/up_suh bindings require "
+                "w13_layout='trellis3_t256_proj'"
+            )
+        assert gate_suh is not None and up_suh is not None
+        for name, scale in (("gate_suh", gate_suh), ("up_suh", up_suh)):
+            if scale.device != resolved_device:
+                raise ValueError(
+                    f"trellis3_t256 {name} must be on {resolved_device}, got {scale.device}"
+                )
+            if scale.dtype != torch.float16:
+                raise TypeError(
+                    f"trellis3_t256 {name} must be torch.float16, got {scale.dtype}"
+                )
+            if tuple(scale.shape) != (num_experts, hidden_size):
+                raise ValueError(
+                    f"trellis3_t256 {name} must have shape "
+                    f"{(num_experts, hidden_size)}, got {tuple(scale.shape)}"
+                )
+            if not scale.is_contiguous():
+                raise ValueError(f"trellis3_t256 {name} must be contiguous")
+
+    if dummy_scale is None:
+        dummy_scale = torch.zeros(4, dtype=torch.uint8, device=resolved_device)
+    else:
+        dummy_device_matches = (
+            dummy_scale.device.type == resolved_device.type
+            and (
+                resolved_device.index is None
+                or dummy_scale.device.index == resolved_device.index
+            )
+        )
+        if not dummy_device_matches:
+            raise ValueError(
+                "trellis3_t256 dummy_scale must share the weight device, got "
+                f"{dummy_scale.device} and {resolved_device}"
+            )
+        if dummy_scale.dtype != torch.uint8:
+            raise TypeError(
+                f"trellis3_t256 dummy_scale must be torch.uint8, got {dummy_scale.dtype}"
+            )
+        if not dummy_scale.is_contiguous() or tuple(dummy_scale.shape) != (4,):
+            raise ValueError(
+                "trellis3_t256 dummy_scale must be a contiguous four-byte "
+                f"tensor with shape (4,), got {tuple(dummy_scale.shape)}"
+            )
+        if int(dummy_scale.data_ptr()) % 16 != 0:
+            raise ValueError(
+                "trellis3_t256 dummy_scale must be at least 16-byte aligned"
+            )
+
+    global_scale = torch.ones(
+        (num_experts,), dtype=torch.float32, device=resolved_device
+    )
+    return PreparedNF3MoeWeights(
+        w13=packed_w13,
+        w13_scale=dummy_scale,
+        w13_global_scale=global_scale,
+        w2=packed_w2,
+        w2_scale=dummy_scale,
+        w2_global_scale=global_scale,
+        workspace=_make_workspace(resolved_device, max_blocks_per_sm=4),
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        num_experts=num_experts,
+        is_gated=is_gated,
+        params_dtype=params_dtype,
+        fc1_tile_n=fc1_tile_n,
+        fc2_tile_n=fc2_tile_n,
+        source_format="trellis3_t256",
+        w13_layout=w13_layout,
+        weight_layout="trellis3_t256",
+        scale_format="e4m3_k32",
+        trellis_codebook=normalized_codebook,
+        trellis_bits=resolved_trellis_bits,
+        gate_suh=gate_suh,
+        up_suh=up_suh,
+    )
+
+
+def _trellis256_marker_codebook(
+    *,
+    mcg: torch.Tensor | None,
+    mul1: torch.Tensor | None,
+    codebook: str | int | None,
+) -> str:
+    if mcg is not None and mul1 is not None:
+        raise ValueError("trellis256 accepts exactly one of mcg or mul1, not both")
+    marker_codebook: str | None = None
+    marker = mcg if mcg is not None else mul1
+    if marker is not None:
+        name = "mcg" if mcg is not None else "mul1"
+        if marker.numel() != 1 or marker.dtype not in (torch.int32, torch.uint32):
+            raise ValueError(
+                f"trellis256 {name} marker must be a scalar int32/uint32 tensor"
+            )
+        marker_value = int(marker.item()) & 0xFFFFFFFF
+        marker_codebook = _normalize_trellis256_codebook(marker_value)
+        if marker_codebook != name:
+            raise ValueError(
+                f"trellis256 {name} marker has unexpected value {marker_value:#010x}"
+            )
+    explicit = (
+        None if codebook is None else _normalize_trellis256_codebook(codebook)
+    )
+    if marker_codebook is not None and explicit is not None and marker_codebook != explicit:
+        raise ValueError(
+            "trellis256 codebook metadata disagrees: "
+            f"marker={marker_codebook!r}, codebook={explicit!r}"
+        )
+    normalized = marker_codebook or explicit
+    if normalized is None:
+        raise ValueError(
+            "trellis256 dense preparation requires mcg=, mul1=, or explicit codebook="
+        )
+    return normalized
+
+
+def prepare_trellis256_dense_weight(
+    trellis: torch.Tensor,
+    suh: torch.Tensor,
+    svh: torch.Tensor,
+    *,
+    mcg: torch.Tensor | None = None,
+    mul1: torch.Tensor | None = None,
+    codebook: str | int | None = None,
+    params_dtype: torch.dtype = torch.float16,
+    dummy_scale: torch.Tensor | None = None,
+) -> PreparedTrellis256DenseWeight:
+    """Prepare one native EXL3 linear for the dense trellis256 entry point.
+
+    The native payload is ``[K/16,N/16,16*bits]i16`` (or the byte-identical
+    ``[...,8*bits]i32`` view), optionally with a leading singleton expert axis.
+    The bitrate is inferred from that final dimension. No trellis or rotation
+    bytes are copied, permuted, stacked, or concatenated.
+    """
+    if params_dtype not in (torch.float16, torch.bfloat16):
+        raise ValueError(
+            "trellis3_t256 dense compute requires fp16 or bf16 MMA inputs"
+        )
+    if trellis.ndim not in (3, 4):
+        raise ValueError(
+            "trellis3_t256 dense payload must have rank 3 or a leading E=1 axis"
+        )
+    if trellis.ndim == 4:
+        if int(trellis.shape[0]) != 1:
+            raise ValueError(
+                f"trellis3_t256 dense payload requires E=1, got {int(trellis.shape[0])}"
+            )
+        k16, n16 = int(trellis.shape[1]), int(trellis.shape[2])
+    else:
+        k16, n16 = int(trellis.shape[0]), int(trellis.shape[1])
+    trellis_bits = _trellis256_bits_from_native_tensor(
+        trellis, name="dense trellis"
+    )
+    in_features = k16 * 16
+    out_features = n16 * 16
+    if in_features <= 0 or out_features <= 0:
+        raise ValueError(
+            f"trellis3_t256 dense dimensions must be positive, got {in_features}x{out_features}"
+        )
+    if in_features % 128 != 0 or out_features % 128 != 0:
+        raise ValueError(
+            "trellis3_t256 dense rotations require K and N divisible by 128; "
+            f"got K={in_features} N={out_features}"
+        )
+    expected_prefix_shape = (
+        (1, k16, n16) if trellis.ndim == 4 else (k16, n16)
+    )
+    device = trellis.device
+    if device.type != "cuda":
+        raise ValueError(
+            f"trellis3_t256 dense weights require CUDA storage, got {device}"
+        )
+    packed = _trellis256_flat_native_view(
+        trellis,
+        name="dense trellis",
+        expected_prefix_shape=expected_prefix_shape,
+        trellis_bits=trellis_bits,
+        device=device,
+    )
+    for name, scale, width in (
+        ("suh", suh, in_features),
+        ("svh", svh, out_features),
+    ):
+        if scale.device != device:
+            raise ValueError(f"trellis3_t256 dense {name} must be on {device}")
+        if scale.dtype != torch.float16:
+            raise TypeError(
+                f"trellis3_t256 dense {name} must be torch.float16, got {scale.dtype}"
+            )
+        if tuple(scale.shape) != (width,):
+            raise ValueError(
+                f"trellis3_t256 dense {name} must have shape {(width,)}, "
+                f"got {tuple(scale.shape)}"
+            )
+        if not scale.is_contiguous():
+            raise ValueError(f"trellis3_t256 dense {name} must be contiguous")
+    normalized_codebook = _trellis256_marker_codebook(
+        mcg=mcg, mul1=mul1, codebook=codebook
+    )
+    if dummy_scale is None:
+        dummy_scale = torch.zeros(4, dtype=torch.uint8, device=device)
+    else:
+        if (
+            dummy_scale.device != device
+            or dummy_scale.dtype != torch.uint8
+            or tuple(dummy_scale.shape) != (4,)
+            or not dummy_scale.is_contiguous()
+            or int(dummy_scale.data_ptr()) % 16 != 0
+        ):
+            raise ValueError(
+                "trellis3_t256 dense dummy_scale must be a contiguous, 16-byte-"
+                "aligned four-byte uint8 tensor on the weight device"
+            )
+    return PreparedTrellis256DenseWeight(
+        trellis=packed,
+        suh=suh,
+        svh=svh,
+        scale=dummy_scale,
+        global_scale=torch.ones(1, dtype=torch.float32, device=device),
+        workspace=_make_workspace(device, max_blocks_per_sm=4),
+        in_features=in_features,
+        out_features=out_features,
+        params_dtype=params_dtype,
+        trellis_bits=trellis_bits,
+        trellis_codebook=normalized_codebook,
+        mcg=mcg,
+        mul1=mul1,
+    )
+
+
 def _nf3_pack_selftest() -> None:
     """Prove the packer matches the kernel's read+dequant chain (torch/CPU).
 
@@ -1544,11 +2178,14 @@ def code(byte2, nib1, j):
 
 __all__ = [
     "PreparedNF3MoeWeights",
+    "PreparedTrellis256DenseWeight",
     "W4A16PackedBuffers",
     "W4A16ModelOptWeights",
     "W4A16PackedWeights",
     "make_w4a16_packed_buffers",
     "prepare_nf3_moe_weights",
+    "prepare_trellis256_moe_weights",
+    "prepare_trellis256_dense_weight",
     "prepare_w4a16_compressed_tensors_weights",
     "prepare_w4a16_e8m0_native_weights",
     "prepare_w4a16_fp4_e8m0_k32_weights",

From 6541ca701563917b9b655b3f54235d0d201f2420 Mon Sep 17 00:00:00 2001
From: Brandon Music <brandon.m.music@gmail.com>
Date: Tue, 21 Jul 2026 18:58:09 -0400
Subject: [PATCH 03/11] moe: expose planned EXL3 Trellis API

Co-authored-by: OpenAI Codex <codex@openai.com>

Signed-off-by: Brandon Music <brandon.m.music@gmail.com>
---
 sparkinfer/__init__.py                        |   1 +
 sparkinfer/moe/__init__.py                    |   3 +-
 sparkinfer/moe/_shared/kernels/w4a16/host.py  |  19 +-
 .../moe/_shared/kernels/w4a16/kernel.py       |   7 +
 .../moe/_shared/kernels/w4a16/route_pack.py   |  31 +-
 sparkinfer/moe/trellis_moe/__init__.py        |  94 +++
 sparkinfer/moe/trellis_moe/_impl.py           | 775 ++++++++++++++++++
 sparkinfer/moe/trellis_moe/api.py             |  52 ++
 tests/moe/test_trellis_moe.py                 | 471 +++++++++++
 9 files changed, 1450 insertions(+), 3 deletions(-)
 create mode 100644 sparkinfer/moe/trellis_moe/__init__.py
 create mode 100644 sparkinfer/moe/trellis_moe/_impl.py
 create mode 100644 sparkinfer/moe/trellis_moe/api.py
 create mode 100644 tests/moe/test_trellis_moe.py

diff --git a/sparkinfer/__init__.py b/sparkinfer/__init__.py
index f3fd439f..65b76389 100644
--- a/sparkinfer/__init__.py
+++ b/sparkinfer/__init__.py
@@ -53,6 +53,7 @@
     "gemm.wo_projection",
     "moe.fused_moe",
     "moe.ep_moe",
+    "moe.trellis_moe",
     "norm.mhc",
     "quantization.mxfp8",
     "quantization.nvfp4",
diff --git a/sparkinfer/moe/__init__.py b/sparkinfer/moe/__init__.py
index faf27bac..7cdb84c1 100644
--- a/sparkinfer/moe/__init__.py
+++ b/sparkinfer/moe/__init__.py
@@ -4,6 +4,7 @@
   activation -> FC2 -> scatter); recipes nvfp4/mxfp4/w4a8_mx/w4a8_nvfp4/w4a16.
 - ``ep_moe``: expert-parallel MoE (replicated input -> local partial;
   cross-rank reduction is the caller's job, typically ``comm.pcie``).
+- ``trellis_moe``: planned full-rotation MCG EXL3 routed experts.
 """
 
 from __future__ import annotations
@@ -11,7 +12,7 @@
 import importlib
 from typing import Any
 
-_OP_MODULES = ("fused_moe", "ep_moe")
+_OP_MODULES = ("fused_moe", "ep_moe", "trellis_moe")
 
 
 def __getattr__(name: str) -> Any:
diff --git a/sparkinfer/moe/_shared/kernels/w4a16/host.py b/sparkinfer/moe/_shared/kernels/w4a16/host.py
index 856a046a..b7970ade 100644
--- a/sparkinfer/moe/_shared/kernels/w4a16/host.py
+++ b/sparkinfer/moe/_shared/kernels/w4a16/host.py
@@ -38,6 +38,7 @@ class W4A16PackedBuffers:
     block_expert_ids: torch.Tensor | None = None
     packed_route_count: torch.Tensor | None = None
     expert_offsets: torch.Tensor | None = None
+    expert_counts: torch.Tensor | None = None
     rotation_a_gate: torch.Tensor | None = None
     rotation_a_up: torch.Tensor | None = None
 
@@ -52,6 +53,7 @@ class W4A16BufferPlan:
     fc2_c_tmp_elements: int
     intermediate_cache13_elements: int
     intermediate_cache2_elements: int
+    block_size_m: int
     rotation_a_elements: int = 0
 
 
@@ -276,6 +278,7 @@ def plan_w4a16_buffers(
     route_num_experts: int | None = None,
     sms: int,
     full_rotation: bool = False,
+    block_size_m: int | None = None,
 ) -> W4A16BufferPlan:
     routed_rows = int(m) * int(topk)
     route_num_experts = (
@@ -286,7 +289,15 @@ def plan_w4a16_buffers(
     intermediate_size = int(prepared.intermediate_size)
     hidden_size = int(prepared.hidden_size)
     fc1_cols = (2 if prepared.is_gated else 1) * intermediate_size
-    block_size_m = select_route_block_size_m(m, topk, route_num_experts)
+    if block_size_m is None:
+        block_size_m = select_route_block_size_m(m, topk, route_num_experts)
+    else:
+        block_size_m = int(block_size_m)
+        if block_size_m not in _W4A16_ALLOWED_ROUTED_SIZES:
+            raise ValueError(
+                "block_size_m must be one of "
+                f"{_W4A16_ALLOWED_ROUTED_SIZES}, got {block_size_m}"
+            )
     route_slots = max_packed_route_slots(routed_rows, block_size_m, route_num_experts)
     route_blocks = (route_slots + block_size_m - 1) // block_size_m
     scratch_sms = int(sms)
@@ -309,6 +320,7 @@ def plan_w4a16_buffers(
         ),
         intermediate_cache13_elements=routed_rows * max(fc1_cols, hidden_size),
         intermediate_cache2_elements=routed_rows * intermediate_size,
+        block_size_m=block_size_m,
         rotation_a_elements=(routed_rows * hidden_size if full_rotation else 0),
     )
 
@@ -322,6 +334,7 @@ def make_w4a16_packed_buffers(
     device: torch.device,
     route_num_experts: int | None = None,
     full_rotation: bool = False,
+    block_size_m: int | None = None,
 ) -> W4A16PackedBuffers:
     route_num_experts = (
         int(prepared.num_experts)
@@ -336,6 +349,7 @@ def make_w4a16_packed_buffers(
         route_num_experts=route_num_experts,
         sms=sms,
         full_rotation=full_rotation,
+        block_size_m=block_size_m,
     )
     fc1_c_tmp = torch.empty(
         (plan.fc1_c_tmp_elements,),
@@ -375,6 +389,9 @@ def make_w4a16_packed_buffers(
         expert_offsets=torch.empty(
             (route_num_experts + 1,), dtype=torch.int32, device=device
         ),
+        expert_counts=torch.empty(
+            (route_num_experts,), dtype=torch.int32, device=device
+        ),
         rotation_a_gate=(
             torch.empty(
                 (plan.routed_rows, int(prepared.hidden_size)),
diff --git a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
index 0c7f59ff..3be092c3 100644
--- a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
+++ b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
@@ -1029,6 +1029,7 @@ def __cache_key__(self) -> tuple[object, ...]:
             self.fused_sum_topk,
             self.cta_m_blocks,
             self.uses_m_block_8,
+            self.max_m_blocks,
             self.shared_words,
             # Launch bounds are part of the compiled kernel.  Keep binaries
             # planned for different residency targets out of the same cache
@@ -9279,6 +9280,7 @@ def pack_topk_routes_by_expert(
     block_expert_ids: torch.Tensor | None = None,
     packed_route_count: torch.Tensor | None = None,
     expert_offsets: torch.Tensor | None = None,
+    expert_counts: torch.Tensor | None = None,
     stream: cuda.CUstream | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """Group top-k routes by expert and pad each group to the GEMM M-block size."""
@@ -9298,6 +9300,7 @@ def pack_topk_routes_by_expert(
         block_expert_ids=block_expert_ids,
         packed_route_count=packed_route_count,
         expert_offsets=expert_offsets,
+        expert_counts=expert_counts,
     )
 
 
@@ -9570,6 +9573,7 @@ def run_w4a16_moe(
     block_expert_ids: torch.Tensor | None = None,
     packed_route_count: torch.Tensor | None = None,
     expert_offsets: torch.Tensor | None = None,
+    expert_counts: torch.Tensor | None = None,
     expert_map: torch.Tensor | None = None,
     output_expert_map: torch.Tensor | None = None,
     activation_amax: torch.Tensor | None = None,
@@ -9990,6 +9994,7 @@ def run_w4a16_moe(
                 block_expert_ids=block_expert_ids,
                 packed_route_count=packed_route_count,
                 expert_offsets=expert_offsets,
+                expert_counts=expert_counts,
                 stream=stream,
             )
         )
@@ -10007,6 +10012,8 @@ def run_w4a16_moe(
         topk=topk,
         route_num_experts=route_num_experts,
         sms=sms,
+        full_rotation=full_rotation,
+        block_size_m=block_size_m,
     )
     intermediate_size = int(prepared.intermediate_size)
     fc1_cols = buffer_plan.fc1_cols
diff --git a/sparkinfer/moe/_shared/kernels/w4a16/route_pack.py b/sparkinfer/moe/_shared/kernels/w4a16/route_pack.py
index 2ba575a9..d3656ada 100644
--- a/sparkinfer/moe/_shared/kernels/w4a16/route_pack.py
+++ b/sparkinfer/moe/_shared/kernels/w4a16/route_pack.py
@@ -302,6 +302,7 @@ def pack_topk_routes_by_expert(
     block_expert_ids: torch.Tensor | None = None,
     packed_route_count: torch.Tensor | None = None,
     expert_offsets: torch.Tensor | None = None,
+    expert_counts: torch.Tensor | None = None,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     numel = int(topk_ids.numel())
     topk = int(topk_ids.shape[-1]) if topk_ids.ndim >= 2 else 1
@@ -406,7 +407,35 @@ def pack_topk_routes_by_expert(
                 triton.cdiv(max_route_blocks, _POST_PREFIX_BLOCK_T),
             ),
         )
-        if not torch.cuda.is_current_stream_capturing():
+        if expert_counts is not None:
+            expert_counts = _workspace_slice(
+                expert_counts,
+                name="expert_counts",
+                elements=int(num_experts),
+                dtype=torch.int32,
+                device=topk_ids.device,
+            )
+            expert_counts.zero_()
+            _w4a16_route_count_kernel[(triton.cdiv(numel, _FAST_COUNT_BLOCK_T),)](
+                topk_ids,
+                expert_map_tensor,
+                expert_counts,
+                numel,
+                NUM_EXPERTS=int(num_experts),
+                HAS_EXPERT_MAP=expert_map is not None,
+                BLOCK_T=_FAST_COUNT_BLOCK_T,
+                num_warps=4,
+            )
+            _w4a16_route_prefix_from_counts_kernel[(1,)](
+                expert_counts,
+                packed_route_count,
+                expert_offsets,
+                BLOCK_SIZE=int(block_size),
+                NUM_EXPERTS=int(num_experts),
+                BLOCK_E=block_e,
+                num_warps=4,
+            )
+        elif not torch.cuda.is_current_stream_capturing():
             # FAST (eager prefill): parallel atomic count + tiny over-experts
             # block-padded prefix, replacing the single-CTA count+prefix
             # (~7-31x measured at prefill). Only the large path (routes > 4096,
diff --git a/sparkinfer/moe/trellis_moe/__init__.py b/sparkinfer/moe/trellis_moe/__init__.py
new file mode 100644
index 00000000..417eaea0
--- /dev/null
+++ b/sparkinfer/moe/trellis_moe/__init__.py
@@ -0,0 +1,94 @@
+"""Planned full-rotation EXL3 Trellis MoE for SM12x.
+
+This op owns the production ``trellis3_t256`` MCG lifecycle. Weight preparation
+wraps projection-major native EXL3 tensors without repacking. Planning fixes the
+token, route, tile, and exact M-block capacities and eagerly compiles both the
+fused MoE launch and full-rotation FP32 top-k reductions. Binding maps one
+caller-owned uint8 scratch arena into stable views; ``run`` performs no tensor
+allocation and is CUDA-graph-capture safe after ordinary eager warmup.
+
+Example:
+    from sparkinfer.moe import trellis_moe
+
+    weights = trellis_moe.prepare_weights(
+        w13, w2,
+        gate_suh=gate_suh,
+        up_suh=up_suh,
+        intermediate_rotations=intermediate_rotations,
+        down_svh=down_svh,
+        codebook="mcg",
+    )
+    plan = trellis_moe.plan(trellis_moe.Caps(
+        max_tokens=32,
+        num_topk=8,
+        num_experts=weights.num_experts,
+        hidden_size=weights.hidden_size,
+        intermediate_size=weights.intermediate_size,
+        route_num_experts=256,
+        block_size_m=8,
+        input_dtype=torch.bfloat16,
+        device=weights.device,
+    ))
+    spec = plan.scratch_specs()[0]
+    scratch = torch.empty(spec.shape, dtype=spec.dtype, device=spec.device)
+    binding = trellis_moe.bind(
+        plan,
+        scratch=scratch,
+        a=x,
+        weights=weights,
+        topk_weights=router_weights,
+        topk_ids=router_ids,
+    )
+    output = trellis_moe.run(binding=binding)
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ..._lib.meta import OpMeta, Provenance, install_lazy_api
+
+META = OpMeta(
+    name="trellis_moe",
+    group="moe",
+    api_style="planned",
+    entry_points=(
+        "Caps",
+        "Plan",
+        "Weights",
+        "Binding",
+        "plan",
+        "prepare_weights",
+        "bind",
+        "run",
+        "is_supported",
+        "clear_caches",
+    ),
+    dtypes=("bf16", "fp16"),
+    recipes=("trellis3_t256_mcg",),
+    requires=("triton",),
+    provenance=Provenance(
+        repo="https://github.com/brandonmmusic-max/b12x",
+        commit="e611971",
+        paths=("b12x/moe/fused/w4a16/",),
+    ),
+    test_path="tests/moe/test_trellis_moe.py",
+    since="1.1.0",
+    notes="Projection-major MCG Trellis weights with fused full rotations.",
+)
+
+if TYPE_CHECKING:  # static analysis only; runtime resolution is lazy
+    from .api import (  # noqa: F401
+        Binding,
+        Caps,
+        Plan,
+        Weights,
+        bind,
+        clear_caches,
+        is_supported,
+        plan,
+        prepare_weights,
+        run,
+    )
+
+install_lazy_api(globals(), META)
diff --git a/sparkinfer/moe/trellis_moe/_impl.py b/sparkinfer/moe/trellis_moe/_impl.py
new file mode 100644
index 00000000..ef16a19e
--- /dev/null
+++ b/sparkinfer/moe/trellis_moe/_impl.py
@@ -0,0 +1,775 @@
+"""Implementation of the public planned full-rotation Trellis MoE op."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass, field, replace
+from math import prod
+
+import torch
+
+from ..._lib.scratch import ScratchBufferSpec, scratch_buffer_spec, scratch_tensor
+from .._shared.kernels.w4a16.host import (
+    W4A16BufferPlan,
+    max_packed_route_slots,
+    plan_w4a16_buffers,
+)
+from .._shared.kernels.w4a16.kernel import (
+    W4A16FusedMoeCompileResult,
+    W4A16TopKSumCompileResult,
+    clear_w4a16_kernel_cache,
+    compile_w4a16_fused_moe,
+    compile_w4a16_topk_sum,
+    run_w4a16_moe,
+)
+from .._shared.kernels.w4a16.prepare import (
+    PreparedNF3MoeWeights,
+    _normalize_trellis256_codebook,
+    prepare_trellis256_moe_weights,
+)
+
+
+_ALLOWED_BLOCK_M = (8, 16, 32, 48, 64)
+_ALLOWED_BITS = (3, 4, 5, 6)
+_DEFAULT_TILE_CONFIG = (64, 256, 64, 256)
+_MCG_SENTINEL = 0xCBAC1FED
+_ARENA_ALIGNMENT = 256
+
+
+def _align_up(value: int, alignment: int = _ARENA_ALIGNMENT) -> int:
+    return ((int(value) + int(alignment) - 1) // int(alignment)) * int(alignment)
+
+
+def _dtype_nbytes(dtype: torch.dtype) -> int:
+    return torch.empty((), dtype=dtype).element_size()
+
+
+def _normalize_tile_config(value: Sequence[int]) -> tuple[int, int, int, int]:
+    if len(value) != 4:
+        raise ValueError(
+            "tile_config must be (fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n)"
+        )
+    tile_config = tuple(int(item) for item in value)
+    if any(item <= 0 or item % 16 != 0 for item in tile_config):
+        raise ValueError(
+            "every tile_config value must be a positive multiple of 16, got "
+            f"{tile_config}"
+        )
+    fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n = tile_config
+    if fc1_tile_k * fc1_tile_n != fc2_tile_k * fc2_tile_n:
+        raise ValueError(
+            "FC1 and FC2 tile_config entries must select the same CTA thread count"
+        )
+    return tile_config
+
+
+def _normalize_input_dtype(dtype: torch.dtype) -> torch.dtype:
+    if dtype not in (torch.bfloat16, torch.float16):
+        raise TypeError(
+            "Trellis MoE input_dtype must be torch.bfloat16 or torch.float16, "
+            f"got {dtype}"
+        )
+    return dtype
+
+
+def _input_dtype_name(dtype: torch.dtype) -> str:
+    return "bf16" if dtype == torch.bfloat16 else "fp16"
+
+
+def _resolve_cuda_device(device: torch.device | str | int) -> torch.device:
+    if isinstance(device, int):
+        result = torch.device("cuda", int(device))
+    else:
+        result = torch.device(device)
+    if result.type != "cuda":
+        raise ValueError(f"Trellis MoE requires a CUDA device, got {result}")
+    if result.index is None and torch.cuda.is_available():
+        result = torch.device("cuda", torch.cuda.current_device())
+    return result
+
+
+@dataclass(frozen=True, kw_only=True)
+class TrellisMoECaps:
+    """Fixed serving capacity and compile policy for one Trellis MoE shape."""
+
+    max_tokens: int
+    num_topk: int
+    num_experts: int
+    hidden_size: int
+    intermediate_size: int
+    device: torch.device | str | int
+    input_dtype: torch.dtype = torch.bfloat16
+    route_num_experts: int | None = None
+    block_size_m: int = 8
+    trellis_bits: int = 3
+    tile_config: tuple[int, int, int, int] = _DEFAULT_TILE_CONFIG
+    activation: str = "silu"
+    fast_math: bool = True
+
+    def __post_init__(self) -> None:
+        for name in (
+            "max_tokens",
+            "num_topk",
+            "num_experts",
+            "hidden_size",
+            "intermediate_size",
+        ):
+            value = int(getattr(self, name))
+            if value <= 0:
+                raise ValueError(f"{name} must be positive, got {value}")
+            object.__setattr__(self, name, value)
+        if self.hidden_size % 128 != 0 or self.intermediate_size % 128 != 0:
+            raise ValueError(
+                "full-rotation Trellis MoE requires hidden_size and "
+                "intermediate_size divisible by 128"
+            )
+        object.__setattr__(self, "device", _resolve_cuda_device(self.device))
+        object.__setattr__(
+            self, "input_dtype", _normalize_input_dtype(self.input_dtype)
+        )
+        route_num_experts = (
+            self.num_experts
+            if self.route_num_experts is None
+            else int(self.route_num_experts)
+        )
+        if route_num_experts <= 0:
+            raise ValueError(
+                f"route_num_experts must be positive, got {route_num_experts}"
+            )
+        if self.num_topk > route_num_experts:
+            raise ValueError(
+                f"num_topk={self.num_topk} exceeds route_num_experts={route_num_experts}"
+            )
+        object.__setattr__(self, "route_num_experts", route_num_experts)
+        block_size_m = int(self.block_size_m)
+        if block_size_m not in _ALLOWED_BLOCK_M:
+            raise ValueError(
+                f"block_size_m must be one of {_ALLOWED_BLOCK_M}, got {block_size_m}"
+            )
+        object.__setattr__(self, "block_size_m", block_size_m)
+        trellis_bits = int(self.trellis_bits)
+        if trellis_bits not in _ALLOWED_BITS:
+            raise ValueError(
+                f"trellis_bits must be one of {_ALLOWED_BITS}, got {trellis_bits}"
+            )
+        object.__setattr__(self, "trellis_bits", trellis_bits)
+        tile_config = _normalize_tile_config(self.tile_config)
+        fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n = tile_config
+        if self.hidden_size % fc1_tile_k != 0:
+            raise ValueError("hidden_size must be divisible by FC1 tile K")
+        if self.intermediate_size % fc1_tile_n != 0:
+            raise ValueError(
+                "projection-major intermediate_size must be divisible by FC1 tile N"
+            )
+        if self.intermediate_size % fc2_tile_k != 0:
+            raise ValueError("intermediate_size must be divisible by FC2 tile K")
+        if self.hidden_size % fc2_tile_n != 0:
+            raise ValueError("hidden_size must be divisible by FC2 tile N")
+        object.__setattr__(self, "tile_config", tile_config)
+        if str(self.activation).strip().lower() != "silu":
+            raise ValueError("the validated full-rotation Trellis recipe requires silu")
+        object.__setattr__(self, "activation", "silu")
+        object.__setattr__(self, "fast_math", bool(self.fast_math))
+
+    @property
+    def is_gated(self) -> bool:
+        return True
+
+
+@dataclass(frozen=True, eq=False)
+class TrellisMoEWeights:
+    """Zero-copy native tensors and persistent full-rotation tables."""
+
+    w13: torch.Tensor
+    w2: torch.Tensor
+    gate_suh: torch.Tensor
+    up_suh: torch.Tensor
+    intermediate_rotations: torch.Tensor
+    down_svh: torch.Tensor
+    hidden_size: int
+    intermediate_size: int
+    num_experts: int
+    trellis_bits: int
+    tile_config: tuple[int, int, int, int]
+    device: torch.device
+    _prepared: PreparedNF3MoeWeights = field(repr=False)
+
+
+@dataclass(frozen=True)
+class _ArenaViewSpec:
+    name: str
+    offset_bytes: int
+    shape: tuple[int, ...]
+    dtype: torch.dtype
+
+    @property
+    def nbytes(self) -> int:
+        return int(prod(self.shape)) * _dtype_nbytes(self.dtype)
+
+
+@dataclass(frozen=True)
+class _ArenaLayout:
+    nbytes: int
+    views: tuple[_ArenaViewSpec, ...]
+
+    def materialize(self, scratch: torch.Tensor) -> dict[str, torch.Tensor]:
+        result: dict[str, torch.Tensor] = {}
+        for spec in self.views:
+            raw = scratch.narrow(0, spec.offset_bytes, spec.nbytes)
+            result[spec.name] = raw.view(spec.dtype).view(spec.shape)
+        return result
+
+
+def _make_arena_layout(
+    caps: TrellisMoECaps,
+    buffers: W4A16BufferPlan,
+    *,
+    sms: int,
+) -> _ArenaLayout:
+    assert caps.route_num_experts is not None
+    specs = (
+        (
+            "intermediate_cache13",
+            (buffers.intermediate_cache13_elements,),
+            torch.float16,
+        ),
+        ("intermediate_cache2", (buffers.intermediate_cache2_elements,), torch.float16),
+        ("output", (caps.max_tokens, caps.hidden_size), torch.float32),
+        ("fc1_c_tmp", (buffers.fc1_c_tmp_elements,), torch.float32),
+        ("fc2_c_tmp", (buffers.fc2_c_tmp_elements,), torch.float32),
+        ("packed_route_indices", (buffers.route_slots,), torch.int32),
+        ("block_expert_ids", (buffers.route_blocks,), torch.int32),
+        ("packed_route_count", (1,), torch.int32),
+        ("expert_offsets", (caps.route_num_experts + 1,), torch.int32),
+        ("expert_counts", (caps.route_num_experts,), torch.int32),
+        ("rotation_a_gate", (buffers.rotation_a_elements,), torch.float16),
+        ("rotation_a_up", (buffers.rotation_a_elements,), torch.float16),
+        ("kernel_workspace", (int(sms) * 4 + 2,), torch.int32),
+    )
+    cursor = 0
+    views: list[_ArenaViewSpec] = []
+    for name, shape, dtype in specs:
+        cursor = _align_up(cursor, max(_ARENA_ALIGNMENT, _dtype_nbytes(dtype)))
+        spec = _ArenaViewSpec(
+            name=name,
+            offset_bytes=cursor,
+            shape=tuple(int(dim) for dim in shape),
+            dtype=dtype,
+        )
+        views.append(spec)
+        cursor += spec.nbytes
+    return _ArenaLayout(nbytes=max(_align_up(cursor), 1), views=tuple(views))
+
+
+@dataclass(frozen=True)
+class TrellisMoEPlan:
+    """Compiled launches and byte-arena layout for one fixed capacity."""
+
+    caps: TrellisMoECaps
+    buffer_plan: W4A16BufferPlan
+    fused_launch: W4A16FusedMoeCompileResult = field(repr=False)
+    identity_sums: tuple[W4A16TopKSumCompileResult, ...] = field(repr=False)
+    mapped_sums: tuple[W4A16TopKSumCompileResult, ...] = field(repr=False)
+    _arena_layout: _ArenaLayout = field(repr=False)
+    _scratch_specs: tuple[ScratchBufferSpec, ...] = field(repr=False)
+
+    @property
+    def scratch_nbytes(self) -> int:
+        return self._arena_layout.nbytes
+
+    def scratch_specs(self) -> tuple[ScratchBufferSpec, ...]:
+        return self._scratch_specs
+
+    def shapes_and_dtypes(self) -> tuple[tuple[tuple[int, ...], torch.dtype], ...]:
+        return tuple((spec.shape, spec.dtype) for spec in self._scratch_specs)
+
+    def bind(self, **kwargs) -> "TrellisMoEBinding":
+        return bind_trellis_moe(self, **kwargs)
+
+    def topk_sum_launch(
+        self, ids_dtype: torch.dtype, *, mapped: bool
+    ) -> W4A16TopKSumCompileResult:
+        launches = self.mapped_sums if mapped else self.identity_sums
+        for launch in launches:
+            if launch.route_ids_dtype == ids_dtype:
+                return launch
+        raise TypeError(f"Trellis MoE does not have a top-k sum for {ids_dtype}")
+
+
+@dataclass(frozen=True, kw_only=True)
+class TrellisMoEBinding:
+    """Stable tensor views for one launch; safe to retain across graph replay."""
+
+    plan: TrellisMoEPlan
+    weights: TrellisMoEWeights
+    a: torch.Tensor
+    topk_weights: torch.Tensor
+    topk_ids: torch.Tensor
+    output: torch.Tensor
+    route_expert_map: torch.Tensor | None
+    output_expert_map: torch.Tensor | None
+    prepared: PreparedNF3MoeWeights = field(repr=False)
+    intermediate_cache13: torch.Tensor = field(repr=False)
+    intermediate_cache2: torch.Tensor = field(repr=False)
+    fc1_c_tmp: torch.Tensor = field(repr=False)
+    fc2_c_tmp: torch.Tensor = field(repr=False)
+    packed_route_indices: torch.Tensor = field(repr=False)
+    block_expert_ids: torch.Tensor = field(repr=False)
+    packed_route_count: torch.Tensor = field(repr=False)
+    expert_offsets: torch.Tensor = field(repr=False)
+    expert_counts: torch.Tensor = field(repr=False)
+    rotation_a_gate: torch.Tensor = field(repr=False)
+    rotation_a_up: torch.Tensor = field(repr=False)
+    topk_sum_launch: W4A16TopKSumCompileResult = field(repr=False)
+
+    def run(self) -> torch.Tensor:
+        return run_trellis_moe(binding=self)
+
+
+def _validate_rotation_table(
+    name: str,
+    tensor: torch.Tensor,
+    *,
+    shape: tuple[int, ...],
+    device: torch.device,
+) -> None:
+    if tensor.dtype != torch.float16:
+        raise TypeError(f"{name} must be torch.float16, got {tensor.dtype}")
+    if tensor.device != device:
+        raise ValueError(f"{name} must be on {device}, got {tensor.device}")
+    if tuple(tensor.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
+    if not tensor.is_contiguous():
+        raise ValueError(f"{name} must be contiguous")
+
+
+def _validate_mcg(codebook: str | int, mcg: torch.Tensor | int | None) -> None:
+    normalized = _normalize_trellis256_codebook(codebook)
+    if normalized != "mcg":
+        raise NotImplementedError(
+            "the production Trellis MoE decoder accepts only the MCG codebook"
+        )
+    if mcg is None:
+        return
+    if isinstance(mcg, torch.Tensor):
+        if mcg.numel() != 1 or mcg.dtype not in (torch.int32, torch.uint32):
+            raise ValueError("mcg must be a scalar int32/uint32 tensor")
+        marker = int(mcg.item()) & 0xFFFFFFFF
+    else:
+        marker = int(mcg) & 0xFFFFFFFF
+    if marker != _MCG_SENTINEL:
+        raise ValueError(
+            f"unexpected MCG marker {marker:#010x}; expected {_MCG_SENTINEL:#010x}"
+        )
+
+
+def prepare_trellis_moe_weights(
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+    *,
+    gate_suh: torch.Tensor,
+    up_suh: torch.Tensor,
+    intermediate_rotations: torch.Tensor,
+    down_svh: torch.Tensor,
+    codebook: str | int = "mcg",
+    mcg: torch.Tensor | int | None = None,
+    tile_config: tuple[int, int, int, int] = _DEFAULT_TILE_CONFIG,
+    dummy_scale: torch.Tensor | None = None,
+) -> TrellisMoEWeights:
+    """Validate and wrap projection-major native EXL3 tensors without copying."""
+    _validate_mcg(codebook, mcg)
+    tile_config = _normalize_tile_config(tile_config)
+    if w13.ndim != 5 or int(w13.shape[0]) != 2:
+        raise ValueError(
+            "w13 must be projection-major [2,E,H/16,I/16,16*bits] int16 "
+            "or the byte-identical [...,8*bits] int32 view"
+        )
+    if w2.ndim != 4:
+        raise ValueError(
+            "w2 must be [E,I/16,H/16,16*bits] int16 or the byte-identical "
+            "[...,8*bits] int32 view"
+        )
+    num_experts = int(w13.shape[1])
+    hidden_size = int(w13.shape[2]) * 16
+    intermediate_size = int(w13.shape[3]) * 16
+    if tuple(w2.shape[:3]) != (
+        num_experts,
+        intermediate_size // 16,
+        hidden_size // 16,
+    ):
+        raise ValueError(
+            "w2 geometry does not match projection-major w13: "
+            f"got {tuple(w2.shape[:3])}"
+        )
+    if hidden_size % 128 != 0 or intermediate_size % 128 != 0:
+        raise ValueError(
+            "full-rotation Trellis weights require hidden and intermediate "
+            "dimensions divisible by 128"
+        )
+    fc1_tile_k, fc1_tile_n, fc2_tile_k, fc2_tile_n = tile_config
+    if (
+        hidden_size % fc1_tile_k != 0
+        or intermediate_size % fc1_tile_n != 0
+        or intermediate_size % fc2_tile_k != 0
+        or hidden_size % fc2_tile_n != 0
+    ):
+        raise ValueError(
+            f"tile_config={tile_config} does not divide H={hidden_size}, "
+            f"I={intermediate_size}"
+        )
+    device = w13.device
+    if w2.device != device:
+        raise ValueError("w13 and w2 must be on the same device")
+    _validate_rotation_table(
+        "gate_suh", gate_suh, shape=(num_experts, hidden_size), device=device
+    )
+    _validate_rotation_table(
+        "up_suh", up_suh, shape=(num_experts, hidden_size), device=device
+    )
+    _validate_rotation_table(
+        "intermediate_rotations",
+        intermediate_rotations,
+        shape=(num_experts, 3 * intermediate_size),
+        device=device,
+    )
+    _validate_rotation_table(
+        "down_svh", down_svh, shape=(num_experts, hidden_size), device=device
+    )
+    prepared = prepare_trellis256_moe_weights(
+        w13,
+        w2,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        num_experts=num_experts,
+        activation="silu",
+        params_dtype=torch.float16,
+        fc1_tile_n=fc1_tile_n,
+        fc2_tile_n=fc2_tile_n,
+        w13_layout="trellis3_t256_proj",
+        codebook=codebook,
+        dummy_scale=dummy_scale,
+        gate_suh=gate_suh,
+        up_suh=up_suh,
+    )
+    if prepared.trellis_codebook != "mcg":
+        raise RuntimeError("Trellis preparation did not preserve the MCG contract")
+    if (
+        prepared.w13.data_ptr() != w13.data_ptr()
+        or prepared.w2.data_ptr() != w2.data_ptr()
+    ):
+        raise RuntimeError("Trellis preparation unexpectedly copied native weights")
+    return TrellisMoEWeights(
+        w13=w13,
+        w2=w2,
+        gate_suh=gate_suh,
+        up_suh=up_suh,
+        intermediate_rotations=intermediate_rotations,
+        down_svh=down_svh,
+        hidden_size=hidden_size,
+        intermediate_size=intermediate_size,
+        num_experts=num_experts,
+        trellis_bits=int(prepared.trellis_bits),
+        tile_config=tile_config,
+        device=device,
+        _prepared=prepared,
+    )
+
+
+def plan_trellis_moe(caps: TrellisMoECaps) -> TrellisMoEPlan:
+    """Compile every launch and produce the fixed caller-scratch layout."""
+    if not isinstance(caps, TrellisMoECaps):
+        raise TypeError("caps must be a TrellisMoECaps")
+    device = _resolve_cuda_device(caps.device)
+    if device.index is None:
+        raise RuntimeError("CUDA must be available before planning Trellis MoE")
+    if device != caps.device:
+        caps = replace(caps, device=device)
+    with torch.cuda.device(device):
+        props = torch.cuda.get_device_properties(device)
+        sms = int(props.multi_processor_count)
+        max_shared_mem = int(getattr(props, "shared_memory_per_block_optin", 101_376))
+        buffer_plan = plan_w4a16_buffers(
+            caps,
+            m=caps.max_tokens,
+            topk=caps.num_topk,
+            route_num_experts=caps.route_num_experts,
+            sms=sms,
+            full_rotation=True,
+            block_size_m=caps.block_size_m,
+        )
+        route_slots = max_packed_route_slots(
+            caps.max_tokens * caps.num_topk,
+            caps.block_size_m,
+            caps.route_num_experts,
+        )
+        max_m_blocks = (route_slots + caps.block_size_m - 1) // caps.block_size_m
+        fused_launch = compile_w4a16_fused_moe(
+            size_m=caps.max_tokens,
+            hidden_size=caps.hidden_size,
+            intermediate_size=caps.intermediate_size,
+            num_experts=caps.num_experts,
+            top_k=caps.num_topk,
+            activation=caps.activation,
+            apply_router_weight_on_input=False,
+            zero_fc2_output=False,
+            moe_block_size=caps.block_size_m,
+            max_m_blocks=max_m_blocks,
+            element_dtype="fp16",
+            fast_math=caps.fast_math,
+            sms=sms,
+            max_shared_mem=max_shared_mem,
+            weight_layout="trellis3_t256",
+            scale_format="e4m3_k32",
+            w13_layout="trellis3_t256_proj",
+            trellis_bits=caps.trellis_bits,
+            force_tile_config=caps.tile_config,
+            intermediate_rotation=True,
+            full_rotation=True,
+            rotation_input_dtype=_input_dtype_name(caps.input_dtype),
+        )
+        identity_sums = tuple(
+            compile_w4a16_topk_sum(
+                m=caps.max_tokens,
+                topk=caps.num_topk,
+                hidden_size=caps.hidden_size,
+                element_dtype="fp16",
+                full_rotation=True,
+                num_experts=caps.num_experts,
+                route_num_experts=0,
+                route_ids_dtype=ids_dtype,
+                use_expert_map=False,
+            )
+            for ids_dtype in (torch.int32, torch.int64)
+        )
+        mapped_sums = tuple(
+            compile_w4a16_topk_sum(
+                m=caps.max_tokens,
+                topk=caps.num_topk,
+                hidden_size=caps.hidden_size,
+                element_dtype="fp16",
+                full_rotation=True,
+                num_experts=caps.num_experts,
+                route_num_experts=caps.route_num_experts,
+                route_ids_dtype=ids_dtype,
+                use_expert_map=True,
+            )
+            for ids_dtype in (torch.int32, torch.int64)
+        )
+    arena_layout = _make_arena_layout(caps, buffer_plan, sms=sms)
+    specs = (
+        scratch_buffer_spec(
+            "trellis_moe",
+            nbytes=arena_layout.nbytes,
+            device=device,
+        ),
+    )
+    return TrellisMoEPlan(
+        caps=caps,
+        buffer_plan=buffer_plan,
+        fused_launch=fused_launch,
+        identity_sums=identity_sums,
+        mapped_sums=mapped_sums,
+        _arena_layout=arena_layout,
+        _scratch_specs=specs,
+    )
+
+
+def _validate_runtime_tensor(
+    name: str,
+    tensor: torch.Tensor,
+    *,
+    shape: tuple[int, ...],
+    dtype: torch.dtype,
+    device: torch.device,
+) -> None:
+    if tensor.dtype != dtype:
+        raise TypeError(f"{name} must be {dtype}, got {tensor.dtype}")
+    if tensor.device != device:
+        raise ValueError(f"{name} must be on {device}, got {tensor.device}")
+    if tuple(tensor.shape) != shape:
+        raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}")
+    if not tensor.is_contiguous():
+        raise ValueError(f"{name} must be contiguous")
+
+
+def _validate_expert_map(
+    name: str,
+    tensor: torch.Tensor | None,
+    *,
+    caps: TrellisMoECaps,
+) -> None:
+    if tensor is None:
+        return
+    assert caps.route_num_experts is not None
+    _validate_runtime_tensor(
+        name,
+        tensor,
+        shape=(caps.route_num_experts,),
+        dtype=torch.int32,
+        device=caps.device,
+    )
+
+
+def bind_trellis_moe(
+    plan: TrellisMoEPlan,
+    *,
+    scratch: torch.Tensor | Mapping[str, torch.Tensor] | Sequence[torch.Tensor],
+    a: torch.Tensor,
+    weights: TrellisMoEWeights,
+    topk_weights: torch.Tensor,
+    topk_ids: torch.Tensor,
+    route_expert_map: torch.Tensor | None = None,
+    output_expert_map: torch.Tensor | None = None,
+    output: torch.Tensor | None = None,
+) -> TrellisMoEBinding:
+    """Bind runtime tensors by carving views from ``scratch`` only."""
+    if not isinstance(plan, TrellisMoEPlan):
+        raise TypeError("plan must come from trellis_moe.plan")
+    if not isinstance(weights, TrellisMoEWeights):
+        raise TypeError("weights must come from trellis_moe.prepare_weights")
+    caps = plan.caps
+    if (
+        weights.hidden_size != caps.hidden_size
+        or weights.intermediate_size != caps.intermediate_size
+        or weights.num_experts != caps.num_experts
+        or weights.trellis_bits != caps.trellis_bits
+        or weights.tile_config != caps.tile_config
+        or weights.device != caps.device
+    ):
+        raise ValueError("weights do not match the Trellis MoE plan")
+    if a.ndim != 2:
+        raise ValueError(f"a must be rank 2, got shape {tuple(a.shape)}")
+    tokens = int(a.shape[0])
+    if tokens < 1 or tokens > caps.max_tokens:
+        raise ValueError(
+            f"input tokens must be in [1, {caps.max_tokens}], got {tokens}"
+        )
+    _validate_runtime_tensor(
+        "a",
+        a,
+        shape=(tokens, caps.hidden_size),
+        dtype=caps.input_dtype,
+        device=caps.device,
+    )
+    _validate_runtime_tensor(
+        "topk_weights",
+        topk_weights,
+        shape=(tokens, caps.num_topk),
+        dtype=torch.float32,
+        device=caps.device,
+    )
+    if topk_ids.dtype not in (torch.int32, torch.int64):
+        raise TypeError("topk_ids must be torch.int32 or torch.int64")
+    _validate_runtime_tensor(
+        "topk_ids",
+        topk_ids,
+        shape=(tokens, caps.num_topk),
+        dtype=topk_ids.dtype,
+        device=caps.device,
+    )
+    _validate_expert_map("route_expert_map", route_expert_map, caps=caps)
+    _validate_expert_map("output_expert_map", output_expert_map, caps=caps)
+
+    scratch_storage = scratch_tensor(scratch, plan._scratch_specs, owner="Trellis MoE")
+    if int(scratch_storage.data_ptr()) % _ARENA_ALIGNMENT != 0:
+        raise ValueError(f"Trellis MoE scratch must be {_ARENA_ALIGNMENT}-byte aligned")
+    views = plan._arena_layout.materialize(scratch_storage)
+    views["kernel_workspace"].zero_()
+    if output is None:
+        output_view = views["output"][:tokens]
+    else:
+        if output.ndim != 2 or tuple(output.shape) not in (
+            (tokens, caps.hidden_size),
+            (caps.max_tokens, caps.hidden_size),
+        ):
+            raise ValueError(
+                "output must be the live or capacity FP32 view: expected "
+                f"{(tokens, caps.hidden_size)} or {(caps.max_tokens, caps.hidden_size)}, "
+                f"got {tuple(output.shape)}"
+            )
+        if output.dtype != torch.float32:
+            raise TypeError(f"output must be torch.float32, got {output.dtype}")
+        if output.device != caps.device or not output.is_contiguous():
+            raise ValueError("output must be contiguous on the planned CUDA device")
+        output_view = output[:tokens]
+    prepared = replace(weights._prepared, workspace=views["kernel_workspace"])
+    return TrellisMoEBinding(
+        plan=plan,
+        weights=weights,
+        a=a,
+        topk_weights=topk_weights,
+        topk_ids=topk_ids,
+        output=output_view,
+        route_expert_map=route_expert_map,
+        output_expert_map=output_expert_map,
+        prepared=prepared,
+        intermediate_cache13=views["intermediate_cache13"],
+        intermediate_cache2=views["intermediate_cache2"],
+        fc1_c_tmp=views["fc1_c_tmp"],
+        fc2_c_tmp=views["fc2_c_tmp"],
+        packed_route_indices=views["packed_route_indices"],
+        block_expert_ids=views["block_expert_ids"],
+        packed_route_count=views["packed_route_count"],
+        expert_offsets=views["expert_offsets"],
+        expert_counts=views["expert_counts"],
+        rotation_a_gate=views["rotation_a_gate"],
+        rotation_a_up=views["rotation_a_up"],
+        topk_sum_launch=plan.topk_sum_launch(
+            topk_ids.dtype, mapped=output_expert_map is not None
+        ),
+    )
+
+
+def run_trellis_moe(*, binding: TrellisMoEBinding) -> torch.Tensor:
+    """Run the preplanned full-rotation path into ``binding.output``."""
+    if not isinstance(binding, TrellisMoEBinding):
+        raise TypeError("binding must come from trellis_moe.bind")
+    caps = binding.plan.caps
+    return run_w4a16_moe(
+        binding.a,
+        binding.prepared,
+        binding.topk_weights,
+        binding.topk_ids,
+        activation=caps.activation,
+        intermediate_cache13=binding.intermediate_cache13,
+        intermediate_cache2=binding.intermediate_cache2,
+        output=binding.output,
+        fc1_c_tmp=binding.fc1_c_tmp,
+        fc2_c_tmp=binding.fc2_c_tmp,
+        packed_route_indices=binding.packed_route_indices,
+        block_expert_ids=binding.block_expert_ids,
+        packed_route_count=binding.packed_route_count,
+        expert_offsets=binding.expert_offsets,
+        expert_counts=binding.expert_counts,
+        expert_map=binding.route_expert_map,
+        output_expert_map=binding.output_expert_map,
+        apply_router_weight_on_input=False,
+        fast_math=caps.fast_math,
+        fused_launch=binding.plan.fused_launch,
+        topk_sum_launch=binding.topk_sum_launch,
+        intermediate_rotation_scales=binding.weights.intermediate_rotations,
+        full_rotation=True,
+        suh_gate_table=binding.weights.gate_suh,
+        suh_up_table=binding.weights.up_suh,
+        svh_table=binding.weights.down_svh,
+        rotation_a_gate=binding.rotation_a_gate,
+        rotation_a_up=binding.rotation_a_up,
+    )
+
+
+def clear_trellis_moe_caches() -> None:
+    """Clear the shared W4A16 compile caches used by this planned op."""
+    clear_w4a16_kernel_cache()
+
+
+__all__ = [
+    "TrellisMoEBinding",
+    "TrellisMoECaps",
+    "TrellisMoEPlan",
+    "TrellisMoEWeights",
+    "bind_trellis_moe",
+    "clear_trellis_moe_caches",
+    "plan_trellis_moe",
+    "prepare_trellis_moe_weights",
+    "run_trellis_moe",
+]
diff --git a/sparkinfer/moe/trellis_moe/api.py b/sparkinfer/moe/trellis_moe/api.py
new file mode 100644
index 00000000..f066cb50
--- /dev/null
+++ b/sparkinfer/moe/trellis_moe/api.py
@@ -0,0 +1,52 @@
+"""Public surface for :mod:`sparkinfer.moe.trellis_moe`."""
+
+from __future__ import annotations
+
+from ..._lib.gating import default_is_supported
+from . import META
+from ._impl import (
+    TrellisMoEBinding as Binding,
+)
+from ._impl import (
+    TrellisMoECaps as Caps,
+)
+from ._impl import (
+    TrellisMoEPlan as Plan,
+)
+from ._impl import (
+    TrellisMoEWeights as Weights,
+)
+from ._impl import (
+    bind_trellis_moe as bind,
+)
+from ._impl import (
+    clear_trellis_moe_caches as clear_caches,
+)
+from ._impl import (
+    plan_trellis_moe as plan,
+)
+from ._impl import (
+    prepare_trellis_moe_weights as prepare_weights,
+)
+from ._impl import (
+    run_trellis_moe as run,
+)
+
+
+def is_supported(device=None) -> bool:
+    """Return whether ``device`` supports the SM12x Trellis kernel stack."""
+    return default_is_supported(device, requires=META.requires)
+
+
+__all__ = [
+    "Caps",
+    "Plan",
+    "Weights",
+    "Binding",
+    "plan",
+    "prepare_weights",
+    "bind",
+    "run",
+    "is_supported",
+    "clear_caches",
+]
diff --git a/tests/moe/test_trellis_moe.py b/tests/moe/test_trellis_moe.py
new file mode 100644
index 00000000..95700554
--- /dev/null
+++ b/tests/moe/test_trellis_moe.py
@@ -0,0 +1,471 @@
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+import torch
+
+pytest.importorskip("cutlass")
+
+from sparkinfer.moe import trellis_moe
+from sparkinfer.moe._shared.kernels.w4a16.host import plan_w4a16_buffers
+
+
+_MCG = np.uint64(0xCBAC1FED)
+_MASK = np.uint32(0x8FFF8FFF)
+_ORC = np.uint32(0x3B603B60)
+
+
+def _caps(**overrides) -> trellis_moe.Caps:
+    values = {
+        "max_tokens": 32,
+        "num_topk": 8,
+        "num_experts": 192,
+        "hidden_size": 6144,
+        "intermediate_size": 512,
+        "route_num_experts": 256,
+        "block_size_m": 8,
+        "input_dtype": torch.bfloat16,
+        "device": "cuda:0",
+    }
+    values.update(overrides)
+    return trellis_moe.Caps(**values)
+
+
+def test_caps_preserve_exact_block_m_and_input_dtypes() -> None:
+    bf16 = _caps(block_size_m=8, input_dtype=torch.bfloat16)
+    fp16 = _caps(block_size_m=64, input_dtype=torch.float16)
+
+    assert bf16.block_size_m == 8
+    assert bf16.route_num_experts == 256
+    assert bf16.input_dtype == torch.bfloat16
+    assert fp16.block_size_m == 64
+    assert fp16.input_dtype == torch.float16
+
+
+@pytest.mark.parametrize("block_size_m", [0, 4, 12, 128])
+def test_caps_reject_unsupported_block_m(block_size_m: int) -> None:
+    with pytest.raises(ValueError, match="block_size_m"):
+        _caps(block_size_m=block_size_m)
+
+
+def test_low_level_buffer_plan_honors_explicit_block_m() -> None:
+    prepared = SimpleNamespace(
+        num_experts=256,
+        hidden_size=6144,
+        intermediate_size=512,
+        is_gated=True,
+    )
+    automatic = plan_w4a16_buffers(
+        prepared,
+        m=1024,
+        topk=8,
+        route_num_experts=256,
+        sms=120,
+    )
+    exact = plan_w4a16_buffers(
+        prepared,
+        m=1024,
+        topk=8,
+        route_num_experts=256,
+        sms=120,
+        full_rotation=True,
+        block_size_m=8,
+    )
+
+    assert automatic.block_size_m != 8
+    assert exact.block_size_m == 8
+    assert exact.rotation_a_elements == 1024 * 8 * 6144
+    with pytest.raises(ValueError, match="block_size_m"):
+        plan_w4a16_buffers(
+            prepared,
+            m=1,
+            topk=8,
+            route_num_experts=256,
+            sms=120,
+            block_size_m=4,
+        )
+
+
+def _cpu_weight_tensors() -> tuple[torch.Tensor, ...]:
+    w13 = torch.zeros((2, 1, 8, 8, 48), dtype=torch.int16)
+    w2 = torch.zeros((1, 8, 8, 48), dtype=torch.int16)
+    edge = torch.ones((1, 128), dtype=torch.float16)
+    intermediate = torch.ones((1, 384), dtype=torch.float16)
+    return w13, w2, edge, edge.clone(), intermediate, edge.clone()
+
+
+def test_prepare_weights_rejects_non_mcg_before_cuda_work() -> None:
+    w13, w2, gate_suh, up_suh, intermediate, down_svh = _cpu_weight_tensors()
+    with pytest.raises(NotImplementedError, match="only the MCG"):
+        trellis_moe.prepare_weights(
+            w13,
+            w2,
+            gate_suh=gate_suh,
+            up_suh=up_suh,
+            intermediate_rotations=intermediate,
+            down_svh=down_svh,
+            codebook="mul1",
+            tile_config=(64, 128, 64, 128),
+        )
+    with pytest.raises(ValueError, match="unexpected MCG marker"):
+        trellis_moe.prepare_weights(
+            w13,
+            w2,
+            gate_suh=gate_suh,
+            up_suh=up_suh,
+            intermediate_rotations=intermediate,
+            down_svh=down_svh,
+            mcg=0x83DCD12D,
+            tile_config=(64, 128, 64, 128),
+        )
+
+
+def test_prepare_weights_validates_all_rotation_shapes() -> None:
+    w13, w2, gate_suh, up_suh, intermediate, down_svh = _cpu_weight_tensors()
+    with pytest.raises(ValueError, match="intermediate_rotations must have shape"):
+        trellis_moe.prepare_weights(
+            w13,
+            w2,
+            gate_suh=gate_suh,
+            up_suh=up_suh,
+            intermediate_rotations=intermediate[:, :-1].contiguous(),
+            down_svh=down_svh,
+            tile_config=(64, 128, 64, 128),
+        )
+
+
+def _decode_3inst_fp16(window: np.ndarray) -> np.ndarray:
+    value = window.astype(np.uint64)
+    value = ((value * _MCG) & np.uint64(0xFFFFFFFF)).astype(np.uint32)
+    value = np.uint32((value & _MASK) ^ _ORC)
+    low = (value & np.uint32(0xFFFF)).astype(np.uint16).view(np.float16)
+    high = (
+        ((value >> np.uint32(16)) & np.uint32(0xFFFF))
+        .astype(np.uint16)
+        .view(np.float16)
+    )
+    return (low.astype(np.float16) + high.astype(np.float16)).astype(np.float16)
+
+
+def _decode_lane(tile_words: np.ndarray, lane: int, bits: int) -> np.ndarray:
+    width = 8 * bits
+    values = []
+    for weight in range(8):
+        end_bit = (lane * 8 + weight + 257) * bits
+        start_bit = end_bit - 16
+        first_word = start_bit // 32
+        last_word = (end_bit - 1) // 32
+        shift = (last_word + 1) * 32 - end_bit
+        first = tile_words[..., first_word % width].astype(np.uint64)
+        last = tile_words[..., last_word % width].astype(np.uint64)
+        merged = (first << np.uint64(32)) | last
+        window = ((merged >> np.uint64(shift)) & np.uint64(0xFFFF)).astype(np.uint32)
+        values.append(_decode_3inst_fp16(window))
+    return np.stack(values, axis=-1).astype(np.float16)
+
+
+def _reconstruct_native(trellis: torch.Tensor) -> torch.Tensor:
+    native = trellis.detach().cpu().numpy()
+    bits = int(native.shape[-1]) // 16
+    k_tiles, n_tiles, _ = native.shape
+    packed = native.view(np.uint16).reshape(k_tiles, n_tiles, 8 * bits, 2)
+    words = packed[..., 0].astype(np.uint32) | (
+        packed[..., 1].astype(np.uint32) << np.uint32(16)
+    )
+    output = np.zeros((k_tiles * 16, n_tiles * 16), dtype=np.float16)
+    for k_tile in range(k_tiles):
+        for n_tile in range(n_tiles):
+            lanes = np.stack(
+                [_decode_lane(words[k_tile, n_tile], lane, bits) for lane in range(32)]
+            )
+            block = np.zeros((16, 16), dtype=np.float16)
+            for lane in range(32):
+                row0 = (lane % 4) * 2
+                rows = (row0, row0 + 1, row0 + 8, row0 + 9)
+                col0 = lane // 8
+                col1 = col0 + 4
+                parity = (lane >> 2) & 1
+                for weight in range(8):
+                    block[
+                        rows[weight % 4], 2 * (col0 if weight < 4 else col1) + parity
+                    ] = lanes[lane, weight]
+            output[
+                k_tile * 16 : (k_tile + 1) * 16,
+                n_tile * 16 : (n_tile + 1) * 16,
+            ] = block
+    return torch.from_numpy(output)
+
+
+def _hadamard_128(device: torch.device) -> torch.Tensor:
+    indices = torch.arange(128, dtype=torch.int64)
+    rows = []
+    for row in range(128):
+        parity = torch.tensor(
+            [(row & int(col)).bit_count() & 1 for col in indices],
+            dtype=torch.bool,
+        )
+        rows.append(torch.where(parity, -1.0, 1.0))
+    return (torch.stack(rows) / (128.0**0.5)).to(device)
+
+
+def _had128(
+    value: torch.Tensor,
+    hadamard: torch.Tensor,
+    *,
+    suh: torch.Tensor | None = None,
+    svh: torch.Tensor | None = None,
+    store_fp16: bool,
+) -> torch.Tensor:
+    work = value.float()
+    if suh is not None:
+        work = (work * suh.float()).to(torch.float16).float()
+    rows, width = work.shape
+    work = (work.view(rows, width // 128, 128) @ hadamard).view(rows, width)
+    if svh is not None:
+        work = work * svh.float()
+    return work.to(torch.float16) if store_fp16 else work
+
+
+def _reference_full_rotation(
+    x: torch.Tensor,
+    local_ids: torch.Tensor,
+    router_weights: torch.Tensor,
+    w13: torch.Tensor,
+    w2: torch.Tensor,
+    gate_suh: torch.Tensor,
+    up_suh: torch.Tensor,
+    intermediate_rotations: torch.Tensor,
+    down_svh: torch.Tensor,
+) -> torch.Tensor:
+    device = x.device
+    experts = int(w2.shape[0])
+    intermediate = int(w2.shape[1]) * 16
+    hidden = int(w2.shape[2]) * 16
+    hadamard = _hadamard_128(device)
+    gate_weights = torch.stack(
+        [_reconstruct_native(w13[0, expert]) for expert in range(experts)]
+    ).to(device)
+    up_weights = torch.stack(
+        [_reconstruct_native(w13[1, expert]) for expert in range(experts)]
+    ).to(device)
+    down_weights = torch.stack(
+        [_reconstruct_native(w2[expert]) for expert in range(experts)]
+    ).to(device)
+    gate_svh = intermediate_rotations[:, :intermediate]
+    up_svh = intermediate_rotations[:, intermediate : 2 * intermediate]
+    down_suh = intermediate_rotations[:, 2 * intermediate :]
+    output = torch.zeros((int(x.shape[0]), hidden), dtype=torch.float32, device=device)
+    for token in range(int(x.shape[0])):
+        for slot in range(int(local_ids.shape[1])):
+            expert = int(local_ids[token, slot])
+            source = x[token : token + 1].to(torch.float16)
+            gate_a = _had128(
+                source,
+                hadamard,
+                suh=gate_suh[expert],
+                store_fp16=True,
+            )
+            up_a = _had128(
+                source,
+                hadamard,
+                suh=up_suh[expert],
+                store_fp16=True,
+            )
+            gate = (gate_a.float() @ gate_weights[expert].float()).to(torch.float16)
+            up = (up_a.float() @ up_weights[expert].float()).to(torch.float16)
+            gate = _had128(
+                gate,
+                hadamard,
+                svh=gate_svh[expert],
+                store_fp16=True,
+            )
+            up = _had128(
+                up,
+                hadamard,
+                svh=up_svh[expert],
+                store_fp16=True,
+            )
+            activated = (gate.float() * torch.sigmoid(gate.float()) * up.float()).to(
+                torch.float16
+            )
+            down_a = _had128(
+                activated,
+                hadamard,
+                suh=down_suh[expert],
+                store_fp16=True,
+            )
+            down = (down_a.float() @ down_weights[expert].float()).to(torch.float16)
+            route = _had128(
+                down,
+                hadamard,
+                svh=down_svh[expert],
+                store_fp16=False,
+            )
+            output[token] += router_weights[token, slot] * route[0]
+    return output
+
+
+def _sm12x_available() -> bool:
+    if not torch.cuda.is_available():
+        return False
+    major, minor = torch.cuda.get_device_capability(torch.cuda.current_device())
+    return major == 12 and minor in (0, 1)
+
+
+@pytest.mark.skipif(not _sm12x_available(), reason="requires an SM120/SM121 GPU")
+@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16])
+def test_planned_full_rotation_matches_reference_and_captures(
+    input_dtype: torch.dtype,
+) -> None:
+    torch.manual_seed(20260721)
+    device = torch.device("cuda", torch.cuda.current_device())
+    experts, hidden, intermediate = 2, 128, 128
+    bits = 3
+    tile_config = (64, 128, 64, 128)
+    w13 = torch.randint(
+        -32768,
+        32767,
+        (2, experts, hidden // 16, intermediate // 16, 16 * bits),
+        dtype=torch.int16,
+        device=device,
+    )
+    w2 = torch.randint(
+        -32768,
+        32767,
+        (experts, intermediate // 16, hidden // 16, 16 * bits),
+        dtype=torch.int16,
+        device=device,
+    )
+
+    def scales(shape: tuple[int, ...]) -> torch.Tensor:
+        return (0.875 + 0.25 * torch.rand(shape, device=device)).to(torch.float16)
+
+    gate_suh = scales((experts, hidden)).contiguous()
+    up_suh = scales((experts, hidden)).contiguous()
+    intermediate_rotations = scales((experts, 3 * intermediate)).contiguous()
+    down_svh = scales((experts, hidden)).contiguous()
+    weights = trellis_moe.prepare_weights(
+        w13,
+        w2,
+        gate_suh=gate_suh,
+        up_suh=up_suh,
+        intermediate_rotations=intermediate_rotations,
+        down_svh=down_svh,
+        codebook="mcg",
+        mcg=0xCBAC1FED,
+        tile_config=tile_config,
+    )
+    assert weights.w13.data_ptr() == w13.data_ptr()
+    assert weights.w2.data_ptr() == w2.data_ptr()
+
+    plan = trellis_moe.plan(
+        trellis_moe.Caps(
+            max_tokens=2,
+            num_topk=2,
+            num_experts=experts,
+            hidden_size=hidden,
+            intermediate_size=intermediate,
+            route_num_experts=4,
+            block_size_m=8,
+            trellis_bits=bits,
+            tile_config=tile_config,
+            input_dtype=input_dtype,
+            device=device,
+        )
+    )
+    assert plan.fused_launch.moe_block_size == 8
+    assert plan.fused_launch.full_rotation
+    assert {launch.route_ids_dtype for launch in plan.identity_sums} == {
+        torch.int32,
+        torch.int64,
+    }
+    assert all(launch.full_rotation for launch in plan.mapped_sums)
+
+    spec = plan.scratch_specs()[0]
+    scratch = torch.empty(spec.shape, dtype=spec.dtype, device=spec.device)
+    x = (torch.randn((2, hidden), device=device) * 1.0e-3).to(input_dtype)
+    local_ids = torch.tensor([[0, 1], [1, 0]], dtype=torch.int32, device=device)
+    global_ids = torch.tensor([[1, 3], [3, 1]], dtype=torch.int64, device=device)
+    router_weights = torch.tensor(
+        [[0.65, 0.35], [0.2, 0.8]], dtype=torch.float32, device=device
+    )
+    route_map = torch.tensor([0, 0, 0, 1], dtype=torch.int32, device=device)
+    output_map = torch.tensor([-1, 0, -1, 1], dtype=torch.int32, device=device)
+    external_output = torch.empty((2, hidden), dtype=torch.float32, device=device)
+
+    mapped = trellis_moe.bind(
+        plan,
+        scratch=scratch,
+        a=x,
+        weights=weights,
+        topk_weights=router_weights,
+        topk_ids=global_ids,
+        route_expert_map=route_map,
+        output_expert_map=output_map,
+        output=external_output,
+    )
+    mapped_output = trellis_moe.run(binding=mapped)
+    torch.cuda.synchronize(device)
+    assert mapped_output.data_ptr() == external_output.data_ptr()
+    mapped_eager = mapped_output.clone()
+
+    identity = trellis_moe.bind(
+        plan,
+        scratch=scratch,
+        a=x,
+        weights=weights,
+        topk_weights=router_weights,
+        topk_ids=local_ids,
+    )
+    identity_output = identity.run()
+    torch.cuda.synchronize(device)
+    assert identity_output.dtype == torch.float32
+    assert torch.allclose(identity_output, mapped_eager, rtol=2.0e-3, atol=2.0e-3)
+
+    reference = _reference_full_rotation(
+        x,
+        local_ids,
+        router_weights,
+        w13,
+        w2,
+        gate_suh,
+        up_suh,
+        intermediate_rotations,
+        down_svh,
+    )
+    relative_error = (mapped_eager - reference).norm() / reference.norm().clamp_min(
+        1.0e-9
+    )
+    cosine = torch.nn.functional.cosine_similarity(
+        mapped_eager.flatten(), reference.flatten(), dim=0
+    )
+    assert float(relative_error) <= 2.0e-2
+    assert float(cosine) >= 0.999
+
+    mapped = trellis_moe.bind(
+        plan,
+        scratch=scratch,
+        a=x,
+        weights=weights,
+        topk_weights=router_weights,
+        topk_ids=global_ids,
+        route_expert_map=route_map,
+        output_expert_map=output_map,
+        output=external_output,
+    )
+    mapped.run()
+    torch.cuda.synchronize(device)
+    allocated_before = torch.cuda.memory_allocated(device)
+    mapped.run()
+    torch.cuda.synchronize(device)
+    assert torch.cuda.memory_allocated(device) == allocated_before
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured_output = mapped.run()
+    graph.replay()
+    torch.cuda.synchronize(device)
+    assert captured_output.data_ptr() == external_output.data_ptr()
+    assert torch.allclose(captured_output, mapped_eager, rtol=2.0e-3, atol=2.0e-3)

From 0367cfa7bbf8bbaa4d0e7872e59d544e732d31aa Mon Sep 17 00:00:00 2001
From: Brandon Music <brandon.m.music@gmail.com>
Date: Tue, 21 Jul 2026 21:19:47 -0400
Subject: [PATCH 04/11] moe(trellis): address CodeRabbit review on PR #49

- w4a16 kernel: compute the trellis256 B-source offset in Int64 in both
  the proj and non-proj branches. The expert-stride term overflows Int32
  for large-but-valid MoE shapes (e.g. E=256, H=8192, bits=6 reaches
  ~3.21e9 > 2^31-1), which would fetch wrong or out-of-bounds B weights.
  Matches the sibling modelopt-native staging path that already builds
  this offset with Int64(...). Value-identical for in-range offsets, so
  existing shapes/tests are unaffected.
- pyproject: declare the `triton` dependency. trellis_moe META.requires
  it and is_supported() gates on it, so a clean install lacking triton
  would silently skip the op.
- trellis_moe.api: sort __all__ (RUF022).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 pyproject.toml                                |  1 +
 .../moe/_shared/kernels/w4a16/kernel.py       | 28 +++++++++----------
 sparkinfer/moe/trellis_moe/api.py             |  8 +++---
 3 files changed, 19 insertions(+), 18 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 449b9c23..7dc9005d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -19,6 +19,7 @@ dependencies = [
   "rich>=13",
   "transformers>=4.51",
   "safetensors>=0.6",
+  "triton",
 ]
 
 [project.optional-dependencies]
diff --git a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
index 3be092c3..90359371 100644
--- a/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
+++ b/sparkinfer/moe/_shared/kernels/w4a16/kernel.py
@@ -3507,25 +3507,25 @@ def _stage_k_tile_async(
                         (self.size_k // 16) * t256_half_n16 * t256_tile_u32
                     )
                     t256_plane_u32 = self.num_experts * t256_proj_expert_u32
-                    b_src_i32 = (
-                        t256_proj * Int32(t256_plane_u32)
-                        + expert_idx * Int32(t256_proj_expert_u32)
-                        + (tile_idx * Int32(self.cta_k_blocks) + t256_kt)
-                        * Int32(t256_half_n16 * t256_tile_u32)
-                        + t256_local_n16 * Int32(t256_tile_u32)
-                        + t256_in_kt * Int32(4)
+                    b_src_i64 = (
+                        Int64(t256_proj) * Int64(t256_plane_u32)
+                        + Int64(expert_idx) * Int64(t256_proj_expert_u32)
+                        + (Int64(tile_idx) * Int64(self.cta_k_blocks) + Int64(t256_kt))
+                        * Int64(t256_half_n16 * t256_tile_u32)
+                        + Int64(t256_local_n16) * Int64(t256_tile_u32)
+                        + Int64(t256_in_kt) * Int64(4)
                     )
                 else:
-                    b_src_i32 = (
-                        Int32(t256_expert_u32) * expert_idx
-                        + (tile_idx * Int32(self.cta_k_blocks) + t256_kt)
-                        * Int32(t256_n16 * t256_tile_u32)
-                        + output_n_tile * Int32(self.cta_n_blocks * t256_tile_u32)
-                        + t256_in_kt * Int32(4)
+                    b_src_i64 = (
+                        Int64(t256_expert_u32) * Int64(expert_idx)
+                        + (Int64(tile_idx) * Int64(self.cta_k_blocks) + Int64(t256_kt))
+                        * Int64(t256_n16 * t256_tile_u32)
+                        + Int64(output_n_tile) * Int64(self.cta_n_blocks * t256_tile_u32)
+                        + Int64(t256_in_kt) * Int64(4)
                     )
                 cp_async4_shared_global_pred(
                     b_dst,
-                    get_ptr_as_int64(b_i32_flat, b_src_i32),
+                    get_ptr_as_int64(b_i32_flat, b_src_i64),
                     (t256_chunk < Int32(t256_total_chunks)).to(Int32),
                 )
 
diff --git a/sparkinfer/moe/trellis_moe/api.py b/sparkinfer/moe/trellis_moe/api.py
index f066cb50..a7a5dcd8 100644
--- a/sparkinfer/moe/trellis_moe/api.py
+++ b/sparkinfer/moe/trellis_moe/api.py
@@ -39,14 +39,14 @@ def is_supported(device=None) -> bool:
 
 
 __all__ = [
+    "Binding",
     "Caps",
     "Plan",
     "Weights",
-    "Binding",
+    "bind",
+    "clear_caches",
+    "is_supported",
     "plan",
     "prepare_weights",
-    "bind",
     "run",
-    "is_supported",
-    "clear_caches",
 ]

From 57dd6915d88b8c5d5219c48cbed99d96f06ed805 Mon Sep 17 00:00:00 2001
From: Brandon Music <brandonmusic@pop-os.tail8674da.ts.net>
Date: Wed, 22 Jul 2026 15:54:57 -0400
Subject: [PATCH 05/11] test(indexer): packed-page direct-K high-page-id
 i64-offset coverage

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 tests/attention/test_fused_indexer.py | 71 +++++++++++++++++++++++++++
 1 file changed, 71 insertions(+)

diff --git a/tests/attention/test_fused_indexer.py b/tests/attention/test_fused_indexer.py
index 5923193c..536ce58c 100644
--- a/tests/attention/test_fused_indexer.py
+++ b/tests/attention/test_fused_indexer.py
@@ -532,3 +532,74 @@ def test_fused_indexer_mla_matches_reference(rows):
         logit = (torch.relu(torch.einsum("hd,td->ht", qf[r], kf[a:b])) * weights[r].unsqueeze(1)).sum(0) * k_scales[a:b]
         gset = set((torch.topk(logit, topk).indices + a).tolist())  # absolute index
         assert set(idx[r].tolist()) == gset
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required for fused indexer")
+def test_fused_indexer_paged_direct_k_high_page_id_i64_offsets():
+    """Direct-K addressing stays exact after the packed-page Int32 boundary."""
+    device = torch.device("cuda")
+    record_bytes = 1_077_120
+    pid_lo, pages_used = 2000, 64
+    pool_pages = pid_lo + pages_used + 4
+    need = pool_pages * record_bytes + (1 << 29)
+    free, _ = torch.cuda.mem_get_info(device)
+    if free < need:
+        pytest.skip(f"needs ~{need / 2**30:.1f} GiB free device memory")
+
+    rows, heads, topk = 4, 64, 512
+    seqlen = pages_used * _PS
+    g = torch.Generator(device="cpu").manual_seed(7)
+    pool = torch.zeros(pool_pages * record_bytes, dtype=torch.uint8, device=device)
+    k_quant = pool.as_strided((pool_pages, _PS, 128), (record_bytes, 128, 1))
+    k_scales = pool.view(torch.float32).as_strided(
+        (pool_pages, _PS), (record_bytes // 4, 1), storage_offset=2048
+    )
+    k_fp8 = (torch.randn((pages_used, _PS, 128), generator=g) / 3).to(
+        torch.float8_e4m3fn
+    )
+    k_quant[pid_lo : pid_lo + pages_used] = k_fp8.view(torch.uint8).to(device)
+    k_scales[pid_lo : pid_lo + pages_used] = (
+        torch.rand((pages_used, _PS), generator=g) + 0.1
+    ).to(device)
+
+    q_fp8 = (torch.randn((rows, heads, 128), generator=g) / 3).to(
+        torch.float8_e4m3fn
+    ).to(device)
+    weights = torch.randn((rows, heads), generator=g, dtype=torch.float32).to(device)
+    page_table = (
+        torch.arange(pid_lo, pid_lo + pages_used, dtype=torch.int32, device=device)
+        .repeat(rows, 1)
+        .contiguous()
+    )
+    seqlens = torch.full((rows,), seqlen, dtype=torch.int32, device=device)
+
+    idx, val = run_fused_paged_indexer(
+        q_bytes=q_fp8.view(torch.uint8),
+        weights=weights,
+        k_quant_bytes=k_quant,
+        k_scales=k_scales,
+        real_page_table=page_table,
+        seqlens=seqlens,
+        num_heads=heads,
+        topk=topk,
+        ctas_per_group=8,
+    )
+    torch.cuda.synchronize(device)
+
+    gold_vals, gold_idx_sets = _golden_topk(
+        q_fp8,
+        weights,
+        k_quant.view(torch.float8_e4m3fn)[pid_lo : pid_lo + pages_used],
+        k_scales[pid_lo : pid_lo + pages_used],
+        page_table - pid_lo,
+        seqlens,
+        topk,
+    )
+    assert torch.allclose(
+        torch.sort(val, dim=1, descending=True).values,
+        gold_vals,
+        atol=1e-2,
+        rtol=0,
+    )
+    for row in range(rows):
+        assert set(idx[row].tolist()) == gold_idx_sets[row]

From e27ae3b34815843a3388468961be451ced9335e0 Mon Sep 17 00:00:00 2001
From: Brandon Music <brandonmusic@pop-os.tail8674da.ts.net>
Date: Wed, 22 Jul 2026 18:14:11 -0400
Subject: [PATCH 06/11] fix(cache): include device_arch in semantic manifest
 identity

device_arch (cache_payload index 4) was read into the disk-cache key but
dropped from _semantic_compile_manifest_payload, so two payloads differing
only by GPU architecture hashed to the same semantic identity. Add it to the
common semantic dict (index 4 in both v6_explicit_spec and v3 formats).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
---
 sparkinfer/_lib/compiler.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/sparkinfer/_lib/compiler.py b/sparkinfer/_lib/compiler.py
index 7fa2a45a..ee62f78c 100644
--- a/sparkinfer/_lib/compiler.py
+++ b/sparkinfer/_lib/compiler.py
@@ -2043,6 +2043,11 @@ def _semantic_compile_manifest_payload(
     semantic: dict[str, Any] = {
         "cache_format": cache_format,
         "target": _semantic_target_key(cache_payload[1]),
+        # device_arch is cache_payload index 4 in both the v6_explicit_spec and
+        # v3 formats. Include it in the semantic identity so two payloads that
+        # differ only by GPU architecture hash to distinct semantic keys and the
+        # device-aware cache format cannot be aliased across architectures.
+        "device_arch": _manifest_json_value(cache_payload[4]),
     }
     if cache_format == "sparkinfer_cute_compile_cache_v6_explicit_spec":
         semantic["compile_spec_hash"] = cache_payload[5]

From da2784b0d28a683fd1122e208d6b5d35bd8952bb Mon Sep 17 00:00:00 2001
From: "Brandon M. Music" <brandon.m.music@gmail.com>
Date: Sat, 25 Jul 2026 18:09:17 -0400
Subject: [PATCH 07/11] test(indexer): assert value-to-index pairing in the
 direct-K top-k test

The direct-K high-page-id test compared the sorted top-k values and the
returned index set independently, so a regression that emitted the correct
values against the wrong indices would still pass.

_golden_topk grows an opt-in return_scores flag (existing call sites keep the
two-tuple contract) and the test now gathers the golden score for each index
the kernel actually returned and compares it position by position. Gathering
by the returned index instead of asserting an index order keeps the check
robust to ties at the top-k boundary.
---
 tests/attention/test_fused_indexer.py | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/tests/attention/test_fused_indexer.py b/tests/attention/test_fused_indexer.py
index 536ce58c..97bad97d 100644
--- a/tests/attention/test_fused_indexer.py
+++ b/tests/attention/test_fused_indexer.py
@@ -236,10 +236,12 @@ def _build_case(rows, heads, seqlen, topk, *, seed, device):
     return q_fp8, weights, k_fp8, k_scales, page_table, seqlens
 
 
-def _golden_topk(q_fp8, weights, k_fp8, k_scales, page_table, seqlens, topk):
+def _golden_topk(
+    q_fp8, weights, k_fp8, k_scales, page_table, seqlens, topk, return_scores=False
+):
     rows, seqlen = int(q_fp8.shape[0]), int(seqlens[0])
     qf, kf = q_fp8.float(), k_fp8.float()
-    vals, idxs = [], []
+    vals, idxs, scores = [], [], []
     for r in range(rows):
         pages = page_table[r].long()
         kr = kf[pages].reshape(-1, 128)[:seqlen]
@@ -248,6 +250,9 @@ def _golden_topk(q_fp8, weights, k_fp8, k_scales, page_table, seqlens, topk):
         tk = torch.topk(logit, topk, largest=True, sorted=True)
         vals.append(tk.values)
         idxs.append(set(tk.indices.tolist()))
+        scores.append(logit)
+    if return_scores:
+        return torch.stack(vals), idxs, torch.stack(scores)
     return torch.stack(vals), idxs
 
 
@@ -586,7 +591,7 @@ def test_fused_indexer_paged_direct_k_high_page_id_i64_offsets():
     )
     torch.cuda.synchronize(device)
 
-    gold_vals, gold_idx_sets = _golden_topk(
+    gold_vals, gold_idx_sets, gold_scores = _golden_topk(
         q_fp8,
         weights,
         k_quant.view(torch.float8_e4m3fn)[pid_lo : pid_lo + pages_used],
@@ -594,6 +599,7 @@ def test_fused_indexer_paged_direct_k_high_page_id_i64_offsets():
         page_table - pid_lo,
         seqlens,
         topk,
+        return_scores=True,
     )
     assert torch.allclose(
         torch.sort(val, dim=1, descending=True).values,
@@ -603,3 +609,11 @@ def test_fused_indexer_paged_direct_k_high_page_id_i64_offsets():
     )
     for row in range(rows):
         assert set(idx[row].tolist()) == gold_idx_sets[row]
+        # Value-to-index pairing. Comparing the sorted value list and the index set
+        # separately would still pass if correct values were emitted against the
+        # wrong indices, so gather the golden score for each index the kernel
+        # actually returned and compare position by position. Gathering by the
+        # returned index (rather than asserting an index order) keeps this robust
+        # to ties at the top-k boundary.
+        paired = gold_scores[row].index_select(0, idx[row].long())
+        assert torch.allclose(val[row].float(), paired.float(), atol=1e-2, rtol=0)

From b916dcc175644f7cc626986a243e5a7c81cd9ad0 Mon Sep 17 00:00:00 2001
From: "Brandon M. Music" <brandon.m.music@gmail.com>
Date: Sat, 25 Jul 2026 19:07:54 -0400
Subject: [PATCH 08/11] address CodeRabbit nitpicks: device-key precision and
 an executable test invariant

compiler: read the arch cache key from the explicitly selected device.
torch.cuda.get_device_capability()/get_device_name() with no argument read the
implicitly current device and the result is memoized process-wide, so a lookup
made before a worker selects its GPU pins the key to the wrong card. Observable
on mixed rigs: Max-Q and non-Max-Q RTX PRO 6000 boards share compute capability
12.0 but report different device names.

prepare: normalise an indexless resolved_device to a concrete ordinal in the
synthesize path, then validate dummy_scale with the same strict comparison used
for the suh/svh tables. Previously the two disagreed -- suh/svh used strict
equality (so cuda:0 != cuda rejected a correctly placed tensor) while
dummy_scale used an index-optional comparison that accepted a tensor from any
ordinal. Normalising once makes both strict and correct.

tests: assert the direct-K case actually crosses the Int32 page-offset boundary.
The i64 path it exists to cover is implied only by the constants, so retuning
record_bytes or pid_lo downward would leave a passing test that no longer
exercises the boundary.
---
 sparkinfer/_lib/compiler.py                   | 10 +++++++--
 .../moe/_shared/kernels/w4a16/prepare.py      | 21 ++++++++++++-------
 tests/attention/test_fused_indexer.py         |  8 +++++++
 3 files changed, 29 insertions(+), 10 deletions(-)

diff --git a/sparkinfer/_lib/compiler.py b/sparkinfer/_lib/compiler.py
index ee62f78c..3835ddd5 100644
--- a/sparkinfer/_lib/compiler.py
+++ b/sparkinfer/_lib/compiler.py
@@ -1467,8 +1467,14 @@ def _device_arch_key() -> tuple[object, ...]:
 
         if not torch.cuda.is_available():
             return ("arch", "unavailable")
-        major, minor = torch.cuda.get_device_capability()
-        name = torch.cuda.get_device_name()
+        # Read the explicitly selected device rather than whatever is implicitly
+        # current. This key is memoized process-wide, so an implicit lookup made
+        # before the worker selects its GPU would pin the key to the wrong card.
+        # That is observable on mixed rigs: Max-Q and non-Max-Q RTX PRO 6000
+        # boards share compute capability 12.0 but report different device names.
+        index = torch.cuda.current_device()
+        major, minor = torch.cuda.get_device_capability(index)
+        name = torch.cuda.get_device_name(index)
     except Exception:
         return ("arch", "unavailable")
     _DEVICE_ARCH_KEY = ("arch", int(major), int(minor), str(name))
diff --git a/sparkinfer/moe/_shared/kernels/w4a16/prepare.py b/sparkinfer/moe/_shared/kernels/w4a16/prepare.py
index a650320b..f0bfacf6 100644
--- a/sparkinfer/moe/_shared/kernels/w4a16/prepare.py
+++ b/sparkinfer/moe/_shared/kernels/w4a16/prepare.py
@@ -1719,6 +1719,14 @@ def prepare_trellis256_moe_weights(
                 "trellis3_t256 W4A16 weights require a CUDA device, got "
                 f"{resolved_device}"
             )
+        if resolved_device.index is None:
+            # Pin an indexless "cuda" to a concrete ordinal so every downstream
+            # device check compares like with like. Without this the suh/svh
+            # tables are validated with strict equality (cuda:0 != cuda, so a
+            # correctly placed tensor is rejected) while dummy_scale is
+            # validated with an index-optional comparison that accepts a tensor
+            # from any ordinal. Normalising once keeps both strict and correct.
+            resolved_device = torch.device("cuda", torch.cuda.current_device())
         generator = torch.Generator(device=resolved_device)
         generator.manual_seed(int(seed))
         if w13_layout == "trellis3_t256_proj":
@@ -1855,14 +1863,11 @@ def prepare_trellis256_moe_weights(
     if dummy_scale is None:
         dummy_scale = torch.zeros(4, dtype=torch.uint8, device=resolved_device)
     else:
-        dummy_device_matches = (
-            dummy_scale.device.type == resolved_device.type
-            and (
-                resolved_device.index is None
-                or dummy_scale.device.index == resolved_device.index
-            )
-        )
-        if not dummy_device_matches:
+        # resolved_device always carries a concrete ordinal by this point (the
+        # synthesize path pins an indexless "cuda" above, the native path takes it
+        # from w13), so use the same strict comparison as the suh/svh tables
+        # instead of accepting a tensor from an arbitrary ordinal.
+        if dummy_scale.device != resolved_device:
             raise ValueError(
                 "trellis3_t256 dummy_scale must share the weight device, got "
                 f"{dummy_scale.device} and {resolved_device}"
diff --git a/tests/attention/test_fused_indexer.py b/tests/attention/test_fused_indexer.py
index 97bad97d..1f03677f 100644
--- a/tests/attention/test_fused_indexer.py
+++ b/tests/attention/test_fused_indexer.py
@@ -545,6 +545,14 @@ def test_fused_indexer_paged_direct_k_high_page_id_i64_offsets():
     device = torch.device("cuda")
     record_bytes = 1_077_120
     pid_lo, pages_used = 2000, 64
+    # The whole point of this case is that the base page offset overflows Int32,
+    # which forces the i64 addressing path. Assert it instead of leaving it
+    # implicit in the constants: retuning record_bytes or pid_lo downward would
+    # otherwise leave a green test that no longer exercises the boundary at all.
+    assert pid_lo * record_bytes >= (1 << 31), (
+        f"page base offset {pid_lo * record_bytes} no longer crosses the Int32 "
+        "boundary; this test would silently stop covering i64 offsets"
+    )
     pool_pages = pid_lo + pages_used + 4
     need = pool_pages * record_bytes + (1 << 29)
     free, _ = torch.cuda.mem_get_info(device)

From 87320905f220c7cd0ea7e8b8934aaa9d08fc7a89 Mon Sep 17 00:00:00 2001
From: "Brandon M. Music" <brandon.m.music@gmail.com>
Date: Sat, 25 Jul 2026 19:23:41 -0400
Subject: [PATCH 09/11] compiler: memoize the arch cache key per device ordinal

The previous change was not a fix. torch resolves get_device_capability() and
get_device_name() against torch.cuda.current_device() when called without an
argument, so passing that same ordinal explicitly is equivalent and left the
process-wide _DEVICE_ARCH_KEY free to freeze whichever GPU happened to be
current on the first call. It also broke
test_device_arch_key_retries_after_unavailable, whose torch stubs take no
device argument.

Replace the single global with a per-ordinal dict so each entry is bound to the
device it was measured on, and thread the ordinal through
_static_compile_cache_context, which is lru_cached on compile_callable and
would otherwise re-freeze the identity at that layer regardless.

The returned key still omits the ordinal: two GPUs with the same capability and
name produce the same key and keep sharing compiled artifacts, so a homogeneous
rig does not fragment its JIT cache. Only the lookup is per ordinal.

Adds test_device_arch_key_is_memoized_per_ordinal, which models a mixed rig
(two boards sharing compute capability 12.0 but reporting different names) and
asserts the second board does not inherit the first board's identity. Restores
the retry test against the per-ordinal cache.
---
 sparkinfer/_lib/compiler.py      | 68 ++++++++++++++++++++++++--------
 tests/_lib/test_compile_cache.py | 53 ++++++++++++++++++++++---
 2 files changed, 99 insertions(+), 22 deletions(-)

diff --git a/sparkinfer/_lib/compiler.py b/sparkinfer/_lib/compiler.py
index 3835ddd5..56f596c1 100644
--- a/sparkinfer/_lib/compiler.py
+++ b/sparkinfer/_lib/compiler.py
@@ -1448,37 +1448,65 @@ def _distribution_version(name: str) -> str:
         return ""
 
 
-_DEVICE_ARCH_KEY: tuple[object, ...] | None = None
+_DEVICE_ARCH_KEYS: dict[int, tuple[object, ...]] = {}
 
 
-def _device_arch_key() -> tuple[object, ...]:
-    """Compute capability of the device this process will compile for.
+def _current_device_ordinal() -> int | None:
+    """Ordinal of the device this process is currently compiling for."""
+    try:
+        import torch
+
+        if not torch.cuda.is_available():
+            return None
+        return int(torch.cuda.current_device())
+    except Exception:
+        return None
+
+
+def _device_arch_key(device_ordinal: int | None = None) -> tuple[object, ...]:
+    """Architecture identity of the device this process will compile for.
+
+    Memoized *per device ordinal* rather than once per process. Torch resolves
+    ``get_device_capability()`` / ``get_device_name()`` against the current
+    device when called without an argument, so a single process-wide cache
+    freezes whichever GPU happened to be current on the first call -- passing
+    ``current_device()`` explicitly is equivalent and does not change that. A
+    process that selects a different GPU afterwards would keep reusing the stale
+    identity. Keying the memo by ordinal binds every entry to the device it was
+    actually measured on, which is what makes the guarantee below hold on a
+    mixed-GPU process (e.g. Max-Q and non-Max-Q RTX PRO 6000 boards share
+    compute capability 12.0 but report different device names).
+
+    The returned key deliberately omits the ordinal: two GPUs with the same
+    capability and name yield the same key and therefore share compiled
+    artifacts, which is the desired behaviour on a homogeneous rig. Only the
+    lookup is per ordinal.
 
     Deliberately NOT lru_cached on failure: if CUDA is not initialized yet the
     query is retried on the next call, so an early call can never memoize a
     bogus "unknown" and let a cubin built for one architecture be served to
     another from the shared on-disk cache.
     """
-    global _DEVICE_ARCH_KEY
-    if _DEVICE_ARCH_KEY is not None:
-        return _DEVICE_ARCH_KEY
     try:
         import torch
 
         if not torch.cuda.is_available():
             return ("arch", "unavailable")
-        # Read the explicitly selected device rather than whatever is implicitly
-        # current. This key is memoized process-wide, so an implicit lookup made
-        # before the worker selects its GPU would pin the key to the wrong card.
-        # That is observable on mixed rigs: Max-Q and non-Max-Q RTX PRO 6000
-        # boards share compute capability 12.0 but report different device names.
-        index = torch.cuda.current_device()
+        index = (
+            int(torch.cuda.current_device())
+            if device_ordinal is None
+            else int(device_ordinal)
+        )
+        cached = _DEVICE_ARCH_KEYS.get(index)
+        if cached is not None:
+            return cached
         major, minor = torch.cuda.get_device_capability(index)
         name = torch.cuda.get_device_name(index)
     except Exception:
         return ("arch", "unavailable")
-    _DEVICE_ARCH_KEY = ("arch", int(major), int(minor), str(name))
-    return _DEVICE_ARCH_KEY
+    key = ("arch", int(major), int(minor), str(name))
+    _DEVICE_ARCH_KEYS[index] = key
+    return key
 
 
 @lru_cache(maxsize=1)
@@ -1572,11 +1600,17 @@ def _compile_environment_key() -> tuple[tuple[str, str], ...]:
 
 
 @lru_cache(maxsize=16)
-def _static_compile_cache_context(compile_callable: Any) -> tuple[object, ...]:
+def _static_compile_cache_context(
+    compile_callable: Any, device_ordinal: int | None = None
+) -> tuple[object, ...]:
+    # device_ordinal participates in the lru_cache key on purpose: without it
+    # this cache would pin the architecture identity to whichever GPU was
+    # current the first time a given compile_callable was seen, re-introducing
+    # the staleness that keying _device_arch_key per ordinal removes.
     return (
         _sparkinfer_package_fingerprint(),
         _runtime_toolchain_key(),
-        _device_arch_key(),
+        _device_arch_key(device_ordinal),
         _compile_options_cache_key(compile_callable),
         _compile_environment_key(),
     )
@@ -1891,7 +1925,7 @@ def _compile_disk_cache_payload(
         device_arch,
         compile_options,
         compile_environment,
-    ) = _static_compile_cache_context(compile_callable)
+    ) = _static_compile_cache_context(compile_callable, _current_device_ordinal())
     if compile_spec is not None:
         kwargs_json_key, kwargs_hash_key = _compile_kwargs_json_key(kwargs)
         return (
diff --git a/tests/_lib/test_compile_cache.py b/tests/_lib/test_compile_cache.py
index bb37c016..cc5fc8c7 100644
--- a/tests/_lib/test_compile_cache.py
+++ b/tests/_lib/test_compile_cache.py
@@ -58,7 +58,7 @@ def test_disk_cache_key_includes_device_arch(monkeypatch):
     monkeypatch.setattr(
         compiler,
         "_static_compile_cache_context",
-        lambda _callable: (
+        lambda _callable, _ordinal=None: (
             "package",
             "toolchain",
             ("arch", 12, 0, "gpu"),
@@ -80,12 +80,55 @@ def test_disk_cache_key_includes_device_arch(monkeypatch):
 
 def test_device_arch_key_retries_after_unavailable(monkeypatch):
     torch = pytest.importorskip("torch")
-    monkeypatch.setattr(compiler, "_DEVICE_ARCH_KEY", None)
+    monkeypatch.setattr(compiler, "_DEVICE_ARCH_KEYS", {})
     monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
     assert compiler._device_arch_key() == ("arch", "unavailable")
-    assert compiler._DEVICE_ARCH_KEY is None
+    # An unavailable probe must not be memoized, or an early call would pin a
+    # bogus identity for the rest of the process.
+    assert compiler._DEVICE_ARCH_KEYS == {}
 
     monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
-    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda: (12, 1))
-    monkeypatch.setattr(torch.cuda, "get_device_name", lambda: "SM121")
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (12, 1))
+    monkeypatch.setattr(torch.cuda, "get_device_name", lambda device=None: "SM121")
     assert compiler._device_arch_key() == ("arch", 12, 1, "SM121")
+    assert compiler._DEVICE_ARCH_KEYS == {0: ("arch", 12, 1, "SM121")}
+
+
+def test_device_arch_key_is_memoized_per_ordinal(monkeypatch):
+    """A second GPU must not inherit the first GPU's architecture identity.
+
+    torch resolves get_device_capability()/get_device_name() against the current
+    device when no argument is given, so a process-wide cache would freeze the
+    first device probed. Mixed rigs make this observable: two boards can share a
+    compute capability but report different device names.
+    """
+    torch = pytest.importorskip("torch")
+    monkeypatch.setattr(compiler, "_DEVICE_ARCH_KEYS", {})
+    monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
+
+    boards = {0: (12, 0, "MAX-Q"), 1: (12, 0, "FULL")}
+    current = {"index": 0}
+    monkeypatch.setattr(torch.cuda, "current_device", lambda: current["index"])
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_capability",
+        lambda device=None: boards[current["index"] if device is None else device][:2],
+    )
+    monkeypatch.setattr(
+        torch.cuda,
+        "get_device_name",
+        lambda device=None: boards[current["index"] if device is None else device][2],
+    )
+
+    assert compiler._device_arch_key() == ("arch", 12, 0, "MAX-Q")
+    # Selecting a different board must re-probe rather than return the memoized
+    # identity of board 0.
+    current["index"] = 1
+    assert compiler._device_arch_key() == ("arch", 12, 0, "FULL")
+    # Explicit ordinals resolve independently of whichever device is current.
+    assert compiler._device_arch_key(0) == ("arch", 12, 0, "MAX-Q")
+    assert compiler._DEVICE_ARCH_KEYS == {
+        0: ("arch", 12, 0, "MAX-Q"),
+        1: ("arch", 12, 0, "FULL"),
+    }

From 2343a317b96a9df518fd3b12ec483cb544e4df97 Mon Sep 17 00:00:00 2001
From: Brandon Music <brandon.m.music@gmail.com>
Date: Sun, 26 Jul 2026 19:12:48 -0400
Subject: [PATCH 10/11] tests: cover small-m partial blocks on the planned
 full-rotation Trellis path

vLLM issue #183 reported that widening the Trellis window
(VLLM_EXL3_TRELLIS_MIN_M=1), so an MTP-N draft's m=1..N GEMMs stay on the fused
path instead of falling to the eager parity path, silently corrupts output at
long context while short prompts pass.

The claim was not reproducible, but it exposed a real coverage hole: the entire
trellis3_t256 + full_rotation production path had exactly one numerical test,
and it ran m == max_tokens == 2 with topk=2 on a non-production tile. Nothing
covered m < plan capacity -- which is precisely the window-widening regime --
and at m < moe_block_size the route packer emits mostly-padding blocks (at m=3,
topk=8, roughly 1-3 real rows out of 8), so padding masking is load-bearing.

Adds a parametrized test over m in {1,2,3,4,8} at production geometry
(tile 64x256x64x256, block_size_m=8, capacity 32, topk=8, 32 experts) with
three oracles:

  * per-token bitwise equality against an m=1 run of the same token, since a
    token's output depends only on its own routes and must not change with
    batch composition;
  * the scratch arena pre-filled with 0xFF (NaN in fp32/fp16) before every
    bind, so any read of a padding or uninitialised slot surfaces as NaN
    rather than as plausible numbers;
  * m strictly below plan capacity, the regime the issue implicates.

All five cases pass, which converts "we believe MIN_M=1 is safe" into a gate.
---
 tests/moe/test_trellis_moe.py | 107 ++++++++++++++++++++++++++++++++++
 1 file changed, 107 insertions(+)

diff --git a/tests/moe/test_trellis_moe.py b/tests/moe/test_trellis_moe.py
index 95700554..a8332cfc 100644
--- a/tests/moe/test_trellis_moe.py
+++ b/tests/moe/test_trellis_moe.py
@@ -469,3 +469,110 @@ def scales(shape: tuple[int, ...]) -> torch.Tensor:
     torch.cuda.synchronize(device)
     assert captured_output.data_ptr() == external_output.data_ptr()
     assert torch.allclose(captured_output, mapped_eager, rtol=2.0e-3, atol=2.0e-3)
+
+
+@pytest.mark.skipif(not _sm12x_available(), reason="requires an SM120/SM121 GPU")
+@pytest.mark.parametrize("m", [1, 2, 3, 4, 8])
+def test_planned_full_rotation_small_m_partial_blocks(m: int) -> None:
+    """Small-m coverage at production geometry, with m < plan capacity.
+
+    Motivated by vLLM issue #183, which reported that widening the Trellis window
+    (VLLM_EXL3_TRELLIS_MIN_M=1) so an MTP-N draft's m=1..N GEMMs stay on the fused
+    path silently corrupts output. m < moe_block_size means the route packer emits
+    mostly-padding blocks (at m=3, topk=8 roughly 1-3 real rows of 8), so padding
+    masking is load-bearing and was previously untested: the only numerical test on
+    this path ran m == max_tokens == 2 with topk=2 and a non-production tile.
+
+    Three oracles:
+      * per-token bitwise equality against an m=1 run of the same token. A token's
+        output depends only on its own routes, so batching must not change a bit.
+      * the scratch arena pre-filled with 0xFF (NaN in fp32/fp16) before every
+        bind, so any read of a padding/uninitialised slot surfaces as NaN.
+      * plan capacity 32 vs m, i.e. the window-widening regime itself.
+    """
+    torch.manual_seed(20260726)
+    device = torch.device("cuda", torch.cuda.current_device())
+    experts, hidden, intermediate = 32, 1024, 1024
+    bits, topk, capacity = 3, 8, 32
+    tile_config = (64, 256, 64, 256)
+
+    w13 = torch.randint(
+        -32768, 32767,
+        (2, experts, hidden // 16, intermediate // 16, 16 * bits),
+        dtype=torch.int16, device=device,
+    )
+    w2 = torch.randint(
+        -32768, 32767,
+        (experts, intermediate // 16, hidden // 16, 16 * bits),
+        dtype=torch.int16, device=device,
+    )
+
+    def scales(shape: tuple[int, ...]) -> torch.Tensor:
+        return (0.875 + 0.25 * torch.rand(shape, device=device)).to(torch.float16)
+
+    weights = trellis_moe.prepare_weights(
+        w13,
+        w2,
+        gate_suh=scales((experts, hidden)).contiguous(),
+        up_suh=scales((experts, hidden)).contiguous(),
+        intermediate_rotations=scales((experts, 3 * intermediate)).contiguous(),
+        down_svh=scales((experts, hidden)).contiguous(),
+        codebook="mcg",
+        mcg=0xCBAC1FED,
+        tile_config=tile_config,
+    )
+    plan = trellis_moe.plan(
+        trellis_moe.Caps(
+            max_tokens=capacity,
+            num_topk=topk,
+            num_experts=experts,
+            hidden_size=hidden,
+            intermediate_size=intermediate,
+            route_num_experts=experts,
+            block_size_m=8,
+            trellis_bits=bits,
+            tile_config=tile_config,
+            input_dtype=torch.bfloat16,
+            device=device,
+        )
+    )
+    assert plan.fused_launch.moe_block_size == 8
+    assert plan.fused_launch.full_rotation
+
+    spec = plan.scratch_specs()[0]
+    scratch = torch.empty(spec.shape, dtype=spec.dtype, device=spec.device)
+
+    x = (torch.randn((m, hidden), device=device) * 1.0e-3).to(torch.bfloat16)
+    ids = torch.stack(
+        [torch.randperm(experts, device=device)[:topk] for _ in range(m)]
+    ).to(torch.int32)
+    router_weights = torch.rand((m, topk), dtype=torch.float32, device=device)
+    router_weights /= router_weights.sum(dim=1, keepdim=True)
+
+    def _run(rows: slice) -> torch.Tensor:
+        # NaN-poison the arena so a padding-slot read cannot pass silently.
+        scratch.fill_(0xFF)
+        binding = trellis_moe.bind(
+            plan,
+            scratch=scratch,
+            a=x[rows].contiguous(),
+            weights=weights,
+            topk_weights=router_weights[rows].contiguous(),
+            topk_ids=ids[rows].contiguous(),
+        )
+        out = trellis_moe.run(binding=binding)
+        torch.cuda.synchronize(device)
+        return out.clone()
+
+    batched = _run(slice(0, m))
+    assert not torch.isnan(batched).any(), (
+        f"NaN in fused output at m={m}: a padding or uninitialised arena slot was read"
+    )
+
+    for row in range(m):
+        single = _run(slice(row, row + 1))
+        assert not torch.isnan(single).any()
+        assert torch.equal(batched[row], single[0]), (
+            f"row {row} of an m={m} batch differs from its own m=1 run; a token's "
+            "output must not depend on how many other tokens share the batch"
+        )

From aec4fc022771ca23df1ae951c5b6523d4760f782 Mon Sep 17 00:00:00 2001
From: Brandon Music <brandon.m.music@gmail.com>
Date: Sun, 26 Jul 2026 19:27:20 -0400
Subject: [PATCH 11/11] tests: cover CUDA-graph capture at m below plan
 capacity

Closes the last surface left untested by vLLM #183. The existing capture test
captures at m == max_tokens, and the eager small-m test does not cover capture:
bind() runs inside the captured region, so the route-pack grid, the
kernel_workspace zero-fill and the active-m constants are baked into the graph.
If any were captured for the wrong m, replay would return stale or mismatched
results while eager execution stayed correct -- which is precisely a "short
prompts pass, long context corrupts" signature once an MTP-N draft replays its
captured graph.

Parametrized over m in {1,2,3} with capacity 32 at production geometry. Asserts
replay matches eager bit-for-bit, no NaN after replay, and that a second replay
is stable (no dependence on scratch left over from the capture pass).

All three pass, so capture at m below capacity is not the reported defect
either.
---
 tests/moe/test_trellis_moe.py | 107 ++++++++++++++++++++++++++++++++++
 1 file changed, 107 insertions(+)

diff --git a/tests/moe/test_trellis_moe.py b/tests/moe/test_trellis_moe.py
index a8332cfc..7d871573 100644
--- a/tests/moe/test_trellis_moe.py
+++ b/tests/moe/test_trellis_moe.py
@@ -576,3 +576,110 @@ def _run(rows: slice) -> torch.Tensor:
             f"row {row} of an m={m} batch differs from its own m=1 run; a token's "
             "output must not depend on how many other tokens share the batch"
         )
+
+
+@pytest.mark.skipif(not _sm12x_available(), reason="requires an SM120/SM121 GPU")
+@pytest.mark.parametrize("m", [1, 2, 3])
+def test_planned_full_rotation_capture_below_capacity(m: int) -> None:
+    """CUDA-graph capture and replay at m strictly below plan capacity.
+
+    This is the surface vLLM #183 actually leaves untested. The existing capture
+    test captures at ``m == max_tokens``, and a small-m eager check does not cover
+    capture: ``bind`` runs inside the captured region, so the route-pack grid, the
+    ``kernel_workspace`` zero-fill and the active-m constants are all baked into
+    the graph. If any of those were captured for the wrong m, replay would produce
+    stale or mismatched results while eager execution stayed correct -- which is
+    exactly a "short prompts pass, long context corrupts" signature once an
+    MTP-N draft replays its captured graph.
+
+    Asserts replay matches the eager result for the same inputs bit-for-bit, and
+    that a second replay is stable (no dependence on scratch left over from the
+    capture pass).
+    """
+    torch.manual_seed(20260726)
+    device = torch.device("cuda", torch.cuda.current_device())
+    experts, hidden, intermediate = 32, 1024, 1024
+    bits, topk, capacity = 3, 8, 32
+    tile_config = (64, 256, 64, 256)
+
+    w13 = torch.randint(
+        -32768, 32767,
+        (2, experts, hidden // 16, intermediate // 16, 16 * bits),
+        dtype=torch.int16, device=device,
+    )
+    w2 = torch.randint(
+        -32768, 32767,
+        (experts, intermediate // 16, hidden // 16, 16 * bits),
+        dtype=torch.int16, device=device,
+    )
+
+    def scales(shape: tuple[int, ...]) -> torch.Tensor:
+        return (0.875 + 0.25 * torch.rand(shape, device=device)).to(torch.float16)
+
+    weights = trellis_moe.prepare_weights(
+        w13,
+        w2,
+        gate_suh=scales((experts, hidden)).contiguous(),
+        up_suh=scales((experts, hidden)).contiguous(),
+        intermediate_rotations=scales((experts, 3 * intermediate)).contiguous(),
+        down_svh=scales((experts, hidden)).contiguous(),
+        codebook="mcg",
+        mcg=0xCBAC1FED,
+        tile_config=tile_config,
+    )
+    plan = trellis_moe.plan(
+        trellis_moe.Caps(
+            max_tokens=capacity,
+            num_topk=topk,
+            num_experts=experts,
+            hidden_size=hidden,
+            intermediate_size=intermediate,
+            route_num_experts=experts,
+            block_size_m=8,
+            trellis_bits=bits,
+            tile_config=tile_config,
+            input_dtype=torch.bfloat16,
+            device=device,
+        )
+    )
+    spec = plan.scratch_specs()[0]
+    scratch = torch.empty(spec.shape, dtype=spec.dtype, device=spec.device)
+
+    x = (torch.randn((m, hidden), device=device) * 1.0e-3).to(torch.bfloat16)
+    ids = torch.stack(
+        [torch.randperm(experts, device=device)[:topk] for _ in range(m)]
+    ).to(torch.int32)
+    router_weights = torch.rand((m, topk), dtype=torch.float32, device=device)
+    router_weights /= router_weights.sum(dim=1, keepdim=True)
+
+    binding = trellis_moe.bind(
+        plan,
+        scratch=scratch,
+        a=x,
+        weights=weights,
+        topk_weights=router_weights,
+        topk_ids=ids,
+    )
+    eager = binding.run().clone()
+    torch.cuda.synchronize(device)
+    assert not torch.isnan(eager).any()
+
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        captured = binding.run()
+    graph.replay()
+    torch.cuda.synchronize(device)
+    assert not torch.isnan(captured).any(), (
+        f"NaN after graph replay at m={m} (capacity {capacity})"
+    )
+    assert torch.equal(captured, eager), (
+        f"graph replay at m={m} below capacity {capacity} disagrees with eager; "
+        "capture baked constants or a route-pack grid for the wrong m"
+    )
+
+    graph.replay()
+    torch.cuda.synchronize(device)
+    assert torch.equal(captured, eager), (
+        f"second replay at m={m} drifted; replay depends on scratch state left "
+        "over from the capture pass"
+    )