diff --git a/MiniMaxAI-msa-blackwell/tests/test_api_surface.py b/MiniMaxAI-msa-blackwell/tests/test_api_surface.py
index 44dbba4..d3f038e 100644
--- a/MiniMaxAI-msa-blackwell/tests/test_api_surface.py
+++ b/MiniMaxAI-msa-blackwell/tests/test_api_surface.py
@@ -1,9 +1,14 @@
 # SPDX-License-Identifier: Apache-2.0
 """API-surface checks for the MiniMax MSA Blackwell Hub package."""
+#!/usr/bin/env python3
 
 from __future__ import annotations
 
-import minimaxai_msa_blackwell as msa
+import argparse
+import importlib
+import json
+import sys
+from pathlib import Path
 
 OFFICIAL_NAMES = {
     "sparse_atten_func",
@@ -22,50 +27,72 @@
 }
 
 
-def test_v1_available_functions_are_exported() -> None:
+def _load_module(artifact):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return importlib.import_module("minimaxai_msa_blackwell")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def _check_v1_available_functions_are_exported(msa) -> None:
     available = set(msa.available_functions())
-    assert available == set(msa.V1_AVAILABLE_FUNCTIONS)
+    if available != set(msa.V1_AVAILABLE_FUNCTIONS):
+        raise AssertionError("available_functions() must equal V1_AVAILABLE_FUNCTIONS")
     for name in available:
-        assert hasattr(msa, name), f"{name} is listed as available but not exported"
+        if not hasattr(msa, name):
+            raise AssertionError(f"{name} is listed as available but not exported")
+    print("PASS v1 available functions are exported")
 
 
-def test_official_api_status_is_complete() -> None:
+def _check_official_api_status_is_complete(msa) -> None:
     tracked = set(msa.official_minimax_msa_functions())
-    assert tracked == OFFICIAL_NAMES
+    if tracked != OFFICIAL_NAMES:
+        raise AssertionError(f"tracked {tracked} != official {OFFICIAL_NAMES}")
     status = msa.official_api_status()
-    assert set(status) == OFFICIAL_NAMES
+    if set(status) != OFFICIAL_NAMES:
+        raise AssertionError("official_api_status() keys must match OFFICIAL_NAMES")
     for name, item in status.items():
-        assert item["status"] in {
-            "available",
-            "available_optional_te",
-        }
-        assert item["target"]
-        assert item["reason"]
+        if item["status"] not in {"available", "available_optional_te"}:
+            raise AssertionError(f"{name}: unexpected status {item['status']}")
+        if not item["target"]:
+            raise AssertionError(f"{name}: missing target")
+        if not item["reason"]:
+            raise AssertionError(f"{name}: missing reason")
+    print("PASS official API status is complete")
 
 
-def test_official_names_are_exported_at_root() -> None:
+def _check_official_names_are_exported_at_root(msa) -> None:
     for name in OFFICIAL_NAMES:
-        assert hasattr(msa, name), f"{name} must be exported for API compatibility"
+        if not hasattr(msa, name):
+            raise AssertionError(f"{name} must be exported for API compatibility")
+    print("PASS official names are exported at root")
 
 
-def test_pure_python_compat_helpers() -> None:
+def _check_pure_python_compat_helpers(msa) -> None:
     import torch
 
     scale = torch.arange(8, dtype=torch.float32).reshape(2, 4)
     swizzled = msa.swizzle_nvfp4_scale_to_128x4(scale, rows=2, cols=4)
-    assert swizzled.shape == (128, 4)
-    assert msa.nvfp4_global_scale_from_amax(torch.tensor([2688.0])).item() == 1.0
+    if swizzled.shape != (128, 4):
+        raise AssertionError("swizzle_nvfp4_scale_to_128x4 shape mismatch")
+    if msa.nvfp4_global_scale_from_amax(torch.tensor([2688.0])).item() != 1.0:
+        raise AssertionError("nvfp4_global_scale_from_amax mismatch")
 
     q2k = torch.tensor([[[0, 1], [1, -1]]], dtype=torch.int32)
     cu_q = torch.tensor([0, 2], dtype=torch.int32)
     cu_k = torch.tensor([0, 256], dtype=torch.int32)
     row_ptr, q_idx = msa.build_k2q_csr(q2k, cu_q, cu_k, 128, total_k=256)
-    assert row_ptr.dtype == torch.int32
-    assert q_idx.dtype == torch.int32
-    assert row_ptr.shape == (1, 3)
+    if row_ptr.dtype != torch.int32 or q_idx.dtype != torch.int32:
+        raise AssertionError("build_k2q_csr must return int32 tensors")
+    if row_ptr.shape != (1, 3):
+        raise AssertionError(f"build_k2q_csr row_ptr shape {row_ptr.shape}")
+    print("PASS pure-python compat helpers")
 
 
-def test_fp4_indexer_block_scores_is_callable() -> None:
+def _check_fp4_indexer_block_scores_is_callable(msa) -> None:
     import torch
 
     total_q, hq, hkv, pages, packed_d = 2, 4, 1, 1, 64
@@ -79,19 +106,48 @@ def test_fp4_indexer_block_scores_is_callable() -> None:
     kv_indices = torch.tensor([0], dtype=torch.int32)
 
     scores = msa.fp4_indexer_block_scores(
-        q_fp4,
-        k_fp4,
-        q_scale,
-        k_scale,
-        cu_q,
-        cu_k,
-        cu_pages,
-        max_seqlen_q=total_q,
-        max_seqlen_k=128,
-        kv_indices=kv_indices,
-        fp4_format="nvfp4",
-        causal=True,
-        scale_layout="public",
+        q_fp4, k_fp4, q_scale, k_scale, cu_q, cu_k, cu_pages,
+        max_seqlen_q=total_q, max_seqlen_k=128, kv_indices=kv_indices,
+        fp4_format="nvfp4", causal=True, scale_layout="public",
     )
-    assert scores.shape == (hq, 1, total_q)
-    assert torch.isfinite(scores).all()
+    if scores.shape != (hq, 1, total_q):
+        raise AssertionError(f"fp4_indexer_block_scores shape {scores.shape}")
+    if not torch.isfinite(scores).all():
+        raise AssertionError("fp4_indexer_block_scores must be finite")
+    print("PASS fp4_indexer_block_scores is callable")
+
+
+def run(args) -> None:
+    msa = _load_module(args.artifact)
+    _check_v1_available_functions_are_exported(msa)
+    _check_official_api_status_is_complete(msa)
+    _check_official_names_are_exported_at_root(msa)
+    _check_pure_python_compat_helpers(msa)
+    _check_fp4_indexer_block_scores_is_callable(msa)
+    print("PASS MiniMaxAI-msa-blackwell API surface: 5 checks")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source",
+                        help="the Hub package is the supported path; source falls back to sys.path")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    try:
+        run(args)
+    except Exception:
+        import traceback
+        traceback.print_exc()
+        return 1
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(
+            json.dumps({"passed": 1, "total": 1, "backend": args.backend}) + "\n"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/MiniMaxAI-msa-blackwell/tests/test_msa_blackwell.py b/MiniMaxAI-msa-blackwell/tests/test_msa_blackwell.py
index 780ecc4..cc4450b 100644
--- a/MiniMaxAI-msa-blackwell/tests/test_msa_blackwell.py
+++ b/MiniMaxAI-msa-blackwell/tests/test_msa_blackwell.py
@@ -38,25 +38,48 @@
 """
 
 import argparse
+import importlib
 import sys
 
 import pytest
 import torch
 
-from minimaxai_msa_blackwell import (  # noqa: E402
-    build_k2q_csr,
-    dequantize_nvfp4_128x4_to_bf16,
-    flash_decode_with_gqa_share_sparse,
-    flash_decode_with_topk_idx,
-    has_native_ops,
-    native_nvfp4_dequant_swizzled_to_bf16,
-    native_topk_from_scores,
-    Nvfp4QuantizedTensor,
-    sparse_atten_func,
-    sparse_atten_nvfp4_kv_func,
-    sparse_decode_atten_func,
-)
-from minimaxai_msa_blackwell.prefill.topk_sparse import flash_prefill_with_gqa_share_sparse  # noqa: E402
+
+def _load_module(artifact):
+    """Import the Hub package, optionally from an installed artifact dir."""
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        module = importlib.import_module("minimaxai_msa_blackwell")
+        prefill = importlib.import_module("minimaxai_msa_blackwell.prefill.topk_sparse")
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+    globals().update(
+        {
+            "build_k2q_csr": module.build_k2q_csr,
+            "dequantize_nvfp4_128x4_to_bf16": module.dequantize_nvfp4_128x4_to_bf16,
+            "flash_decode_with_gqa_share_sparse": module.flash_decode_with_gqa_share_sparse,
+            "flash_decode_with_topk_idx": module.flash_decode_with_topk_idx,
+            "has_native_ops": module.has_native_ops,
+            "native_nvfp4_dequant_swizzled_to_bf16": module.native_nvfp4_dequant_swizzled_to_bf16,
+            "native_topk_from_scores": module.native_topk_from_scores,
+            "Nvfp4QuantizedTensor": module.Nvfp4QuantizedTensor,
+            "sparse_atten_func": module.sparse_atten_func,
+            "sparse_atten_nvfp4_kv_func": module.sparse_atten_nvfp4_kv_func,
+            "sparse_decode_atten_func": module.sparse_decode_atten_func,
+            "flash_prefill_with_gqa_share_sparse": prefill.flash_prefill_with_gqa_share_sparse,
+        }
+    )
+    return module
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _pytest_loads_hub_package():
+    """pytest mode loads the Hub package once so the parametrized entrypoints
+    below see the same globals script mode injects from main()."""
+    _load_module(None)
+    yield
 
 DEVICE = "cuda"
 DTYPE = torch.bfloat16
@@ -658,8 +681,14 @@ def test_native_nvfp4_dequant_swizzled_to_bf16(rows, cols):
 # ===========================================================================
 def main():
     ap = argparse.ArgumentParser()
+    ap.add_argument("--backend", choices=["source", "installed"], default="source",
+                    help="the installed Hub package is the supported path for this "
+                         "Triton-based package; source mode falls back to sys.path")
+    ap.add_argument("--artifact", default=None)
+    ap.add_argument("--mode", choices=["smoke", "full"], default="full",
+                    help="smoke skips the 32768-ctx cases")
     ap.add_argument("--quick", action="store_true",
-                    help="skip the 32768-ctx cases")
+                    help="skip the 32768-ctx cases (alias for --mode smoke)")
     ap.add_argument("--long-context", action="store_true",
                     help="also run 65536/131072 standalone long-context cases")
     args = ap.parse_args()
@@ -668,6 +697,10 @@ def main():
         print("CUDA not available; these Triton kernels require a CUDA GPU.")
         return 1
 
+    _load_module(args.artifact)
+    if args.quick:
+        args.mode = "smoke"
+
     print(f"device={torch.cuda.get_device_name()} dtype={DTYPE}")
     print(f"M3 config: Hq={M3_HQ} Hkv={M3_HKV} D={M3_D} block={M3_BLOCK} "
           f"topk={M3_TOPK}")
diff --git a/adaptive-layernorm-producers/SYNC.md b/adaptive-layernorm-producers/SYNC.md
index f304441..36f886a 100644
--- a/adaptive-layernorm-producers/SYNC.md
+++ b/adaptive-layernorm-producers/SYNC.md
@@ -1,5 +1,41 @@
 # Source Sync
 
+- Package: `adaptive-layernorm-producers` (Adaptive/diT LayerNorm and AdaLN-modulation FP8 producers).
+- Upstream FlashRT source: `../official/FlashRT`
+- Upstream revision: pending confirmation; packaged source is maintained in the
+  flashrt-project FlashRT-HF-kernels repository
+  (https://github.com/flashrt-project/FlashRT-HF-kernels).
+
+Copied source files:
+
+- `csrc/ada_layer_norm_fp8.cu`
+- `csrc/ada_layer_norm_fp8.cuh`
+- `csrc/ada_layer_norm_fp8_ptok.cu`
+- `csrc/adaln_modulation6.cu`
+- `csrc/adaln_modulation6.cuh`
+- `csrc/dit_layer_norm_fp8.cu`
+- `csrc/dit_layer_norm_fp8.cuh`
+
+Local packaging edits:
+
+- Added Tensor-facing PyTorch custom ops in `torch-ext/torch_binding.cpp`.
+- Added Python wrappers and fake registrations in `torch-ext/adaptive_layernorm_producers`.
+- Kept public APIs Tensor-facing; no raw pointer or stream arguments.
+- Includes rewritten to be package-local; serving-runtime dependencies removed.
+- CUDA launchers kept graph-safe: no dynamic allocation inside hot kernels.
+
+Architecture assumptions:
+
+- CUDA 12.8+ / 13.0+ (CUDA 13.2 validated on NVIDIA Thor, sm_110a).
+- NVIDIA Blackwell-family targets; Thor sm_110a validated on real hardware.
+
+Runtime constraints:
+
+- Inputs and outputs are `torch.Tensor`; shapes and dtypes are validated in the binding.
+- Benchmarks cap CUDA memory at 30 GB per process via `set_per_process_memory_fraction`.
+
+Additional FP4 producer provenance:
+
 - Upstream: `flashrt-project/FlashRT`
 - GROOT N1.7 sync commit:
   `24df793f4fa2d50780aea03b644208c6e0cb4162`
diff --git a/adaptive-layernorm-producers/benchmarks/benchmark.py b/adaptive-layernorm-producers/benchmarks/benchmark.py
index 40a099a..d52b0fb 100644
--- a/adaptive-layernorm-producers/benchmarks/benchmark.py
+++ b/adaptive-layernorm-producers/benchmarks/benchmark.py
@@ -22,6 +22,16 @@
 )
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def load_installed_ops(artifact: str | None):
     if artifact:
         sys.path.insert(0, artifact)
@@ -175,7 +185,9 @@ def main() -> None:
     parser.add_argument("--artifact", default=None)
     parser.add_argument("--iters", type=int, default=200)
     parser.add_argument("--markdown", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
diff --git a/audio-codebook-primitives/benchmarks/benchmark.py b/audio-codebook-primitives/benchmarks/benchmark.py
index 00a8223..22ae761 100644
--- a/audio-codebook-primitives/benchmarks/benchmark.py
+++ b/audio-codebook-primitives/benchmarks/benchmark.py
@@ -11,6 +11,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def elapsed_us(fn, warmup, iterations):
     for _ in range(warmup):
         fn()
@@ -31,7 +41,9 @@ def main():
     parser.add_argument("--artifact")
     parser.add_argument("--warmup", type=int, default=50)
     parser.add_argument("--iterations", type=int, default=500)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     if args.backend == "source":
         tests = Path(__file__).resolve().parents[1] / "tests"
         sys.path.insert(0, str(tests))
diff --git a/bf16-linear-gemv/SYNC.md b/bf16-linear-gemv/SYNC.md
new file mode 100644
index 0000000..e8a70b9
--- /dev/null
+++ b/bf16-linear-gemv/SYNC.md
@@ -0,0 +1,32 @@
+# Source Sync
+
+- Package: `bf16-linear-gemv` (BF16 M=1 decode GEMV and GEMV-family kernels).
+- Upstream FlashRT source: `../official/FlashRT`
+- Upstream revision: pending confirmation; packaged source is maintained in the
+  flashrt-project FlashRT-HF-kernels repository
+  (https://github.com/flashrt-project/FlashRT-HF-kernels).
+
+Copied source files:
+
+- `csrc/gemm/bf16_gemv_m1_sm120.cu`
+- `csrc/gemm/bf16_gemv_m1_sm120.cuh`
+- `csrc/kernels/nexn2_bf16_gemv.cu`
+- `csrc/kernels/nexn2_bf16_gemv.cuh`
+
+Local packaging edits:
+
+- Added Tensor-facing PyTorch custom ops in `torch-ext/torch_binding.cpp`.
+- Added Python wrappers and fake registrations in `torch-ext/bf16_linear_gemv`.
+- Kept public APIs Tensor-facing; no raw pointer or stream arguments.
+- Includes rewritten to be package-local; serving-runtime dependencies removed.
+- CUDA launchers kept graph-safe: no dynamic allocation inside hot kernels.
+
+Architecture assumptions:
+
+- CUDA 12.8+ / 13.0+ (CUDA 13.2 validated on NVIDIA Thor, sm_110a).
+- NVIDIA Blackwell-family targets; Thor sm_110a validated on real hardware.
+
+Runtime constraints:
+
+- Inputs and outputs are `torch.Tensor`; shapes and dtypes are validated in the binding.
+- Benchmarks cap CUDA memory at 30 GB per process via `set_per_process_memory_fraction`.
diff --git a/bf16-linear-gemv/benchmarks/benchmark.py b/bf16-linear-gemv/benchmarks/benchmark.py
new file mode 100644
index 0000000..fbb89f9
--- /dev/null
+++ b/bf16-linear-gemv/benchmarks/benchmark.py
@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+"""BF16 M=1 decode GEMV benchmark (kernel vs eager vs torch.compile)."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from pathlib import Path
+
+import torch
+
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
+def elapsed_us(fn, warmup: int = 20, repeats: int = 100) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(repeats):
+        fn()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) * 1000.0 / repeats
+
+
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        root = Path(__file__).resolve().parents[1]
+        sys.path.insert(0, str(root / "torch-ext"))
+        return importlib.import_module("bf16_linear_gemv")
+    if artifact:
+        sys.path.insert(0, artifact)
+    return importlib.import_module("bf16_linear_gemv")
+
+
+def run_case(ops, label: str, n: int, k: int) -> dict:
+    gen = torch.Generator(device="cuda").manual_seed(0)
+    x = (torch.randn((k,), device="cuda", generator=gen) * 0.05).to(torch.bfloat16)
+    w = (torch.randn((n, k), device="cuda", generator=gen) * 0.05).to(torch.bfloat16)
+    out = torch.empty((n,), device="cuda", dtype=torch.bfloat16)
+    eager = lambda: x.float() @ w.float().t()
+    compiled = torch.compile(eager, mode="reduce-overhead")
+    compiled()
+
+    def kernel_v0():
+        ops.bf16_decode_gemv_bf16(x, w, out=out)
+
+    def kernel_v1():
+        ops.bf16_decode_gemv_bf16(x, w, variant=16, out=out)
+
+    def kernel_unrolled():
+        ops.bf16_decode_gemv_unrolled_bf16(x, w, out=out)
+
+    kernel_us = elapsed_us(kernel_v0)
+    kernel_v1_us = elapsed_us(kernel_v1)
+    kernel_ur_us = elapsed_us(kernel_unrolled)
+    eager_us = elapsed_us(eager)
+    compile_us = elapsed_us(lambda: compiled())
+    return {
+        "label": label, "N": n, "K": k,
+        "kernel_v0_us": kernel_us, "kernel_v1_us": kernel_v1_us,
+        "kernel_unrolled_us": kernel_ur_us, "eager_us": eager_us,
+        "compile_us": compile_us,
+        "vs_eager_v0": eager_us / kernel_us,
+        "vs_eager_v1": eager_us / kernel_v1_us,
+        "vs_eager_ur": eager_us / kernel_ur_us,
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="installed")
+    parser.add_argument("--artifact")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
+    args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
+    ops = load_ops(args.backend, args.artifact)
+    print("label,N,K,kernel_v0_us,kernel_v1_us,kernel_unrolled_us,eager_us,compile_us,vs_eager_v0,vs_eager_v1,vs_eager_ur")
+    for label, n, k in [
+        ("decode_8k", 8192, 4096),
+        ("decode_12k", 12288, 8192),
+        ("decode_16k", 16384, 8192),
+    ]:
+        r = run_case(ops, label, n, k)
+        print(
+            f"{r['label']},{r['N']},{r['K']},{r['kernel_v0_us']:.3f},"
+            f"{r['kernel_v1_us']:.3f},{r['kernel_unrolled_us']:.3f},"
+            f"{r['eager_us']:.3f},{r['compile_us']:.3f},"
+            f"{r['vs_eager_v0']:.2f}x,{r['vs_eager_v1']:.2f}x,{r['vs_eager_ur']:.2f}x"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/blockwise-fp8-producers/benchmarks/benchmark.py b/blockwise-fp8-producers/benchmarks/benchmark.py
index 450335a..9745e5e 100644
--- a/blockwise-fp8-producers/benchmarks/benchmark.py
+++ b/blockwise-fp8-producers/benchmarks/benchmark.py
@@ -16,6 +16,16 @@
 from test_blockwise_fp8_producers import load_source_ops  # noqa: E402
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def load_ops(backend: str, artifact: str | None):
     if backend == "source":
         return load_source_ops()
@@ -72,7 +82,9 @@ def main() -> int:
     parser.add_argument("--mode", choices=["headline", "full"], default="headline")
     parser.add_argument("--warmup", type=int, default=30)
     parser.add_argument("--iters", type=int, default=200)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     ops = load_ops(args.backend, args.artifact)
     shapes = [(51, 4096), (277, 9216), (1024, 1152)]
     if args.mode == "full":
diff --git a/causal-conv1d-state/benchmarks/benchmark.py b/causal-conv1d-state/benchmarks/benchmark.py
index c41dda6..db2f62f 100644
--- a/causal-conv1d-state/benchmarks/benchmark.py
+++ b/causal-conv1d-state/benchmarks/benchmark.py
@@ -15,6 +15,16 @@
 from test_causal_conv1d_state import MODES, SHAPES, load_installed_ops, load_source_ops, make_inputs, ref_chunk  # noqa: E402
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def time_cuda(fn, warmup: int, iters: int) -> float:
     for _ in range(warmup):
         fn()
@@ -37,7 +47,9 @@ def main() -> int:
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iters", type=int, default=200)
     parser.add_argument("--json-out", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
     rows = []
     for name in MODES[args.mode]:
diff --git a/demos/AOTI-demo-pi05/run_benchmark.py b/demos/AOTI-demo-pi05/run_benchmark.py
index 5188705..dfbda61 100644
--- a/demos/AOTI-demo-pi05/run_benchmark.py
+++ b/demos/AOTI-demo-pi05/run_benchmark.py
@@ -32,6 +32,16 @@
 SEED = 0
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def build_policy():
     from lerobot.policies.pi05.modeling_pi05 import PI05Policy
 
@@ -107,7 +117,9 @@ def main() -> None:
     parser.add_argument("--no-inductor-flags", action="store_true")
     parser.add_argument("--safety", type=float, default=1.0)
     parser.add_argument("--single", action="store_true", help="run only the configured rung, not the full ladder")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
diff --git a/demos/pi05-groot-ffn-epilogue/benchmark.py b/demos/pi05-groot-ffn-epilogue/benchmark.py
index c066ef6..c5a9853 100644
--- a/demos/pi05-groot-ffn-epilogue/benchmark.py
+++ b/demos/pi05-groot-ffn-epilogue/benchmark.py
@@ -27,6 +27,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 PACKAGE = ROOT / "flashrt-gemm-epilogues"
 REGISTRATION_INCLUDE = (
@@ -433,7 +443,9 @@ def main() -> None:
     parser.add_argument("--compile-baseline", action="store_true")
     parser.add_argument("--output", type=Path, default=None)
     parser.add_argument("--markdown", type=Path, default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
diff --git a/demos/pi05-hf-runtime/README.md b/demos/pi05-hf-runtime/README.md
index 48346ff..50e9fde 100644
--- a/demos/pi05-hf-runtime/README.md
+++ b/demos/pi05-hf-runtime/README.md
@@ -36,7 +36,7 @@ Use an environment matching the published Hub artifact. On the local validation
 machine this is the torch 2.11 CUDA 12.8 smoke environment:
 
 ```bash
-/home/heima/suliang/PI/.flashrt-hub-smoke-torch211/bin/python \
+<workspace>/.flashrt-hub-smoke-torch211/bin/python \
   demos/pi05-hf-runtime/benchmark_runtime.py \
   --shape pi05_decoder \
   --warmup 10 \
@@ -47,7 +47,7 @@ machine this is the torch 2.11 CUDA 12.8 smoke environment:
 Try a larger PI0.5 vision-shaped chain:
 
 ```bash
-/home/heima/suliang/PI/.flashrt-hub-smoke-torch211/bin/python \
+<workspace>/.flashrt-hub-smoke-torch211/bin/python \
   demos/pi05-hf-runtime/benchmark_runtime.py \
   --shape pi05_vision \
   --warmup 5 \
diff --git a/demos/pi05-hf-runtime/benchmark_runtime.py b/demos/pi05-hf-runtime/benchmark_runtime.py
index e549196..413a052 100644
--- a/demos/pi05-hf-runtime/benchmark_runtime.py
+++ b/demos/pi05-hf-runtime/benchmark_runtime.py
@@ -14,6 +14,16 @@
 from kernels import get_kernel
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 SHAPES: dict[str, tuple[int, int, int, int, int]] = {
     "pi05_decoder": (10, 1024, 4096, 1024, 18),
     "pi05_vision": (512, 1152, 4304, 1152, 27),
@@ -426,7 +436,9 @@ def main() -> None:
     parser.add_argument("--cosine-limit", type=float, default=0.999)
     parser.add_argument("--output", type=Path)
     parser.add_argument("--markdown", type=Path)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     result = run(args)
     print(json.dumps(asdict(result), indent=2))
diff --git a/demos/runtime-demo/README.md b/demos/runtime-demo/README.md
index 212914d..1767002 100644
--- a/demos/runtime-demo/README.md
+++ b/demos/runtime-demo/README.md
@@ -11,6 +11,9 @@ It is intentionally separate from the package-level demos:
 - the hot path loads kernels once, owns persistent buffers, avoids timed-loop
   allocation, and supports CUDA Graph replay.
 
+> `<workspace>` below is a placeholder for your checkout root (the directory
+> containing the FlashRT-HF-kernels repository and any sibling data trees).
+
 This is not the upstream FlashRT serving runtime. The current Hub path is a
 checkpoint-backed PI0.5 runtime bridge for validating whether Hugging Face
 Kernel Hub packages can drive a clean model pipeline without losing
@@ -67,7 +70,7 @@ OpenPI baseline now live under `demos/runtime-demo/test/`.
 1. **Python env matching a published Hub variant** — `torch 2.11` + `CUDA 12.8`
    with the `kernels` package installed. The Hub packages download automatically
    via `kernels.get_kernel("flashrt/...", version=1)`. Local validation env:
-   `/home/heima/suliang/PI/.flashrt-hub-smoke-torch211/bin/python` (substitute
+   `<workspace>/.flashrt-hub-smoke-torch211/bin/python` (substitute
    your own).
 2. **PI0.5 LIBERO checkpoint** — the OpenPI PI0.5-LIBERO model as PyTorch
    `model.safetensors`. Point `--checkpoint` at the directory holding it
@@ -86,7 +89,7 @@ The default path runs the QKV / O / vision projection GEMMs in FP8 (published
 Hub kernels only):
 
 ```bash
-PY=/home/heima/suliang/PI/.flashrt-hub-smoke-torch211/bin/python
+PY=<workspace>/.flashrt-hub-smoke-torch211/bin/python
 $PY demos/runtime-demo/pi05_hf_decoder_e2e.py run-vision-encoder-decoder \
   --encoder-bundle internal-tests/runtime-demo/pi05-real-images-encoder-x-kv-frame50.pt \
   --checkpoint /path/to/pi05_libero_pytorch \
@@ -127,7 +130,7 @@ The synthetic fixed-shape runtime profiles moved to `test/pi05_runtime_demo.py`.
 They are a microbench of the composed Hub path, not the real-input E2E:
 
 ```bash
-/home/heima/suliang/PI/.flashrt-hub-smoke-torch211/bin/python \
+<workspace>/.flashrt-hub-smoke-torch211/bin/python \
   demos/runtime-demo/test/pi05_runtime_demo.py \
   --profile pi05_hotpath --layers 4 --ffn-activation gelu \
   --attention-backend sdpa --warmup 10 --iters 50 --cuda-graph
@@ -137,7 +140,7 @@ Before a rebuilt GeGLU artifact is uploaded to the Hub, validate the local
 `build/<variant>` directory explicitly:
 
 ```bash
-/home/heima/suliang/PI/.flashrt-hub-smoke-torch211/bin/python \
+<workspace>/.flashrt-hub-smoke-torch211/bin/python \
   demos/runtime-demo/test/pi05_runtime_demo.py \
   --profile pi05_hotpath \
   --layers 4 \
@@ -156,7 +159,7 @@ To validate the rebuilt QKV package before upload, add the local qkv artifact
 and switch the decoder QKV path to the PI0.5 GQA cache API:
 
 ```bash
-/home/heima/suliang/PI/.flashrt-hub-smoke-torch211/bin/python \
+<workspace>/.flashrt-hub-smoke-torch211/bin/python \
   demos/runtime-demo/test/pi05_runtime_demo.py \
   --profile pi05_hotpath \
   --layers 1 \
@@ -178,10 +181,10 @@ For an end-to-end staging report that keeps the baselines separate:
 python demos/runtime-demo/test/pi05_e2e_runner.py \
   --openpi-baseline-mode docker \
   --container pi0-stablehlo-test \
-  --container-repo /workspace/PI/FlashRT-HF-kernels \
-  --container-python /workspace/PI/FlashRT-HF-kernels/internal-tests/envs/openpi-baseline/bin/python \
-  --container-openpi-root /workspace/PI/openpi_src/src \
-  --container-checkpoint /workspace/PI/checkpoints/pi05_libero_pytorch \
+  --container-repo <workspace>/FlashRT-HF-kernels \
+  --container-python <workspace>/FlashRT-HF-kernels/internal-tests/envs/openpi-baseline/bin/python \
+  --container-openpi-root <workspace>/openpi_src/src \
+  --container-checkpoint <workspace>/checkpoints/pi05_libero_pytorch \
   --cuda-graph \
   --output internal-tests/runtime-demo/pi05-e2e-staging.json
 ```
@@ -202,11 +205,11 @@ To run only the official OpenPI/PyTorch baseline inside the local container:
 
 ```bash
 docker exec pi0-stablehlo-test bash -lc '
-cd /workspace/PI/FlashRT-HF-kernels &&
-PYTHONPATH=/workspace/PI/openpi_src/src \
+cd <workspace>/FlashRT-HF-kernels &&
+PYTHONPATH=<workspace>/openpi_src/src \
 python3 demos/runtime-demo/test/pi05_openpi_baseline.py \
-  --openpi-root /workspace/PI/openpi_src/src \
-  --checkpoint /workspace/PI/checkpoints/pi05_libero_pytorch \
+  --openpi-root <workspace>/openpi_src/src \
+  --checkpoint <workspace>/checkpoints/pi05_libero_pytorch \
   --num-views 2 \
   --steps 10 \
   --warmup 5 \
@@ -222,12 +225,12 @@ global Python. The local venv is intentionally under ignored `internal-tests/`:
 ```bash
 docker exec pi0-stablehlo-test bash -lc '
 set -euo pipefail
-VENV=/workspace/PI/FlashRT-HF-kernels/internal-tests/envs/openpi-baseline
+VENV=<workspace>/FlashRT-HF-kernels/internal-tests/envs/openpi-baseline
 python3 -m venv --system-site-packages "$VENV"
 "$VENV/bin/python" -m pip install "transformers==4.53.2"
 "$VENV/bin/python" - <<PY
 import pathlib, shutil, transformers
-src = pathlib.Path("/workspace/PI/openpi_src/src/openpi/models_pytorch/transformers_replace")
+src = pathlib.Path("<workspace>/openpi_src/src/openpi/models_pytorch/transformers_replace")
 dst = pathlib.Path(transformers.__file__).resolve().parent
 for item in src.iterdir():
     target = dst / item.name
@@ -251,12 +254,12 @@ static activation scales, first capture real OpenPI decoder activations:
 
 ```bash
 docker exec pi0-stablehlo-test bash -lc '
-cd /workspace/PI/FlashRT-HF-kernels &&
-PYTHONPATH=/workspace/PI/openpi_src/src \
-/workspace/PI/FlashRT-HF-kernels/internal-tests/envs/openpi-baseline/bin/python \
+cd <workspace>/FlashRT-HF-kernels &&
+PYTHONPATH=<workspace>/openpi_src/src \
+<workspace>/FlashRT-HF-kernels/internal-tests/envs/openpi-baseline/bin/python \
   demos/runtime-demo/pi05_capture_openpi_ffn_activations.py \
-  --openpi-root /workspace/PI/openpi_src/src \
-  --checkpoint /workspace/PI/checkpoints/pi05_libero_pytorch \
+  --openpi-root <workspace>/openpi_src/src \
+  --checkpoint <workspace>/checkpoints/pi05_libero_pytorch \
   --family decoder \
   --layer 0 \
   --num-views 2 \
@@ -268,9 +271,9 @@ PYTHONPATH=/workspace/PI/openpi_src/src \
 Then run the Hub-kernel FFN island in the HF kernel environment:
 
 ```bash
-/home/heima/suliang/PI/.flashrt-hub-smoke-torch211/bin/python \
+<workspace>/.flashrt-hub-smoke-torch211/bin/python \
   demos/runtime-demo/test/pi05_real_weight_swiglu.py \
-  --checkpoint /home/heima/suliang/PI/checkpoints/pi05_libero_pytorch \
+  --checkpoint <workspace>/checkpoints/pi05_libero_pytorch \
   --family decoder \
   --layer 0 \
   --rows 100 \
@@ -290,9 +293,9 @@ weights, official FlashRT-style time/style precompute, the rebuilt QKV cache
 kernel, FP8 GeGLU, SDPA attention, and CUDA Graph:
 
 ```bash
-/home/heima/suliang/PI/.flashrt-hub-smoke-torch211/bin/python \
+<workspace>/.flashrt-hub-smoke-torch211/bin/python \
   demos/runtime-demo/pi05_decoder_loop_hub.py \
-  --checkpoint /home/heima/suliang/PI/checkpoints/pi05_libero_pytorch \
+  --checkpoint <workspace>/checkpoints/pi05_libero_pytorch \
   --layers 18 \
   --steps 10 \
   --local-qkv-artifact flashrt-qkv-cache-rope/build/torch211-cxx11-cu128-x86_64-linux \
@@ -323,12 +326,12 @@ path and replaces selected Gemma MLP layers:
 
 ```bash
 docker exec pi0-stablehlo-test bash -lc '
-cd /workspace/PI/FlashRT-HF-kernels &&
-PYTHONPATH=/workspace/PI/openpi_src/src \
-/workspace/PI/FlashRT-HF-kernels/internal-tests/envs/openpi-baseline/bin/python \
+cd <workspace>/FlashRT-HF-kernels &&
+PYTHONPATH=<workspace>/openpi_src/src \
+<workspace>/FlashRT-HF-kernels/internal-tests/envs/openpi-baseline/bin/python \
   demos/runtime-demo/test/pi05_openpi_hub_ffn_e2e.py \
-  --openpi-root /workspace/PI/openpi_src/src \
-  --checkpoint /workspace/PI/checkpoints/pi05_libero_pytorch \
+  --openpi-root <workspace>/openpi_src/src \
+  --checkpoint <workspace>/checkpoints/pi05_libero_pytorch \
   --num-views 2 \
   --steps 10 \
   --warmup 1 \
diff --git a/demos/runtime-demo/dyn_full.py b/demos/runtime-demo/dyn_full.py
index 0586bf9..44e5161 100755
--- a/demos/runtime-demo/dyn_full.py
+++ b/demos/runtime-demo/dyn_full.py
@@ -7,6 +7,15 @@
 HERE = Path(__file__).resolve().parent
 ROOT = HERE.parents[1]
 FP8=448.0
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
 class DynDecoder(dec.HubDecoderLoop):
     def _sc(self,a): return torch.clamp(a/FP8,min=1e-12).reshape(1).to(self.w.device,torch.float32)
     def __call__(self):
@@ -59,7 +68,9 @@ def __call__(self):
 parser.add_argument("--warmup", type=int, default=8)
 parser.add_argument("--iters", type=int, default=30)
 parser.add_argument("--no-cuda-graph", action="store_true")
+parser.add_argument("--max-mem-gb", type=float, default=30.0)
 args = parser.parse_args()
+_apply_mem_cap(args.max_mem_gb)
 mode=args.mode
 if mode=="dynamic": dec.HubDecoderLoop=DynDecoder
 import pi05_hf_decoder_e2e as e2e
diff --git a/demos/runtime-demo/dyn_full_all_fp8.py b/demos/runtime-demo/dyn_full_all_fp8.py
index 9739dc9..ee34238 100755
--- a/demos/runtime-demo/dyn_full_all_fp8.py
+++ b/demos/runtime-demo/dyn_full_all_fp8.py
@@ -13,6 +13,17 @@
 import pi05_decoder_loop_hub as dec  # noqa: E402
 import pi05_hf_decoder_e2e as e2e  # noqa: E402
 
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 FP8_MAX = 448.0
 
 
@@ -217,7 +228,9 @@ def main() -> None:
     parser.add_argument("--warmup", type=int, default=8)
     parser.add_argument("--iters", type=int, default=30)
     parser.add_argument("--no-cuda-graph", action="store_true")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     mode = args.mode
     if mode == "dynamic":
diff --git a/demos/runtime-demo/pi05_decoder_loop_hub.py b/demos/runtime-demo/pi05_decoder_loop_hub.py
index ef32c3f..a66d3e5 100644
--- a/demos/runtime-demo/pi05_decoder_loop_hub.py
+++ b/demos/runtime-demo/pi05_decoder_loop_hub.py
@@ -22,6 +22,7 @@
 import importlib
 import json
 import math
+import os
 import sys
 from dataclasses import asdict, dataclass
 from pathlib import Path
@@ -33,6 +34,16 @@
 from safetensors import safe_open
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 DEC_L = 18
 DEC_D = 1024
 DEC_H = 4096
@@ -927,7 +938,7 @@ def run(args: argparse.Namespace) -> Result:
 
 def main() -> None:
     parser = argparse.ArgumentParser()
-    parser.add_argument("--checkpoint", default="/home/heima/suliang/PI/checkpoints/pi05_libero_pytorch")
+    parser.add_argument("--checkpoint", default=os.environ.get("PI05_CHECKPOINT_DIR", ""))
     parser.add_argument("--layers", type=int, default=1)
     parser.add_argument("--steps", type=int, default=10)
     parser.add_argument("--encoder-seq-len", type=int, default=560)
@@ -946,7 +957,9 @@ def main() -> None:
     parser.add_argument("--p99-abs-limit", type=float, default=0.5)
     parser.add_argument("--cosine-limit", type=float, default=0.9)
     parser.add_argument("--output", type=Path)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     if not 1 <= args.layers <= DEC_L:
         raise ValueError(f"--layers must be in [1, {DEC_L}]")
     if args.steps <= 0:
diff --git a/demos/runtime-demo/pi05_hf_decoder_e2e.py b/demos/runtime-demo/pi05_hf_decoder_e2e.py
index 8067482..b81a087 100644
--- a/demos/runtime-demo/pi05_hf_decoder_e2e.py
+++ b/demos/runtime-demo/pi05_hf_decoder_e2e.py
@@ -34,6 +34,16 @@
 import torch.nn.functional as F
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 PI_ROOT = ROOT.parent
 DEFAULT_CKPT = PI_ROOT / "checkpoints" / "pi05_libero_pytorch"
@@ -1717,6 +1727,7 @@ def run_hf_vision_encoder_decoder(args: argparse.Namespace) -> BridgeResult:
 def main() -> None:
     parser = argparse.ArgumentParser()
     sub = parser.add_subparsers(dest="cmd", required=True)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
 
     exp = sub.add_parser("export-encoder")
     exp.add_argument("--checkpoint", type=Path, default=DEFAULT_CKPT)
@@ -1815,6 +1826,7 @@ def main() -> None:
     visencdec.add_argument("--output", type=Path)
 
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     if args.cmd == "export-encoder":
         export_encoder(args)
         return
diff --git a/demos/runtime-demo/sd_decoder.py b/demos/runtime-demo/sd_decoder.py
index b21bb51..616dfbf 100755
--- a/demos/runtime-demo/sd_decoder.py
+++ b/demos/runtime-demo/sd_decoder.py
@@ -5,6 +5,15 @@
 sys.path.insert(0, str(HERE))
 import pi05_decoder_loop_hub as dec
 FP8=448.0
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
 class DynDecoder(dec.HubDecoderLoop):
     """Same as base (BF16 QKV/O, FP8 FFN) but FFN scales computed per-forward -> split GeGLU."""
     def _sc(self, a): return torch.clamp(a/FP8, min=1e-12).reshape(1).to(self.w.device, torch.float32)
@@ -51,7 +60,9 @@ def __call__(self):
 parser.add_argument("--calibration-input", default=str(ROOT / "internal-tests/runtime-demo/pi05-decoder-loop-hub-static-scales.json"))
 parser.add_argument("--warmup", type=int, default=8)
 parser.add_argument("--iters", type=int, default=30)
+parser.add_argument("--max-mem-gb", type=float, default=30.0)
 args = parser.parse_args()
+_apply_mem_cap(args.max_mem_gb)
 B=args.encoder_kv_bundle
 CK=args.checkpoint or os.environ.get("PI05_CHECKPOINT") or str(ROOT.parent / "checkpoints/pi05_libero_pytorch")
 CAL=args.calibration_input
diff --git a/demos/runtime-demo/test/pi05_e2e_runner.py b/demos/runtime-demo/test/pi05_e2e_runner.py
index 6c653f9..1592168 100644
--- a/demos/runtime-demo/test/pi05_e2e_runner.py
+++ b/demos/runtime-demo/test/pi05_e2e_runner.py
@@ -31,13 +31,12 @@
 DEFAULT_CKPT = PI_ROOT / "checkpoints" / "pi05_libero_pytorch"
 DEFAULT_HUB_PY = PI_ROOT / ".flashrt-hub-smoke-torch211" / "bin" / "python"
 DEFAULT_OPENPI_ROOT = PI_ROOT / "openpi_src" / "src"
-DEFAULT_CONTAINER_REPO = "/workspace/PI/FlashRT-HF-kernels"
-DEFAULT_CONTAINER_OPENPI_ROOT = "/workspace/PI/openpi_src/src"
-DEFAULT_CONTAINER_CKPT = "/workspace/PI/checkpoints/pi05_libero_pytorch"
-DEFAULT_CONTAINER_FLASHRT_ROOT = "/workspace/PI/official/FlashRT"
-DEFAULT_CONTAINER_OPENPI_PY = (
-    "/workspace/PI/FlashRT-HF-kernels/internal-tests/envs/"
-    "openpi-baseline/bin/python"
+DEFAULT_CONTAINER_REPO = str(ROOT)
+DEFAULT_CONTAINER_OPENPI_ROOT = str(PI_ROOT / "openpi_src" / "src")
+DEFAULT_CONTAINER_CKPT = str(PI_ROOT / "checkpoints" / "pi05_libero_pytorch")
+DEFAULT_CONTAINER_FLASHRT_ROOT = str(PI_ROOT / "official" / "FlashRT")
+DEFAULT_CONTAINER_OPENPI_PY = str(
+    ROOT / "internal-tests" / "envs" / "openpi-baseline" / "bin" / "python"
 )
 
 
diff --git a/demos/runtime-demo/test/pi05_openpi_hub_ffn_e2e.py b/demos/runtime-demo/test/pi05_openpi_hub_ffn_e2e.py
index 9d0e5c3..d2299be 100644
--- a/demos/runtime-demo/test/pi05_openpi_hub_ffn_e2e.py
+++ b/demos/runtime-demo/test/pi05_openpi_hub_ffn_e2e.py
@@ -21,6 +21,7 @@
 import argparse
 import json
 import math
+import os
 import shutil
 import statistics
 import sys
@@ -152,9 +153,17 @@ def _compile_direct_extension(repo_dir: Path, name: str):
     import torch
     from torch.utils.cpp_extension import load
 
-    builder_root = Path("/home/heima/suliang/PI/kernels/kernel-builder/src/pyproject/templates/torch")
-    if not builder_root.exists():
-        builder_root = Path("/workspace/PI/kernels/kernel-builder/src/pyproject/templates/torch")
+    builder_root = Path(os.environ.get("KERNEL_BUILDER_REGISTRATION_INCLUDE", ""))
+    if not builder_root.is_dir():
+        builder_root = (
+            repo_dir.parent.parent
+            / "kernels" / "kernel-builder" / "src" / "pyproject" / "templates" / "torch"
+        )
+    if not builder_root.is_dir():
+        raise RuntimeError(
+            "kernel-builder registration include not found; set "
+            "KERNEL_BUILDER_REGISTRATION_INCLUDE"
+        )
     sources = [
         repo_dir / "torch-ext/torch_binding.cpp",
         *sorted((repo_dir / "csrc").glob("*.cu")),
diff --git a/demos/static-vs-dynamic-fp8/geglu_static_dynamic_microbench.py b/demos/static-vs-dynamic-fp8/geglu_static_dynamic_microbench.py
index 06bf2ff..8840f2f 100755
--- a/demos/static-vs-dynamic-fp8/geglu_static_dynamic_microbench.py
+++ b/demos/static-vs-dynamic-fp8/geglu_static_dynamic_microbench.py
@@ -16,6 +16,16 @@
 from kernels import get_kernel
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 FP8_MAX = 448.0
 
 
@@ -183,7 +193,9 @@ def main() -> None:
     parser.add_argument("--H", type=int, default=16384)
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iters", type=int, default=100)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required")
diff --git a/demos/static-vs-dynamic-fp8/run_static_dynamic_fp8.py b/demos/static-vs-dynamic-fp8/run_static_dynamic_fp8.py
index 3e58bc6..74a8b84 100755
--- a/demos/static-vs-dynamic-fp8/run_static_dynamic_fp8.py
+++ b/demos/static-vs-dynamic-fp8/run_static_dynamic_fp8.py
@@ -55,11 +55,13 @@ def main() -> None:
     )
     parser.add_argument("--microbench-warmup", type=int, default=20)
     parser.add_argument("--microbench-iters", type=int, default=100)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
 
     out_dir = Path(args.out_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
     py = sys.executable
+    mem_args = ["--max-mem-gb", str(args.max_mem_gb)]
     summary = []
 
     if args.suite in {"all", "all-fp8-e2e"}:
@@ -77,6 +79,7 @@ def main() -> None:
                         str(args.warmup),
                         "--iters",
                         str(args.iters),
+                        *mem_args,
                     ],
                     out_dir,
                 )
@@ -97,6 +100,7 @@ def main() -> None:
                         str(args.warmup),
                         "--iters",
                         str(args.iters),
+                        *mem_args,
                     ],
                     out_dir,
                 )
@@ -115,6 +119,7 @@ def main() -> None:
                     str(args.warmup),
                     "--iters",
                     str(args.iters),
+                    *mem_args,
                 ],
                 out_dir,
             )
@@ -133,6 +138,7 @@ def main() -> None:
                     str(args.microbench_warmup),
                     "--iters",
                     str(args.microbench_iters),
+                    *mem_args,
                 ],
                 out_dir,
             )
diff --git a/demos/wan-qkv-postprocess/benchmark.py b/demos/wan-qkv-postprocess/benchmark.py
index f8bd29a..bf69151 100644
--- a/demos/wan-qkv-postprocess/benchmark.py
+++ b/demos/wan-qkv-postprocess/benchmark.py
@@ -18,6 +18,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 PACKAGE = ROOT / "flashrt-vla-video"
 REGISTRATION_INCLUDE = (
@@ -799,7 +809,9 @@ def main() -> int:
     parser.add_argument("--compile-mode", default="default")
     parser.add_argument("--output", default="internal-tests/demos/wan-qkv-postprocess/results.json")
     parser.add_argument("--markdown", default="internal-tests/demos/wan-qkv-postprocess/results.md")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
diff --git a/diffusion-step-ops/benchmarks/benchmark.py b/diffusion-step-ops/benchmarks/benchmark.py
index ac770e6..08db9c8 100644
--- a/diffusion-step-ops/benchmarks/benchmark.py
+++ b/diffusion-step-ops/benchmarks/benchmark.py
@@ -15,6 +15,16 @@
 from test_diffusion_step_ops import load_installed_ops, load_source_ops  # noqa: E402
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def bench(fn, warmup: int, iters: int) -> float:
     for _ in range(warmup):
         fn()
@@ -35,7 +45,9 @@ def main() -> int:
     parser.add_argument("--artifact", default=None)
     parser.add_argument("--warmup", type=int, default=100)
     parser.add_argument("--iters", type=int, default=1000)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required")
diff --git a/docs/environment.md b/docs/environment.md
index 3788d94..c3fe123 100644
--- a/docs/environment.md
+++ b/docs/environment.md
@@ -37,6 +37,23 @@ PYTHONPATH=../official/FlashRT pytest internal-tests
 These tests are for source sync confidence and FlashRT parity. They are not
 Hub-compatible CI tests and should not be copied into package `tests/`.
 
+## SM110 (Thor) Triton note
+
+Triton bundles its own `ptxas` (e.g. CUDA 12.8 build) which does not accept
+`--gpu-name=sm_110a`, so any Triton kernel fails to compile on Thor with
+`ptxas fatal: Value 'sm_110a' is not defined for option 'gpu-name'`. Point
+Triton at a ptxas that knows `sm_110a` (the system CUDA 13.2 `ptxas` works):
+
+```bash
+ln -sf /usr/local/cuda-13.2/bin/ptxas \
+  <venv>/lib/python3.10/site-packages/triton/backends/nvidia/bin/ptxas
+# or per-invocation:
+TRITON_PTXAS_PATH=/usr/local/cuda-13.2/bin/ptxas python ...
+```
+
+Without this, Triton-based paths (e.g. MiniMaxAI-msa-blackwell decode/indexer
+tests) fail to JIT on SM110 even though the native CUDA kernels are correct.
+
 ## Dependency Policy
 
 - Do not reuse FlashRT editable install state as a hidden dependency.
diff --git a/fa2-seqused-runtime/benchmarks/benchmark.py b/fa2-seqused-runtime/benchmarks/benchmark.py
index bdcab37..cb87504 100644
--- a/fa2-seqused-runtime/benchmarks/benchmark.py
+++ b/fa2-seqused-runtime/benchmarks/benchmark.py
@@ -9,6 +9,16 @@
 from fa2_seqused_runtime import allocate_outputs, allocate_workspace, forward_static
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 SHAPES = [
     # GROOT DiT self/cross attention.
     ("groot-dit-self", 1, 51, 51, 32, 32, 48, False),
@@ -47,7 +57,9 @@ def time_us(fn, warmup=50, repeats=200):
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--dtype", choices=("bf16", "fp16"), default="bf16")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float16
     print(
         "Workload,Mode,B,Sq,Sk,Hq,Hkv,D,FlashRT_us,SDPA_expandedGQA_us,Speedup,"
diff --git a/fa2-seqused-runtime/tests/test_fa2_seqused_runtime.py b/fa2-seqused-runtime/tests/test_fa2_seqused_runtime.py
index b978d15..067463b 100644
--- a/fa2-seqused-runtime/tests/test_fa2_seqused_runtime.py
+++ b/fa2-seqused-runtime/tests/test_fa2_seqused_runtime.py
@@ -1,27 +1,218 @@
+#!/usr/bin/env python3
+"""Correctness tests for fa2-seqused-runtime (static-buffer FA2 forward)."""
 from __future__ import annotations
 
+import argparse
+import importlib
+import json
 import math
+import os
+import sys
+from pathlib import Path
 
-import pytest
 import torch
 import torch.nn.functional as F
 
-from fa2_seqused_runtime import (
-    FA2Workspace,
-    SUPPORTED_HEAD_DIMS,
-    SPLIT_HEAD_DIMS,
-    allocate_outputs,
-    allocate_workspace,
-    forward,
-    forward_seqused_static,
-    forward_static,
+
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "fa2-seqused-runtime"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
 )
+CUTLASS_INCLUDE = Path(os.environ.get("FA2_CUTLASS_INCLUDE", ""))
+
+SUPPORTED_HEAD_DIMS = tuple(range(8, 257, 8))
+SPLIT_HEAD_DIMS = tuple(range(40, 129, 8)) + tuple(range(232, 257, 8))
+
+_FA2_SOURCES = [
+    "torch-ext/torch_binding.cpp",
+    "csrc/fa2_wrapper.cu",
+    "csrc/fa2_wrapper_causal.cu",
+    "csrc/flash_attn/flash_fwd_hdim64_fp16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_hdim64_bf16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_hdim96_fp16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_hdim96_bf16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_hdim128_fp16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_hdim128_bf16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_hdim256_fp16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_hdim256_bf16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_split_hdim64_fp16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_split_hdim64_bf16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_split_hdim96_fp16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_split_hdim96_bf16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_split_hdim128_fp16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_split_hdim128_bf16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_split_hdim256_fp16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_split_hdim256_bf16_sm80.cu",
+    "csrc/flash_attn/flash_fwd_hdim128_bf16_sm80_causal.cu",
+    "csrc/flash_attn/flash_fwd_split_hdim128_bf16_sm80_causal.cu",
+    "csrc/flash_attn/flash_fwd_hdim256_bf16_sm80_causal.cu",
+    "csrc/flash_attn/flash_fwd_split_hdim256_bf16_sm80_causal.cu",
+]
+
+_FA2_CUDA_FLAGS = [
+    "-O3",
+    "--use_fast_math",
+    "--expt-relaxed-constexpr",
+    "--expt-extended-lambda",
+    "-U__CUDA_NO_HALF_OPERATORS__",
+    "-U__CUDA_NO_HALF_CONVERSIONS__",
+    "-U__CUDA_NO_BFLOAT16_CONVERSIONS__",
+    "-U__CUDA_NO_BFLOAT16_OPERATORS__",
+    "-U__CUDA_NO_BFLOAT162_OPERATORS__",
+    "-DFA2_HAS_HDIM_64=1",
+    "-DFA2_HAS_HDIM_96=1",
+    "-DFA2_HAS_HDIM_128=1",
+    "-DFA2_HAS_HDIM_256=1",
+    "-DFA2_HAS_FP16=1",
+    "-DFA2_HAS_BF16=1",
+]
+
+
+def _arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if major >= 12:
+        return "12.0a"
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    return f"{major}.{minor}"
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+
+    def allocate_outputs(self, q):
+        out = torch.empty_strided(q.shape, q.stride(), device=q.device, dtype=q.dtype)
+        lse = torch.empty((q.shape[0], q.shape[2], q.shape[1]), device=q.device, dtype=torch.float32)
+        return out, lse
+
+    def forward(self, q, k, v, *, softmax_scale=None, causal=False, use_split_kv=False):
+        out, lse = self.allocate_outputs(q)
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** -0.5
+        self._ops.forward_static(
+            q, k, v, out, lse, None, None, float(softmax_scale), bool(causal), 0
+        )
+        return out
+
+    def forward_static(self, q, k, v, *, out, softmax_lse, workspace=None, softmax_scale=None, causal=False):
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** -0.5
+        lse_accum = None
+        out_accum = None
+        num_sms = 0
+        if workspace is not None:
+            lse_accum, out_accum, num_sms = workspace
+        self._ops.forward_static(
+            q, k, v, out, softmax_lse, lse_accum, out_accum,
+            float(softmax_scale), bool(causal), int(num_sms),
+        )
+        return out
+
+    def forward_seqused_static(self, q, k, v, seqused_k, *, out, softmax_lse, workspace=None, softmax_scale=None):
+        if softmax_scale is None:
+            softmax_scale = q.shape[-1] ** -0.5
+        lse_accum = None
+        out_accum = None
+        num_sms = 0
+        if workspace is not None:
+            lse_accum, out_accum, num_sms = workspace
+        if lse_accum is not None:
+            lse_accum.fill_(-torch.inf)
+        self._ops.forward_seqused_static(
+            q, k, v, seqused_k, out, softmax_lse, lse_accum, out_accum,
+            float(softmax_scale), int(num_sms),
+        )
+        return out
+
+    def allocate_workspace(self, q, k, *, num_sms=None):
+        # The installed package computes the heuristic in Python; the source
+        # backend does not expose it. Always return None so the no-split path
+        # is exercised in source mode.
+        return None
+
+
+class InstalledOps:
+    def __init__(self, module) -> None:
+        self._module = module
+
+    def allocate_outputs(self, q):
+        return self._module.allocate_outputs(q)
+
+    def forward(self, q, k, v, *, softmax_scale=None, causal=False, use_split_kv=True):
+        return self._module.forward(q, k, v, softmax_scale=softmax_scale, causal=causal, use_split_kv=use_split_kv)
+
+    def forward_static(self, q, k, v, *, out, softmax_lse, workspace=None, softmax_scale=None, causal=False):
+        return self._module.forward_static(q, k, v, out=out, softmax_lse=softmax_lse, workspace=workspace, softmax_scale=softmax_scale, causal=causal)
+
+    def forward_seqused_static(self, q, k, v, seqused_k, *, out, softmax_lse, workspace=None, softmax_scale=None):
+        return self._module.forward_seqused_static(q, k, v, seqused_k, out=out, softmax_lse=softmax_lse, workspace=workspace, softmax_scale=softmax_scale)
+
+    def allocate_workspace(self, q, k, *, num_sms=None):
+        return self._module.allocate_workspace(q, k, num_sms=num_sms)
+
+
+def _preload_cublaslt() -> None:
+    import ctypes
+    import ctypes.util
+
+    for parent in Path(torch.__file__).resolve().parents:
+        candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12"
+        if candidate.exists():
+            ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL)
+            return
+    library = ctypes.util.find_library("cublasLt")
+    if library:
+        ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL)
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    if not CUTLASS_INCLUDE.is_dir():
+        raise RuntimeError(
+            f"set FA2_CUTLASS_INCLUDE to a CUTLASS include directory "
+            f"(got {CUTLASS_INCLUDE!r})"
+        )
+    _preload_cublaslt()
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
+    namespace = "fa2_seqused_runtime_test"
+    load(
+        name=namespace,
+        sources=[str(PACKAGE / s) for s in _FA2_SOURCES],
+        extra_include_paths=[
+            str(PACKAGE / "csrc"),
+            str(PACKAGE / "csrc" / "flash_attn"),
+            str(CUTLASS_INCLUDE),
+            str(REGISTRATION_INCLUDE),
+        ],
+        extra_cflags=["-O3", "-DCUDA_KERNEL", "-DFA2_HAS_HDIM_128=1", "-DFA2_HAS_HDIM_256=1", "-DFA2_HAS_BF16=1", "-DFA2_HAS_FP16=1"],
+        extra_cuda_cflags=["-DCUDA_KERNEL", *_FA2_CUDA_FLAGS],
+        verbose=False,
+    )
+    return SourceOps(namespace)
 
 
-pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return InstalledOps(importlib.import_module("fa2_seqused_runtime"))
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
 
 
-def _reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool = False) -> torch.Tensor:
+def _reference(q, k, v, causal=False):
     hq, hkv = q.shape[2], k.shape[2]
     if hq != hkv:
         repeat = hq // hkv
@@ -34,71 +225,58 @@ def _reference(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, causal: bool =
         k_idx = torch.arange(sk, device=q.device).view(1, sk)
         attn_mask = k_idx <= q_idx + sk - sq
     return F.scaled_dot_product_attention(
-        q.permute(0, 2, 1, 3),
-        k.permute(0, 2, 1, 3),
-        v.permute(0, 2, 1, 3),
-        attn_mask=attn_mask,
-        is_causal=False,
+        q.permute(0, 2, 1, 3), k.permute(0, 2, 1, 3), v.permute(0, 2, 1, 3),
+        attn_mask=attn_mask, is_causal=False,
     ).permute(0, 2, 1, 3)
 
 
-def _assert_close(actual: torch.Tensor, expected: torch.Tensor) -> None:
+def _assert_close(actual, expected):
     diff = (actual.float() - expected.float()).abs()
     cosine = F.cosine_similarity(actual.float().flatten(), expected.float().flatten(), dim=0)
     if actual.dtype == torch.float16:
-        assert diff.max().item() <= 8e-3
-        assert diff.mean().item() <= 4e-4
-        assert cosine.item() >= 0.9999
+        assert diff.max().item() <= 8e-3, f"max_abs {diff.max().item()}"
+        assert diff.mean().item() <= 4e-4, f"mean_abs {diff.mean().item()}"
+        assert cosine.item() >= 0.9999, f"cosine {cosine.item()}"
     else:
-        assert diff.max().item() <= 6.5e-2
-        assert diff.mean().item() <= 3e-3
-        assert cosine.item() >= 0.999
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("head_dim", [48, 64, 72, 80, 96, 128, 256])
-@pytest.mark.parametrize(
-    "shape",
-    [
-        (1, 1, 1, 8, 8),
-        (1, 17, 63, 16, 4),
-        (2, 64, 129, 12, 4),
-        (1, 257, 511, 8, 2),
-    ],
-)
-def test_noncausal_full_matrix(dtype, head_dim, shape):
+        assert diff.max().item() <= 6.5e-2, f"max_abs {diff.max().item()}"
+        assert diff.mean().item() <= 3e-3, f"mean_abs {diff.mean().item()}"
+        assert cosine.item() >= 0.999, f"cosine {cosine.item()}"
+
+
+def _check_noncausal(ops, dtype, head_dim, shape):
     batch, sq, sk, hq, hkv = shape
     q = torch.randn(batch, sq, hq, head_dim, device="cuda", dtype=dtype) * 0.5
     k = torch.randn(batch, sk, hkv, head_dim, device="cuda", dtype=dtype) * 0.5
     v = torch.randn_like(k)
-    actual = forward(q, k, v, use_split_kv=False)
-    expected = _reference(q, k, v)
-    _assert_close(actual, expected)
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("head_dim", SUPPORTED_HEAD_DIMS)
-def test_all_aligned_logical_head_dims(dtype, head_dim):
-    q = torch.randn(1, 7, 4, head_dim, device="cuda", dtype=dtype) * 0.5
-    k = torch.randn(1, 13, 2, head_dim, device="cuda", dtype=dtype) * 0.5
-    v = torch.randn_like(k)
-    _assert_close(
-        forward(q, k, v, use_split_kv=False),
-        _reference(q, k, v),
-    )
+    actual = ops.forward(q, k, v, use_split_kv=False)
+    _assert_close(actual, _reference(q, k, v))
+    print(f"PASS noncausal dtype={dtype} head_dim={head_dim} shape={shape}")
 
 
-@pytest.mark.parametrize("head_dim", [128, 256])
-@pytest.mark.parametrize("seqlen", [1, 17, 128, 257, 1024])
-def test_causal_bf16(head_dim, seqlen):
+def _check_causal(ops, head_dim, seqlen):
     q = torch.randn(1, seqlen, 8, head_dim, device="cuda", dtype=torch.bfloat16) * 0.5
     k = torch.randn(1, seqlen, 2, head_dim, device="cuda", dtype=torch.bfloat16) * 0.5
     v = torch.randn_like(k)
-    actual = forward(q, k, v, causal=True, use_split_kv=False)
+    actual = ops.forward(q, k, v, causal=True, use_split_kv=False)
     _assert_close(actual, _reference(q, k, v, causal=True))
+    print(f"PASS causal head_dim={head_dim} seqlen={seqlen}")
 
 
-def test_aligned_padded_strides_bf16():
+def _check_seqused(ops):
+    q = torch.randn(2, 17, 8, 128, device="cuda", dtype=torch.bfloat16) * 0.5
+    k = torch.randn(2, 513, 2, 128, device="cuda", dtype=torch.bfloat16) * 0.5
+    v = torch.randn_like(k)
+    used = torch.tensor([127, 513], device="cuda", dtype=torch.int32)
+    out, lse = ops.allocate_outputs(q)
+    ops.forward_seqused_static(q, k, v, used, out=out, softmax_lse=lse)
+    refs = []
+    for batch, n_used in enumerate((127, 513)):
+        refs.append(_reference(q[batch : batch + 1], k[batch : batch + 1, :n_used], v[batch : batch + 1, :n_used]))
+    _assert_close(out, torch.cat(refs, dim=0))
+    print("PASS device seqused per-batch")
+
+
+def _check_padded_strides(ops):
     def padded(shape):
         storage = torch.randn(*shape[:-1], shape[-1] + 8, device="cuda", dtype=torch.bfloat16)
         return storage[..., : shape[-1]]
@@ -108,150 +286,109 @@ def padded(shape):
     v = padded((1, 257, 2, 128))
     out = torch.empty_strided(q.shape, q.stride(), device=q.device, dtype=q.dtype)
     lse = torch.empty((1, 8, 49), device="cuda", dtype=torch.float32)
-    forward_static(q, k, v, out=out, softmax_lse=lse)
-    _assert_close(out, _reference(q, k, v))
-
-
-def test_device_seqused_per_batch():
-    q = torch.randn(2, 17, 8, 128, device="cuda", dtype=torch.bfloat16) * 0.5
-    k = torch.randn(2, 513, 2, 128, device="cuda", dtype=torch.bfloat16) * 0.5
-    v = torch.randn_like(k)
-    seqused = torch.tensor([127, 513], device="cuda", dtype=torch.int32)
-    out, lse = allocate_outputs(q)
-    forward_seqused_static(q, k, v, seqused, out=out, softmax_lse=lse)
-    refs = []
-    for batch, used in enumerate((127, 513)):
-        refs.append(_reference(q[batch : batch + 1], k[batch : batch + 1, :used], v[batch : batch + 1, :used]))
-    _assert_close(out, torch.cat(refs, dim=0))
-
-
-@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
-@pytest.mark.parametrize("head_dim", SPLIT_HEAD_DIMS)
-def test_split_kv_noncausal(dtype, head_dim):
-    q = torch.randn(1, 1, 8, head_dim, device="cuda", dtype=dtype) * 0.5
-    k = torch.randn(1, 4096, 2, head_dim, device="cuda", dtype=dtype) * 0.5
-    v = torch.randn_like(k)
-    workspace = allocate_workspace(q, k)
-    assert workspace is not None and workspace.num_splits > 1
-    assert workspace.out_accum.shape[-1] == (head_dim + 31) & ~31
-    out, lse = allocate_outputs(q)
-    forward_static(q, k, v, out=out, softmax_lse=lse, workspace=workspace)
+    ops.forward_static(q, k, v, out=out, softmax_lse=lse)
     _assert_close(out, _reference(q, k, v))
+    print("PASS aligned padded strides")
 
 
-@pytest.mark.parametrize("head_dim", SPLIT_HEAD_DIMS)
-def test_split_kv_seqused(head_dim):
-    q = torch.randn(1, 1, 8, head_dim, device="cuda", dtype=torch.bfloat16) * 0.5
-    k = torch.randn(1, 4096, 2, head_dim, device="cuda", dtype=torch.bfloat16) * 0.5
-    v = torch.randn_like(k)
-    workspace = allocate_workspace(q, k)
-    assert workspace is not None and workspace.num_splits > 1
-    out, lse = allocate_outputs(q)
-    used = torch.tensor([3073], device="cuda", dtype=torch.int32)
-    forward_seqused_static(
-        q, k, v, used, out=out, softmax_lse=lse, workspace=workspace
-    )
-    _assert_close(out, _reference(q, k[:, :3073], v[:, :3073]))
-
-
-@pytest.mark.parametrize(
-    "head_dim",
-    tuple(dim for dim in SUPPORTED_HEAD_DIMS if dim not in SPLIT_HEAD_DIMS),
-)
-def test_unsafe_split_dimensions_fall_back_to_no_split(head_dim):
-    q = torch.randn(1, 1, 8, head_dim, device="cuda", dtype=torch.bfloat16) * 0.5
-    k = torch.randn(1, 1024, 2, head_dim, device="cuda", dtype=torch.bfloat16) * 0.5
-    v = torch.randn_like(k)
-    assert allocate_workspace(q, k) is None
-    _assert_close(forward(q, k, v, use_split_kv=True), _reference(q, k, v))
-
-
-def test_rejects_manual_unsafe_split_workspace():
-    head_dim = 32
-    q = torch.randn(1, 1, 8, head_dim, device="cuda", dtype=torch.bfloat16)
-    k = torch.randn(1, 1024, 2, head_dim, device="cuda", dtype=torch.bfloat16)
-    v = torch.randn_like(k)
-    out, lse = allocate_outputs(q)
-    workspace = FA2Workspace(
-        torch.empty((4, 1, 8, 1), device="cuda", dtype=torch.float32),
-        torch.empty((4, 1, 8, 1, 32), device="cuda", dtype=torch.float32),
-        128,
-        4,
-    )
-    with pytest.raises(RuntimeError, match="split-KV does not support"):
-        forward_static(q, k, v, out=out, softmax_lse=lse, workspace=workspace)
-
-
-@pytest.mark.parametrize("head_dim", [128, 256])
-def test_split_kv_causal(head_dim):
-    q = torch.randn(1, 257, 8, head_dim, device="cuda", dtype=torch.bfloat16) * 0.5
-    k = torch.randn(1, 4096, 2, head_dim, device="cuda", dtype=torch.bfloat16) * 0.5
-    v = torch.randn_like(k)
-    workspace = allocate_workspace(q, k)
-    assert workspace is not None and workspace.num_splits > 1
-    out, lse = allocate_outputs(q)
-    forward_static(q, k, v, out=out, softmax_lse=lse, workspace=workspace, causal=True)
-    _assert_close(out, _reference(q, k, v, causal=True))
-
-
-def test_cuda_graph_static_and_device_length_update():
-    q = torch.randn(1, 8, 8, 128, device="cuda", dtype=torch.bfloat16) * 0.5
-    k = torch.randn(1, 256, 2, 128, device="cuda", dtype=torch.bfloat16) * 0.5
-    v = torch.randn_like(k)
-    used = torch.tensor([128], device="cuda", dtype=torch.int32)
-    out, lse = allocate_outputs(q)
-    for _ in range(3):
-        forward_seqused_static(q, k, v, used, out=out, softmax_lse=lse)
-    torch.cuda.synchronize()
-    graph = torch.cuda.CUDAGraph()
-    with torch.cuda.graph(graph):
-        forward_seqused_static(q, k, v, used, out=out, softmax_lse=lse)
-    graph.replay()
-    _assert_close(out, _reference(q, k[:, :128], v[:, :128]))
-    used.fill_(255)
-    graph.replay()
-    _assert_close(out, _reference(q, k[:, :255], v[:, :255]))
-
-
-def test_torch_compile_trace():
-    q = torch.randn(1, 16, 8, 128, device="cuda", dtype=torch.bfloat16)
-    k = torch.randn(1, 64, 2, 128, device="cuda", dtype=torch.bfloat16)
-    v = torch.randn_like(k)
-    out, lse = allocate_outputs(q)
-
-    def call(q_, k_, v_, out_, lse_):
-        return forward_static(q_, k_, v_, out=out_, softmax_lse=lse_)
-
-    compiled = torch.compile(call, fullgraph=True)
-    actual = compiled(q, k, v, out, lse)
-    _assert_close(actual, _reference(q, k, v))
-
-
-@pytest.mark.parametrize("bad_dim", [7, 44, 260])
-def test_rejects_unbuilt_head_dim(bad_dim):
+def _check_rejections(ops):
+    bad_dim = 44
     q = torch.randn(1, 4, 4, bad_dim, device="cuda", dtype=torch.bfloat16)
     k = torch.randn(1, 4, 4, bad_dim, device="cuda", dtype=torch.bfloat16)
     v = torch.randn_like(k)
     out = torch.empty_like(q)
     lse = torch.empty((1, 4, 4), device="cuda", dtype=torch.float32)
-    with pytest.raises(RuntimeError, match="head_dim"):
-        forward_static(q, k, v, out=out, softmax_lse=lse)
-
-
-def test_rejects_fp16_causal():
-    q = torch.randn(1, 4, 4, 128, device="cuda", dtype=torch.float16)
-    k = torch.randn(1, 4, 4, 128, device="cuda", dtype=torch.float16)
-    v = torch.randn_like(k)
-    out, lse = allocate_outputs(q)
-    with pytest.raises(RuntimeError, match="causal v1 supports bf16"):
-        forward_static(q, k, v, out=out, softmax_lse=lse, causal=True)
+    try:
+        ops.forward_static(q, k, v, out=out, softmax_lse=lse)
+    except RuntimeError as exc:
+        if "head_dim" not in str(exc):
+            raise
+    else:
+        raise AssertionError("unbuilt head_dim must be rejected")
+
+    q16 = torch.randn(1, 4, 4, 128, device="cuda", dtype=torch.float16)
+    k16 = torch.randn(1, 4, 4, 128, device="cuda", dtype=torch.float16)
+    v16 = torch.randn_like(k16)
+    out16, lse16 = ops.allocate_outputs(q16)
+    try:
+        ops.forward_static(q16, k16, v16, out=out16, softmax_lse=lse16, causal=True)
+    except RuntimeError as exc:
+        if "causal v1 supports bf16" not in str(exc):
+            raise
+    else:
+        raise AssertionError("fp16 causal must be rejected")
 
-def test_rejects_misaligned_head_stride():
     storage = torch.randn(1, 8, 4, 129, device="cuda", dtype=torch.bfloat16)
-    q = storage[..., :128]
-    k = storage[..., :128]
-    v = storage[..., :128]
-    out = torch.empty_strided(q.shape, q.stride(), device=q.device, dtype=q.dtype)
-    lse = torch.empty((1, 4, 8), device="cuda", dtype=torch.float32)
-    with pytest.raises(RuntimeError, match="preserve 16-byte alignment"):
-        forward_static(q, k, v, out=out, softmax_lse=lse)
+    qs = storage[..., :128]
+    ks = storage[..., :128]
+    vs = storage[..., :128]
+    outs = torch.empty_strided(qs.shape, qs.stride(), device=qs.device, dtype=qs.dtype)
+    lses = torch.empty((1, 4, 8), device="cuda", dtype=torch.float32)
+    try:
+        ops.forward_static(qs, ks, vs, out=outs, softmax_lse=lses)
+    except RuntimeError as exc:
+        if "16-byte alignment" not in str(exc):
+            raise
+    else:
+        raise AssertionError("misaligned head stride must be rejected")
+    print("PASS rejections (unbuilt head_dim, fp16 causal, misaligned stride)")
+
+
+def run(args) -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    if args.mode == "full":
+        for dtype in (torch.float16, torch.bfloat16):
+            for head_dim in (48, 64, 128):
+                for shape in [(1, 17, 63, 16, 4), (2, 64, 129, 12, 4)]:
+                    _check_noncausal(ops, dtype, head_dim, shape)
+        for head_dim in (128, 256):
+            for seqlen in (1, 17, 257):
+                _check_causal(ops, head_dim, seqlen)
+    else:
+        _check_noncausal(ops, torch.bfloat16, 128, (1, 17, 63, 16, 4))
+        _check_noncausal(ops, torch.float16, 64, (1, 17, 63, 16, 4))
+        _check_causal(ops, 128, 17)
+    _check_seqused(ops)
+    _check_padded_strides(ops)
+    _check_rejections(ops)
+    if args.backend == "installed":
+        if args.mode == "full":
+            for head_dim in SPLIT_HEAD_DIMS:
+                if head_dim not in (48, 64, 96, 128):
+                    continue
+                q = torch.randn(1, 1, 8, head_dim, device="cuda", dtype=torch.bfloat16) * 0.5
+                k = torch.randn(1, 4096, 2, head_dim, device="cuda", dtype=torch.bfloat16) * 0.5
+                v = torch.randn_like(k)
+                workspace = ops.allocate_workspace(q, k)
+                out, lse = ops.allocate_outputs(q)
+                ops.forward_static(q, k, v, out=out, softmax_lse=lse, workspace=workspace)
+                _assert_close(out, _reference(q, k, v))
+            print("PASS split-KV workspace (installed)")
+    print(f"PASS fa2-seqused-runtime {args.backend} mode={args.mode}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    try:
+        run(args)
+    except Exception:
+        import traceback
+        traceback.print_exc()
+        return 1
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(
+            json.dumps({"passed": 1, "total": 1, "backend": args.backend}) + "\n"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-adaptive-norms/benchmarks/benchmark.py b/flashrt-adaptive-norms/benchmarks/benchmark.py
index a27e5a1..4bf9410 100644
--- a/flashrt-adaptive-norms/benchmarks/benchmark.py
+++ b/flashrt-adaptive-norms/benchmarks/benchmark.py
@@ -17,6 +17,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 PACKAGE = ROOT / "flashrt-adaptive-norms"
 REGISTRATION_INCLUDE = (
@@ -251,7 +261,9 @@ def main():
     )
     parser.add_argument("--output", default=None)
     parser.add_argument("--markdown", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
     torch.manual_seed(53)
diff --git a/flashrt-adarms-train/SYNC.md b/flashrt-adarms-train/SYNC.md
new file mode 100644
index 0000000..82371f6
--- /dev/null
+++ b/flashrt-adarms-train/SYNC.md
@@ -0,0 +1,31 @@
+# Source Sync
+
+- Package: `flashrt-adarms-train` (AdaLRS/AdARMS optimizer training step kernels).
+- Upstream FlashRT source: `../official/FlashRT`
+- Upstream revision: pending confirmation; packaged source is maintained in the
+  flashrt-project FlashRT-HF-kernels repository
+  (https://github.com/flashrt-project/FlashRT-HF-kernels).
+
+Copied source files:
+
+- `csrc/README.md`
+- `csrc/adarms_train.cu`
+- `csrc/adarms_train.cuh`
+
+Local packaging edits:
+
+- Added Tensor-facing PyTorch custom ops in `torch-ext/torch_binding.cpp`.
+- Added Python wrappers and fake registrations in `torch-ext/flashrt_adarms_train`.
+- Kept public APIs Tensor-facing; no raw pointer or stream arguments.
+- Includes rewritten to be package-local; serving-runtime dependencies removed.
+- CUDA launchers kept graph-safe: no dynamic allocation inside hot kernels.
+
+Architecture assumptions:
+
+- CUDA 12.8+ / 13.0+ (CUDA 13.2 validated on NVIDIA Thor, sm_110a).
+- NVIDIA Blackwell-family targets; Thor sm_110a validated on real hardware.
+
+Runtime constraints:
+
+- Inputs and outputs are `torch.Tensor`; shapes and dtypes are validated in the binding.
+- Benchmarks cap CUDA memory at 30 GB per process via `set_per_process_memory_fraction`.
diff --git a/flashrt-adarms-train/benchmarks/benchmark.py b/flashrt-adarms-train/benchmarks/benchmark.py
new file mode 100644
index 0000000..200727f
--- /dev/null
+++ b/flashrt-adarms-train/benchmarks/benchmark.py
@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+"""AdaRMS / gated-residual AdaRMS training benchmark (kernel vs torch.compile)."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from pathlib import Path
+
+import torch
+
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
+def elapsed_us(fn, warmup: int = 20, repeats: int = 100) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(repeats):
+        fn()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) * 1000.0 / repeats
+
+
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        root = Path(__file__).resolve().parents[1]
+        sys.path.insert(0, str(root / "torch-ext"))
+        return importlib.import_module("flashrt_adarms_train")
+    if artifact:
+        sys.path.insert(0, artifact)
+    return importlib.import_module("flashrt_adarms_train")
+
+
+def run_case(ops, label: str, b: int, t: int, h: int) -> dict:
+    torch.manual_seed(0)
+    x = torch.randn(b, t, h, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    mod = torch.randn(b, 1, 3 * h, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    xb = torch.randn(b, t, h, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+    hb = torch.randn_like(xb, requires_grad=True)
+    gb = torch.randn_like(xb, requires_grad=True)
+    mb = torch.randn(b, 3 * h, device="cuda", dtype=torch.bfloat16, requires_grad=True)
+
+    def fwd():
+        y, _g = ops.adarms(x, mod)
+        return y
+
+    def fwd_bwd():
+        y, _g = ops.adarms(x, mod)
+        y.float().sum().backward()
+        return y
+
+    def resgate():
+        return ops.resgate_adarms(xb, hb, gb, mb)[1]
+
+    def resgate_bwd():
+        y = ops.resgate_adarms(xb, hb, gb, mb)[1]
+        y.float().sum().backward()
+        return y
+
+    fwd_us = elapsed_us(fwd)
+    fwd_bwd_us = elapsed_us(fwd_bwd)
+    resgate_us = elapsed_us(resgate)
+    resgate_bwd_us = elapsed_us(resgate_bwd)
+
+    comp_fwd = torch.compile(fwd, mode="reduce-overhead")
+    comp_fwd()
+    compile_us = elapsed_us(lambda: comp_fwd())
+
+    comp_rg = torch.compile(resgate, mode="reduce-overhead")
+    comp_rg()
+    compile_resgate_us = elapsed_us(lambda: comp_rg())
+
+    return {
+        "label": label, "B": b, "T": t, "H": h,
+        "adarms_fwd_us": fwd_us, "adarms_fwd_bwd_us": fwd_bwd_us,
+        "resgate_fwd_us": resgate_us, "resgate_fwd_bwd_us": resgate_bwd_us,
+        "compile_fwd_us": compile_us, "compile_resgate_us": compile_resgate_us,
+        "vs_compile_fwd": compile_us / fwd_us,
+        "vs_compile_resgate": compile_resgate_us / resgate_us,
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="installed")
+    parser.add_argument("--artifact")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
+    args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
+    ops = load_ops(args.backend, args.artifact)
+    print("label,B,T,H,adarms_fwd_us,adarms_fwd_bwd_us,resgate_fwd_us,resgate_fwd_bwd_us,compile_fwd_us,compile_resgate_us,vs_compile_fwd,vs_compile_resgate")
+    for label, b, t, h in [
+        ("vla_1k", 1, 1024, 4096),
+        ("vla_2k", 1, 2048, 4096),
+        ("vla_4k_b4", 4, 1024, 4096),
+    ]:
+        r = run_case(ops, label, b, t, h)
+        print(
+            f"{r['label']},{r['B']},{r['T']},{r['H']},{r['adarms_fwd_us']:.3f},"
+            f"{r['adarms_fwd_bwd_us']:.3f},{r['resgate_fwd_us']:.3f},"
+            f"{r['resgate_fwd_bwd_us']:.3f},{r['compile_fwd_us']:.3f},"
+            f"{r['compile_resgate_us']:.3f},"
+            f"{r['vs_compile_fwd']:.2f}x,{r['vs_compile_resgate']:.2f}x"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-flex-attention-train/benchmarks/benchmark.py b/flashrt-flex-attention-train/benchmarks/benchmark.py
index 750ba3c..87aa986 100644
--- a/flashrt-flex-attention-train/benchmarks/benchmark.py
+++ b/flashrt-flex-attention-train/benchmarks/benchmark.py
@@ -11,6 +11,16 @@
 import torch.nn.functional as F
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[1]
 sys.path.insert(0, str(ROOT / "torch-ext"))
 import flashrt_flex_attention_train as flex_ops  # noqa: E402
@@ -365,11 +375,13 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--output")
     parser.add_argument("--require-gates", action="store_true")
     parser.add_argument("--no-prefix-mask", action="store_true")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     return parser.parse_args()
 
 
 def main() -> None:
     args = parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     if args.device.startswith("cuda") and not torch.cuda.is_available():
         raise SystemExit("CUDA requested but not available")
     backends = (
diff --git a/flashrt-fp8-ffn/benchmarks/benchmark.py b/flashrt-fp8-ffn/benchmarks/benchmark.py
index aab1d5f..3ee83dc 100644
--- a/flashrt-fp8-ffn/benchmarks/benchmark.py
+++ b/flashrt-fp8-ffn/benchmarks/benchmark.py
@@ -17,6 +17,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 PACKAGE = ROOT / "flashrt-fp8-ffn"
 REGISTRATION_INCLUDE = (
@@ -522,7 +532,9 @@ def main() -> None:
     parser.add_argument("--output", type=Path, default=None)
     parser.add_argument("--markdown", type=Path, default=None)
     parser.add_argument("--list-shapes", action="store_true")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if args.list_shapes:
         print("Shape groups:")
diff --git a/flashrt-fp8-swiglu-ffn/benchmarks/benchmark.py b/flashrt-fp8-swiglu-ffn/benchmarks/benchmark.py
index b038df2..9bf38a8 100644
--- a/flashrt-fp8-swiglu-ffn/benchmarks/benchmark.py
+++ b/flashrt-fp8-swiglu-ffn/benchmarks/benchmark.py
@@ -17,6 +17,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 PACKAGE = ROOT / "flashrt-fp8-swiglu-ffn"
 REGISTRATION_INCLUDE = (
@@ -512,7 +522,9 @@ def main() -> None:
     parser.add_argument("--p99-rel-limit", type=float, default=0.05)
     parser.add_argument("--output", default=None)
     parser.add_argument("--markdown", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
diff --git a/flashrt-fused-quant/VALIDATION.md b/flashrt-fused-quant/VALIDATION.md
index 627fafa..002f3b8 100644
--- a/flashrt-fused-quant/VALIDATION.md
+++ b/flashrt-fused-quant/VALIDATION.md
@@ -76,8 +76,7 @@ over the full grid.
 ## Known Gaps
 
 - `build.toml`, `flake.nix`, and `flake.lock` are present.
-- `/home/heima/suliang/PI/.hf-kernel-env/bin/kernel-builder-docker
-  check-config .` passed for this package.
+- `kernel-builder-docker check-config .` passed for this package.
 - `kernel-builder build --variant torch211-cxx11-cu128-x86_64-linux` passed
   for this package, and the copied artifact passed package tests, examples,
   installed accuracy sweep, and the local release-candidate benchmark runner.
diff --git a/flashrt-fused-quant/tests/test_silu_mul_quant_nvfp4.py b/flashrt-fused-quant/tests/test_silu_mul_quant_nvfp4.py
index 568c30e..a8c1a63 100644
--- a/flashrt-fused-quant/tests/test_silu_mul_quant_nvfp4.py
+++ b/flashrt-fused-quant/tests/test_silu_mul_quant_nvfp4.py
@@ -1,16 +1,124 @@
+#!/usr/bin/env python3
+"""Correctness tests for flashrt-fused-quant (SiLU*up fused into NVFP4)."""
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
 import math
+import os
 import struct
+import sys
+from pathlib import Path
 
-import pytest
 import torch
 
-from flashrt_fused_quant import (
-    nvfp4_swizzled_scale_bytes,
-    silu_mul_merged_quant_nvfp4_swizzled_bf16,
-    silu_mul_quant_nvfp4_swizzled_bf16,
+
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "flashrt-fused-quant"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
 )
 
 
+def _arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if major >= 12:
+        return "12.0a"
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    return f"{major}.{minor}"
+
+
+def _swizzled_bytes(rows: int, cols: int) -> int:
+    if rows <= 0:
+        raise ValueError("rows must be positive")
+    if cols <= 0 or cols % 16 != 0:
+        raise ValueError("cols must be positive and divisible by 16")
+    n_blocks = cols // 16
+    n_row_super = (rows + 127) // 128
+    n_col_super = (n_blocks + 3) // 4
+    return n_row_super * n_col_super * 512
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+
+    def scale_bytes(self, rows: int, cols: int) -> int:
+        return _swizzled_bytes(rows, cols)
+
+    def silu_mul_quant(self, gate, up, *, packed=None, scales=None):
+        rows, cols = gate.shape
+        if packed is None:
+            packed = torch.empty((rows, cols // 2), device=gate.device, dtype=torch.uint8)
+        if scales is None:
+            scales = torch.zeros(_swizzled_bytes(rows, cols), device=gate.device, dtype=torch.uint8)
+        self._ops.silu_mul_quant_nvfp4_swizzled_bf16(gate, up, packed, scales)
+        return packed, scales
+
+    def silu_mul_merged_quant(self, merged, *, packed=None, scales=None):
+        rows, merged_cols = merged.shape
+        cols = merged_cols // 2
+        if packed is None:
+            packed = torch.empty((rows, cols // 2), device=merged.device, dtype=torch.uint8)
+        if scales is None:
+            scales = torch.zeros(_swizzled_bytes(rows, cols), device=merged.device, dtype=torch.uint8)
+        self._ops.silu_mul_merged_quant_nvfp4_swizzled_bf16(merged, packed, scales)
+        return packed, scales
+
+
+class InstalledOps:
+    def __init__(self, module) -> None:
+        self._module = module
+
+    def scale_bytes(self, rows: int, cols: int) -> int:
+        return int(self._module.nvfp4_swizzled_scale_bytes(rows, cols))
+
+    def silu_mul_quant(self, gate, up, *, packed=None, scales=None):
+        return self._module.silu_mul_quant_nvfp4_swizzled_bf16(gate, up, packed=packed, scales=scales)
+
+    def silu_mul_merged_quant(self, merged, *, packed=None, scales=None):
+        return self._module.silu_mul_merged_quant_nvfp4_swizzled_bf16(merged, packed=packed, scales=scales)
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
+    namespace = "flashrt_fused_quant_test"
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            str(PACKAGE / "csrc" / "silu_mul_to_nvfp4_swizzled.cu"),
+        ],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+
+
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return InstalledOps(importlib.import_module("flashrt_fused_quant"))
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
 def _float_to_fp4_e2m1(v: float) -> int:
     sign = 0x8 if v < 0 else 0x0
     a = abs(v)
@@ -70,6 +178,22 @@ def _ue4m3_to_float(byte: int) -> float:
     return math.ldexp(1.0 + m / 8.0, e - 7)
 
 
+def _swizzle(scales: torch.Tensor) -> torch.Tensor:
+    rows, n_blocks = scales.shape
+    n_col_super = (n_blocks + 3) // 4
+    out = torch.zeros(_swizzled_bytes(rows, n_blocks * 16), dtype=torch.uint8)
+    for row in range(rows):
+        rb = row // 128
+        ri = row % 128
+        for block in range(n_blocks):
+            cb = block // 4
+            ci = block % 4
+            super_idx = rb * n_col_super + cb
+            inner_off = (ri % 32) * 16 + (ri // 32) * 4 + ci
+            out[super_idx * 512 + inner_off] = scales[row, block]
+    return out
+
+
 def _reference(gate: torch.Tensor, up: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
     rows, cols = gate.shape
     gate_cpu = gate.cpu()
@@ -82,7 +206,6 @@ def _reference(gate: torch.Tensor, up: torch.Tensor) -> tuple[torch.Tensor, torc
             silu = g / (1.0 + math.exp(-g))
             silu_bf = float(torch.tensor(silu, dtype=torch.bfloat16))
             vals[row, col] = torch.tensor(silu_bf * u, dtype=torch.bfloat16)
-
     packed = torch.empty((rows, cols // 2), dtype=torch.uint8)
     scale_linear = torch.empty((rows, cols // 16), dtype=torch.uint8)
     for row in range(rows):
@@ -98,61 +221,78 @@ def _reference(gate: torch.Tensor, up: torch.Tensor) -> tuple[torch.Tensor, torc
                 lo = _float_to_fp4_e2m1(float(vals[row, i]) * inv_scale)
                 hi = _float_to_fp4_e2m1(float(vals[row, i + 1]) * inv_scale)
                 packed[row, i // 2] = (hi << 4) | lo
+    return packed, _swizzle(scale_linear)
 
-    scales = _swizzle(scale_linear)
-    return packed, scales
 
-
-def _swizzle(scales: torch.Tensor) -> torch.Tensor:
-    rows, n_blocks = scales.shape
-    n_col_super = (n_blocks + 3) // 4
-    out = torch.zeros(
-        (nvfp4_swizzled_scale_bytes(rows, n_blocks * 16),),
-        dtype=torch.uint8,
-    )
-    for row in range(rows):
-        rb = row // 128
-        ri = row % 128
-        for block in range(n_blocks):
-            cb = block // 4
-            ci = block % 4
-            super_idx = rb * n_col_super + cb
-            inner_off = (ri % 32) * 16 + (ri // 32) * 4 + ci
-            out[super_idx * 512 + inner_off] = scales[row, block]
-    return out
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-@pytest.mark.parametrize(("rows", "cols"), [(1, 16), (3, 64), (33, 128)])
-def test_silu_mul_quant_nvfp4_swizzled_bf16(rows, cols):
+def _check_silu_mul_quant(ops, rows: int, cols: int) -> None:
     torch.manual_seed(0)
     gate = (torch.randn((rows, cols), device="cuda", dtype=torch.bfloat16) * 0.5).contiguous()
     up = (torch.randn((rows, cols), device="cuda", dtype=torch.bfloat16) * 0.5).contiguous()
+    packed, scales = ops.silu_mul_quant(gate, up)
+    exp_packed, exp_scales = _reference(gate, up)
+    torch.testing.assert_close(packed.cpu(), exp_packed)
+    torch.testing.assert_close(scales.cpu(), exp_scales)
+    print(f"PASS silu_mul_quant rows={rows} cols={cols}: exact")
 
-    packed, scales = silu_mul_quant_nvfp4_swizzled_bf16(gate, up)
-    expected_packed, expected_scales = _reference(gate, up)
-
-    torch.testing.assert_close(packed.cpu(), expected_packed)
-    torch.testing.assert_close(scales.cpu(), expected_scales)
 
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_silu_mul_merged_quant_nvfp4_swizzled_bf16():
+def _check_merged_quant(ops) -> None:
     torch.manual_seed(1)
     rows, cols = 4, 64
     gate = (torch.randn((rows, cols), device="cuda", dtype=torch.bfloat16) * 0.5).contiguous()
     up = (torch.randn((rows, cols), device="cuda", dtype=torch.bfloat16) * 0.5).contiguous()
     merged = torch.cat([gate, up], dim=1).contiguous()
+    packed, scales = ops.silu_mul_merged_quant(merged)
+    exp_packed, exp_scales = _reference(gate, up)
+    torch.testing.assert_close(packed.cpu(), exp_packed)
+    torch.testing.assert_close(scales.cpu(), exp_scales)
+    print("PASS silu_mul_merged_quant: exact")
+
+
+def _check_scale_bytes(ops) -> None:
+    for fn_args, match in [((0, 64), "rows"), ((1, 15), "cols")]:
+        try:
+            ops.scale_bytes(*fn_args)
+        except (ValueError, RuntimeError) as exc:
+            if match not in str(exc):
+                raise AssertionError(f"expected error containing {match!r}, got {exc}")
+        else:
+            raise AssertionError(f"expected error containing {match!r}")
+    print("PASS scale_bytes rejects invalid shape")
+
+
+def run(args) -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    cases = [(1, 16), (3, 64)] if args.mode == "smoke" else [(1, 16), (3, 64), (33, 128)]
+    for rows, cols in cases:
+        _check_silu_mul_quant(ops, rows, cols)
+    _check_merged_quant(ops)
+    _check_scale_bytes(ops)
+    print(f"PASS flashrt-fused-quant {args.backend} mode={args.mode}: "
+          f"{len(cases) + 2} checks")
 
-    packed, scales = silu_mul_merged_quant_nvfp4_swizzled_bf16(merged)
-    expected_packed, expected_scales = _reference(gate, up)
 
-    torch.testing.assert_close(packed.cpu(), expected_packed)
-    torch.testing.assert_close(scales.cpu(), expected_scales)
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    try:
+        run(args)
+    except Exception:
+        import traceback
+        traceback.print_exc()
+        return 1
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(
+            json.dumps({"passed": 1, "total": 1, "backend": args.backend}) + "\n"
+        )
+    return 0
 
 
-def test_nvfp4_swizzled_scale_bytes_rejects_invalid_shape():
-    with pytest.raises(ValueError, match="rows"):
-        nvfp4_swizzled_scale_bytes(0, 64)
-    with pytest.raises(ValueError, match="cols"):
-        nvfp4_swizzled_scale_bytes(1, 15)
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-gemm-epilogues/VALIDATION.md b/flashrt-gemm-epilogues/VALIDATION.md
index 10f7147..7b76166 100644
--- a/flashrt-gemm-epilogues/VALIDATION.md
+++ b/flashrt-gemm-epilogues/VALIDATION.md
@@ -34,11 +34,11 @@ Runtime smoke environment:
 From this package directory:
 
 ```bash
-/home/heima/suliang/PI/.hf-kernel-env/bin/kernel-builder-docker check-config .
-/home/heima/suliang/PI/.hf-kernel-env/bin/kernel-builder-docker build --variant torch211-cxx11-cu128-x86_64-linux --max-jobs 1 --cores 8 -L .
-/home/heima/suliang/PI/.hf-kernel-env/bin/kernel-builder-docker build --variant torch211-cxx11-cu126-x86_64-linux --max-jobs 1 --cores 8 -L .
-/home/heima/suliang/PI/.hf-kernel-env/bin/kernel-builder-docker build --variant torch211-cxx11-cu130-x86_64-linux --max-jobs 1 --cores 8 -L .
-/home/heima/suliang/PI/.hf-kernel-env/bin/kernel-builder-docker check-builds .
+kernel-builder-docker check-config .
+kernel-builder-docker build --variant torch211-cxx11-cu128-x86_64-linux --max-jobs 1 --cores 8 -L .
+kernel-builder-docker build --variant torch211-cxx11-cu126-x86_64-linux --max-jobs 1 --cores 8 -L .
+kernel-builder-docker build --variant torch211-cxx11-cu130-x86_64-linux --max-jobs 1 --cores 8 -L .
+kernel-builder-docker check-builds .
 ```
 
 Host-side correctness smoke from the repository root:
diff --git a/flashrt-gemm-epilogues/benchmarks/benchmark.py b/flashrt-gemm-epilogues/benchmarks/benchmark.py
index bea32d9..2ca92fe 100644
--- a/flashrt-gemm-epilogues/benchmarks/benchmark.py
+++ b/flashrt-gemm-epilogues/benchmarks/benchmark.py
@@ -3,6 +3,19 @@
 from kernels.benchmark import Benchmark
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
+_apply_mem_cap()
+
+
 _original_allclose = torch.allclose
 
 
diff --git a/flashrt-gemm-epilogues/tests/test_bf16_gemm_bias_gelu.py b/flashrt-gemm-epilogues/tests/test_bf16_gemm_bias_gelu.py
index 07df7fc..9b895d8 100644
--- a/flashrt-gemm-epilogues/tests/test_bf16_gemm_bias_gelu.py
+++ b/flashrt-gemm-epilogues/tests/test_bf16_gemm_bias_gelu.py
@@ -1,92 +1,197 @@
-import pytest
-import torch
-
-import flashrt_gemm_epilogues as flashrt_ops
-
-
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="CUDA is required",
-)
-
+#!/usr/bin/env python3
+"""Correctness tests for flashrt-gemm-epilogues BF16 GEMM bias/GELU ops."""
+from __future__ import annotations
 
-def _reference(a, b, bias):
-    y = a @ b
-    y = y + bias
-    return torch.nn.functional.gelu(y).to(torch.bfloat16)
+import argparse
+import importlib
+import json
+import os
+import sys
+from pathlib import Path
 
-
-def _bias_reference(a, b, bias):
-    return ((a @ b) + bias).to(torch.bfloat16)
+import torch
 
 
-@pytest.mark.parametrize(
-    ("m", "n", "k"),
-    [
-        (16, 64, 32),
-        (32, 128, 64),
-    ],
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "flashrt-gemm-epilogues"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
 )
-def test_bf16_gemm_bias(m, n, k):
-    torch.manual_seed(2)
-    device = torch.device("cuda")
-    a = torch.randn((m, k), device=device, dtype=torch.bfloat16).contiguous()
-    b = torch.randn((k, n), device=device, dtype=torch.bfloat16).contiguous()
-    bias = torch.randn((n,), device=device, dtype=torch.bfloat16).contiguous()
-
-    out = flashrt_ops.bf16_gemm_bias(a, b, bias)
-    expected = _bias_reference(a, b, bias)
 
+_SOURCE_LIST = [
+    "torch-ext/torch_binding.cpp",
+    "csrc/bf16_gemm_bias_gelu.cu",
+    "csrc/bias_gelu_quantize_fp8.cu",
+    "csrc/channel_scale_quantize_fp8.cu",
+]
+
+
+def _arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if major >= 12:
+        return "12.0a"
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    return f"{major}.{minor}"
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+
+    def gemm_bias(self, a, b, bias, *, out=None):
+        if out is None:
+            out = torch.empty((a.shape[0], b.shape[1]), device=a.device, dtype=torch.bfloat16)
+        self._ops.bf16_gemm_bias(a, b, bias, out)
+        return out
+
+    def gemm_bias_gelu(self, a, b, bias, *, out=None):
+        if out is None:
+            out = torch.empty((a.shape[0], b.shape[1]), device=a.device, dtype=torch.bfloat16)
+        self._ops.bf16_gemm_bias_gelu(a, b, bias, out)
+        return out
+
+
+class InstalledOps:
+    def __init__(self, module) -> None:
+        self._module = module
+
+    def gemm_bias(self, a, b, bias, *, out=None):
+        return self._module.bf16_gemm_bias(a, b, bias, out=out)
+
+    def gemm_bias_gelu(self, a, b, bias, *, out=None):
+        return self._module.bf16_gemm_bias_gelu(a, b, bias, out=out)
+
+
+def _preload_cublaslt() -> None:
+    import ctypes
+    import ctypes.util
+
+    for parent in Path(torch.__file__).resolve().parents:
+        candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12"
+        if candidate.exists():
+            ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL)
+            return
+    library = ctypes.util.find_library("cublasLt")
+    if library:
+        ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL)
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    _preload_cublaslt()
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
+    namespace = "flashrt_gemm_epilogues_test"
+    load(
+        name=namespace,
+        sources=[str(PACKAGE / s) for s in _SOURCE_LIST],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+
+
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return InstalledOps(importlib.import_module("flashrt_gemm_epilogues"))
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def _check_gemm_bias(ops, m: int, n: int, k: int) -> None:
+    torch.manual_seed(2)
+    a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16).contiguous()
+    b = torch.randn((k, n), device="cuda", dtype=torch.bfloat16).contiguous()
+    bias = torch.randn((n,), device="cuda", dtype=torch.bfloat16).contiguous()
+    out = ops.gemm_bias(a, b, bias)
+    expected = ((a @ b) + bias).to(torch.bfloat16)
     torch.testing.assert_close(out.float(), expected.float(), rtol=3e-2, atol=1.25e-1)
+    print(f"PASS bf16_gemm_bias ({m},{n},{k})")
 
 
-@pytest.mark.parametrize(
-    ("m", "n", "k"),
-    [
-        (16, 64, 32),
-        (32, 128, 64),
-    ],
-)
-def test_bf16_gemm_bias_gelu(m, n, k):
+def _check_gemm_bias_gelu(ops, m: int, n: int, k: int) -> None:
     torch.manual_seed(3)
-    device = torch.device("cuda")
-    a = torch.randn((m, k), device=device, dtype=torch.bfloat16).contiguous()
-    b = torch.randn((k, n), device=device, dtype=torch.bfloat16).contiguous()
-    bias = torch.randn((n,), device=device, dtype=torch.bfloat16).contiguous()
-
-    out = flashrt_ops.bf16_gemm_bias_gelu(a, b, bias)
-    expected = _reference(a, b, bias)
-
+    a = torch.randn((m, k), device="cuda", dtype=torch.bfloat16).contiguous()
+    b = torch.randn((k, n), device="cuda", dtype=torch.bfloat16).contiguous()
+    bias = torch.randn((n,), device="cuda", dtype=torch.bfloat16).contiguous()
+    out = ops.gemm_bias_gelu(a, b, bias)
+    expected = torch.nn.functional.gelu(a @ b + bias).to(torch.bfloat16)
     torch.testing.assert_close(out.float(), expected.float(), rtol=3e-2, atol=1.25e-1)
-
-
-def test_bf16_gemm_bias_gelu_out_tensor_is_reused():
-    device = torch.device("cuda")
-    a = torch.randn((16, 32), device=device, dtype=torch.bfloat16).contiguous()
-    b = torch.randn((32, 64), device=device, dtype=torch.bfloat16).contiguous()
-    bias = torch.randn((64,), device=device, dtype=torch.bfloat16).contiguous()
-    out = torch.empty((16, 64), device=device, dtype=torch.bfloat16)
-
-    returned = flashrt_ops.bf16_gemm_bias_gelu(a, b, bias, out=out)
-
-    assert returned is out
-
-
-def test_bf16_gemm_bias_gelu_rejects_wrong_b_shape():
-    device = torch.device("cuda")
-    a = torch.randn((16, 32), device=device, dtype=torch.bfloat16).contiguous()
-    b = torch.randn((31, 64), device=device, dtype=torch.bfloat16).contiguous()
-    bias = torch.randn((64,), device=device, dtype=torch.bfloat16).contiguous()
-
-    with pytest.raises(RuntimeError, match="a.shape\\[1\\]"):
-        flashrt_ops.bf16_gemm_bias_gelu(a, b, bias)
-
-
-def test_bf16_gemm_bias_gelu_rejects_wrong_bias_shape():
-    device = torch.device("cuda")
-    a = torch.randn((16, 32), device=device, dtype=torch.bfloat16).contiguous()
-    b = torch.randn((32, 64), device=device, dtype=torch.bfloat16).contiguous()
-    bias = torch.randn((63,), device=device, dtype=torch.bfloat16).contiguous()
-
-    with pytest.raises(RuntimeError, match="bias length"):
-        flashrt_ops.bf16_gemm_bias_gelu(a, b, bias)
+    print(f"PASS bf16_gemm_bias_gelu ({m},{n},{k})")
+
+
+def _check_reuse_and_reject(ops) -> None:
+    a = torch.randn((16, 32), device="cuda", dtype=torch.bfloat16).contiguous()
+    b = torch.randn((32, 64), device="cuda", dtype=torch.bfloat16).contiguous()
+    bias = torch.randn((64,), device="cuda", dtype=torch.bfloat16).contiguous()
+    out = torch.empty((16, 64), device="cuda", dtype=torch.bfloat16)
+    returned = ops.gemm_bias_gelu(a, b, bias, out=out)
+    if returned is not out:
+        raise AssertionError("out tensor must be reused in place")
+    try:
+        ops.gemm_bias_gelu(a, torch.randn((31, 64), device="cuda", dtype=torch.bfloat16).contiguous(), bias)
+    except RuntimeError as exc:
+        if "a.shape[1]" not in str(exc):
+            raise
+    else:
+        raise AssertionError("wrong b shape must be rejected")
+    try:
+        ops.gemm_bias_gelu(a, b, torch.randn((63,), device="cuda", dtype=torch.bfloat16).contiguous())
+    except RuntimeError as exc:
+        if "bias length" not in str(exc):
+            raise
+    else:
+        raise AssertionError("wrong bias shape must be rejected")
+    print("PASS bf16_gemm_bias_gelu reuse + rejections")
+
+
+def run(args) -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    shapes = [(16, 64, 32)] if args.mode == "smoke" else [(16, 64, 32), (32, 128, 64)]
+    for m, n, k in shapes:
+        _check_gemm_bias(ops, m, n, k)
+        _check_gemm_bias_gelu(ops, m, n, k)
+    _check_reuse_and_reject(ops)
+    print(f"PASS flashrt-gemm-epilogues/bf16-gemm {args.backend} mode={args.mode}: "
+          f"{2 * len(shapes) + 1} checks")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    try:
+        run(args)
+    except Exception:
+        import traceback
+        traceback.print_exc()
+        return 1
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(
+            json.dumps({"passed": 1, "total": 1, "backend": args.backend}) + "\n"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-gemm-epilogues/tests/test_bf16_linear.py b/flashrt-gemm-epilogues/tests/test_bf16_linear.py
index de724b8..d4e729d 100644
--- a/flashrt-gemm-epilogues/tests/test_bf16_linear.py
+++ b/flashrt-gemm-epilogues/tests/test_bf16_linear.py
@@ -1,36 +1,126 @@
+#!/usr/bin/env python3
+"""Correctness tests for flashrt-gemm-epilogues BF16 linear (bias) ops."""
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
 import math
+import os
+import sys
+from pathlib import Path
 
-import pytest
 import torch
 
-import flashrt_gemm_epilogues as flashrt_ops
-
 
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available(),
-    reason="CUDA is required",
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "flashrt-gemm-epilogues"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
 )
 
-
 LINEAR_SHAPES = [
-    ("pi05_action_in", 10, 32, 1024),
-    ("pi05_qkv", 10, 1024, 2560),
-    ("pi05_o_proj", 10, 2048, 1024),
-    ("pi05_action_out", 10, 1024, 32),
     ("decode_m1_1024", 1, 1024, 1024),
     ("decode_m1_qkv", 1, 1024, 2560),
     ("decode_m8_1024", 8, 1024, 1024),
-    ("decode_m8_qkv", 8, 1024, 2560),
-    ("decode_m10_1024", 10, 1024, 1024),
     ("decode_m10_qkv", 10, 1024, 2560),
-    ("decode_m16_1024", 16, 1024, 1024),
-    ("decode_m16_qkv", 16, 1024, 2560),
     ("vlm_m512_square", 512, 1152, 1152),
-    ("vlm_m512_wide", 512, 1152, 4304),
     ("vla_m1024_square", 1024, 2048, 2048),
-    ("vla_m1024_wide", 1024, 2048, 8192),
 ]
 
+_SOURCE_LIST = [
+    "torch-ext/torch_binding.cpp",
+    "csrc/bf16_gemm_bias_gelu.cu",
+    "csrc/bias_gelu_quantize_fp8.cu",
+    "csrc/channel_scale_quantize_fp8.cu",
+]
+
+
+def _arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if major >= 12:
+        return "12.0a"
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    return f"{major}.{minor}"
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+
+    def linear(self, x, w, *, out=None):
+        if out is None:
+            out = torch.empty((x.shape[0], w.shape[1]), device=x.device, dtype=torch.bfloat16)
+        self._ops.bf16_linear_bf16(x, w, out)
+        return out
+
+    def linear_bias(self, x, w, bias, *, out=None):
+        if out is None:
+            out = torch.empty((x.shape[0], w.shape[1]), device=x.device, dtype=torch.bfloat16)
+        self._ops.bf16_linear_bias_bf16(x, w, bias, out)
+        return out
+
+
+class InstalledOps:
+    def __init__(self, module) -> None:
+        self._module = module
+
+    def linear(self, x, w, *, out=None):
+        return self._module.bf16_linear_bf16(x, w, out=out)
+
+    def linear_bias(self, x, w, bias, *, out=None):
+        return self._module.bf16_linear_bias_bf16(x, w, bias, out=out)
+
+
+def _preload_cublaslt() -> None:
+    import ctypes
+    import ctypes.util
+
+    for parent in Path(torch.__file__).resolve().parents:
+        candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12"
+        if candidate.exists():
+            ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL)
+            return
+    library = ctypes.util.find_library("cublasLt")
+    if library:
+        ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL)
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    _preload_cublaslt()
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
+    namespace = "flashrt_gemm_epilogues_test"
+    load(
+        name=namespace,
+        sources=[str(PACKAGE / s) for s in _SOURCE_LIST],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+
+
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return InstalledOps(importlib.import_module("flashrt_gemm_epilogues"))
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
 
 def _percentile(x: torch.Tensor, q: float) -> torch.Tensor:
     flat = x.flatten()
@@ -38,107 +128,103 @@ def _percentile(x: torch.Tensor, q: float) -> torch.Tensor:
     return flat.kthvalue(k).values
 
 
-def _metrics(out: torch.Tensor, expected: torch.Tensor) -> dict[str, float | str]:
-    diff = (out.float() - expected.float()).abs()
-    return {
-        "max_abs": float(diff.max().item()),
-        "mean_abs": float(diff.mean().item()),
-        "p99_abs": float(_percentile(diff, 0.99).item()),
-        "cosine": float(
-            torch.nn.functional.cosine_similarity(
-                out.float().flatten(), expected.float().flatten(), dim=0
-            ).item()
-        ),
-        "dtype": str(out.dtype),
-    }
-
-
-def _print_result(
-    name: str,
-    op: str,
-    shape: tuple[int, int, int],
-    metrics: dict[str, float | str],
-    *,
-    p99_limit: float,
-    cosine_limit: float,
-) -> None:
-    passed = metrics["p99_abs"] <= p99_limit and metrics["cosine"] >= cosine_limit
-    print(
-        "bf16_linear_correctness "
-        f"name={name} op={op} shape={shape} dtype={metrics['dtype']} "
-        f"max_abs={metrics['max_abs']:.6f} mean_abs={metrics['mean_abs']:.6f} "
-        f"p99_abs={metrics['p99_abs']:.6f} cosine={metrics['cosine']:.8f} "
-        f"tolerance=p99<={p99_limit},cosine>={cosine_limit} "
-        f"verified={'PASS' if passed else 'FAIL'}"
-    )
-
-
-@pytest.mark.parametrize(("name", "m", "k", "n"), LINEAR_SHAPES)
-def test_bf16_linear_correctness(name, m, k, n):
+def _check_linear(ops, name: str, m: int, k: int, n: int) -> None:
     torch.manual_seed(11)
-    device = torch.device("cuda")
-    x = torch.randn((m, k), device=device, dtype=torch.bfloat16).contiguous()
-    w = torch.randn((k, n), device=device, dtype=torch.bfloat16).contiguous()
-    out = torch.empty((m, n), device=device, dtype=torch.bfloat16)
-
-    returned = flashrt_ops.bf16_linear_bf16(x, w, out=out)
+    x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16).contiguous()
+    w = torch.randn((k, n), device="cuda", dtype=torch.bfloat16).contiguous()
+    out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+    returned = ops.linear(x, w, out=out)
+    if returned is not out:
+        raise AssertionError("out tensor must be reused in place")
     expected = (x @ w).to(torch.bfloat16)
-
-    assert returned is out
-    metrics = _metrics(out, expected)
-    p99_limit = 0.5
-    cosine_limit = 0.999
-    _print_result(name, "linear", (m, k, n), metrics, p99_limit=p99_limit, cosine_limit=cosine_limit)
-    assert metrics["p99_abs"] <= p99_limit
-    assert metrics["cosine"] >= cosine_limit
+    diff = (out.float() - expected.float()).abs()
+    p99 = float(_percentile(diff, 0.99).item())
+    cosine = float(torch.nn.functional.cosine_similarity(out.float().flatten(), expected.float().flatten(), dim=0).item())
+    if not (p99 <= 0.5 and cosine >= 0.999):
+        raise AssertionError(f"{name} linear failed: p99={p99} cosine={cosine}")
+    print(f"PASS bf16_linear {name} ({m},{k},{n}): p99={p99:.6f} cosine={cosine:.8f}")
 
 
-@pytest.mark.parametrize(("name", "m", "k", "n"), LINEAR_SHAPES)
-def test_bf16_linear_bias_correctness(name, m, k, n):
+def _check_linear_bias(ops, name: str, m: int, k: int, n: int) -> None:
     torch.manual_seed(17)
-    device = torch.device("cuda")
-    x = torch.randn((m, k), device=device, dtype=torch.bfloat16).contiguous()
-    w = torch.randn((k, n), device=device, dtype=torch.bfloat16).contiguous()
-    bias = torch.randn((n,), device=device, dtype=torch.bfloat16).contiguous()
-    out = torch.empty((m, n), device=device, dtype=torch.bfloat16)
-
-    returned = flashrt_ops.bf16_linear_bias_bf16(x, w, bias, out=out)
+    x = torch.randn((m, k), device="cuda", dtype=torch.bfloat16).contiguous()
+    w = torch.randn((k, n), device="cuda", dtype=torch.bfloat16).contiguous()
+    bias = torch.randn((n,), device="cuda", dtype=torch.bfloat16).contiguous()
+    out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
+    returned = ops.linear_bias(x, w, bias, out=out)
+    if returned is not out:
+        raise AssertionError("out tensor must be reused in place")
     expected = torch.addmm(bias, x, w).to(torch.bfloat16)
-
-    assert returned is out
-    metrics = _metrics(out, expected)
-    p99_limit = 0.5
-    cosine_limit = 0.999
-    _print_result(
-        name, "linear_bias", (m, k, n), metrics, p99_limit=p99_limit, cosine_limit=cosine_limit
-    )
-    assert metrics["p99_abs"] <= p99_limit
-    assert metrics["cosine"] >= cosine_limit
-
-
-def test_bf16_linear_rejects_wrong_w_shape():
-    device = torch.device("cuda")
-    x = torch.randn((10, 32), device=device, dtype=torch.bfloat16).contiguous()
-    w = torch.randn((31, 1024), device=device, dtype=torch.bfloat16).contiguous()
-
-    with pytest.raises(RuntimeError, match="x.shape\\[1\\]"):
-        flashrt_ops.bf16_linear_bf16(x, w)
-
-
-def test_bf16_linear_bias_rejects_wrong_bias_shape():
-    device = torch.device("cuda")
-    x = torch.randn((10, 32), device=device, dtype=torch.bfloat16).contiguous()
-    w = torch.randn((32, 1024), device=device, dtype=torch.bfloat16).contiguous()
-    bias = torch.randn((1023,), device=device, dtype=torch.bfloat16).contiguous()
-
-    with pytest.raises(RuntimeError, match="bias length"):
-        flashrt_ops.bf16_linear_bias_bf16(x, w, bias)
-
-
-def test_bf16_linear_rejects_noncontiguous_input():
-    device = torch.device("cuda")
-    x = torch.randn((32, 10), device=device, dtype=torch.bfloat16).t()
-    w = torch.randn((32, 1024), device=device, dtype=torch.bfloat16).contiguous()
-
-    with pytest.raises(RuntimeError, match="x must be contiguous"):
-        flashrt_ops.bf16_linear_bf16(x, w)
+    diff = (out.float() - expected.float()).abs()
+    p99 = float(_percentile(diff, 0.99).item())
+    cosine = float(torch.nn.functional.cosine_similarity(out.float().flatten(), expected.float().flatten(), dim=0).item())
+    if not (p99 <= 0.5 and cosine >= 0.999):
+        raise AssertionError(f"{name} linear_bias failed: p99={p99} cosine={cosine}")
+    print(f"PASS bf16_linear_bias {name} ({m},{k},{n}): p99={p99:.6f} cosine={cosine:.8f}")
+
+
+def _check_rejections(ops) -> None:
+    x = torch.randn((10, 32), device="cuda", dtype=torch.bfloat16).contiguous()
+    w_bad = torch.randn((31, 1024), device="cuda", dtype=torch.bfloat16).contiguous()
+    try:
+        ops.linear(x, w_bad)
+    except RuntimeError as exc:
+        if "x.shape[1]" not in str(exc):
+            raise
+    else:
+        raise AssertionError("wrong w shape must be rejected")
+    bias_bad = torch.randn((1023,), device="cuda", dtype=torch.bfloat16).contiguous()
+    w = torch.randn((32, 1024), device="cuda", dtype=torch.bfloat16).contiguous()
+    try:
+        ops.linear_bias(x, w, bias_bad)
+    except RuntimeError as exc:
+        if "bias length" not in str(exc):
+            raise
+    else:
+        raise AssertionError("wrong bias shape must be rejected")
+    try:
+        ops.linear(torch.randn((32, 10), device="cuda", dtype=torch.bfloat16).t().contiguous() if False else torch.randn((32, 10), device="cuda", dtype=torch.bfloat16).t(), w)
+    except RuntimeError as exc:
+        if "contiguous" not in str(exc):
+            raise
+    else:
+        raise AssertionError("non-contiguous x must be rejected")
+    print("PASS bf16_linear rejections")
+
+
+def run(args) -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    shapes = LINEAR_SHAPES if args.mode == "full" else LINEAR_SHAPES[:3]
+    for name, m, k, n in shapes:
+        _check_linear(ops, name, m, k, n)
+        _check_linear_bias(ops, name, m, k, n)
+    _check_rejections(ops)
+    print(f"PASS flashrt-gemm-epilogues/bf16-linear {args.backend} mode={args.mode}: "
+          f"{2 * len(shapes) + 1} checks")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    try:
+        run(args)
+    except Exception:
+        import traceback
+        traceback.print_exc()
+        return 1
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(
+            json.dumps({"passed": 1, "total": 1, "backend": args.backend}) + "\n"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-gemm-epilogues/tests/test_bias_gelu_quantize_fp8.py b/flashrt-gemm-epilogues/tests/test_bias_gelu_quantize_fp8.py
index 5fd02a6..fc13559 100644
--- a/flashrt-gemm-epilogues/tests/test_bias_gelu_quantize_fp8.py
+++ b/flashrt-gemm-epilogues/tests/test_bias_gelu_quantize_fp8.py
@@ -1,15 +1,36 @@
-import pytest
-import torch
+#!/usr/bin/env python3
+"""Correctness tests for flashrt-gemm-epilogues FP8 quantization epilogues."""
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import os
+import sys
+from pathlib import Path
 
-import flashrt_gemm_epilogues as flashrt_ops
+import torch
 
 
-pytestmark = pytest.mark.skipif(
-    not torch.cuda.is_available()
-    or not (hasattr(torch, "float8_e4m3fn") or hasattr(torch, "float8_e4m3fnuz")),
-    reason="CUDA/ROCm with FP8 support is required",
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "flashrt-gemm-epilogues"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
 )
 
+_SOURCE_LIST = [
+    "torch-ext/torch_binding.cpp",
+    "csrc/bf16_gemm_bias_gelu.cu",
+    "csrc/bias_gelu_quantize_fp8.cu",
+    "csrc/channel_scale_quantize_fp8.cu",
+]
+
 
 def _fp8_dtype():
     if torch.version.hip is not None and hasattr(torch, "float8_e4m3fnuz"):
@@ -21,6 +42,97 @@ def _fp8_max() -> float:
     return 240.0 if torch.version.hip is not None else 448.0
 
 
+def _arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if major >= 12:
+        return "12.0a"
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    return f"{major}.{minor}"
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+
+    def bias_gelu_quantize(self, input, bias, scale, *, out=None):
+        if out is None:
+            out = torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
+        self._ops.bias_gelu_quantize_fp8_static_bf16(input, bias, scale, out)
+        return out
+
+    def gelu_quantize(self, input, scale, *, out=None):
+        if out is None:
+            out = torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
+        self._ops.gelu_quantize_fp8_static_bf16(input, scale, out)
+        return out
+
+    def channel_scale_quantize(self, input, channel_scale, scale, *, out=None):
+        if out is None:
+            out = torch.empty(input.shape, device=input.device, dtype=_fp8_dtype())
+        self._ops.channel_scale_quantize_fp8_static_bf16(input, channel_scale, scale, out)
+        return out
+
+
+class InstalledOps:
+    def __init__(self, module) -> None:
+        self._module = module
+
+    def bias_gelu_quantize(self, input, bias, scale, *, out=None):
+        return self._module.bias_gelu_quantize_fp8_static_bf16(input, bias, scale, out=out)
+
+    def gelu_quantize(self, input, scale, *, out=None):
+        return self._module.gelu_quantize_fp8_static_bf16(input, scale, out=out)
+
+    def channel_scale_quantize(self, input, channel_scale, scale, *, out=None):
+        return self._module.channel_scale_quantize_fp8_static_bf16(
+            input, channel_scale, scale, out=out
+        )
+
+
+def _preload_cublaslt() -> None:
+    import ctypes
+    import ctypes.util
+
+    for parent in Path(torch.__file__).resolve().parents:
+        candidate = parent / "nvidia" / "cublas" / "lib" / "libcublasLt.so.12"
+        if candidate.exists():
+            ctypes.CDLL(str(candidate), mode=ctypes.RTLD_GLOBAL)
+            return
+    library = ctypes.util.find_library("cublasLt")
+    if library:
+        ctypes.CDLL(library, mode=ctypes.RTLD_GLOBAL)
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    _preload_cublaslt()
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
+    namespace = "flashrt_gemm_epilogues_test"
+    load(
+        name=namespace,
+        sources=[str(PACKAGE / s) for s in _SOURCE_LIST],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+
+
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return InstalledOps(importlib.import_module("flashrt_gemm_epilogues"))
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
 def _reference(input, bias, scale):
     y = input.float()
     if bias is not None:
@@ -30,87 +142,111 @@ def _reference(input, bias, scale):
     return y.to(_fp8_dtype())
 
 
-def _channel_scale_reference(input, channel_scale, scale):
+def _channel_reference(input, channel_scale, scale):
     y = input.float() * channel_scale.float()
     y = torch.clamp(y / scale.float(), -_fp8_max(), _fp8_max())
     return y.to(_fp8_dtype())
 
 
-@pytest.mark.parametrize("shape", [(4, 16), (2, 3, 32)])
-def test_bias_gelu_quantize_fp8_static_bf16(shape):
+def _check_bias_gelu(ops, shape) -> None:
     torch.manual_seed(0)
-    device = torch.device("cuda")
-    input = torch.randn(shape, device=device, dtype=torch.bfloat16).contiguous()
-    bias = torch.randn((shape[-1],), device=device, dtype=torch.bfloat16).contiguous()
-    scale = torch.tensor([0.25], device=device, dtype=torch.float32)
-
-    out = flashrt_ops.bias_gelu_quantize_fp8_static_bf16(input, bias, scale)
+    input = torch.randn(shape, device="cuda", dtype=torch.bfloat16).contiguous()
+    bias = torch.randn((shape[-1],), device="cuda", dtype=torch.bfloat16).contiguous()
+    scale = torch.tensor([0.25], device="cuda", dtype=torch.float32)
+    out = ops.bias_gelu_quantize(input, bias, scale)
     expected = _reference(input, bias, scale)
-
     torch.testing.assert_close(out.float(), expected.float(), rtol=0, atol=0)
+    print(f"PASS bias_gelu_quantize_fp8 shape={shape}: exact")
 
 
-def test_gelu_quantize_fp8_static_bf16_no_bias():
+def _check_gelu(ops) -> None:
     torch.manual_seed(1)
-    device = torch.device("cuda")
-    input = torch.randn((8, 64), device=device, dtype=torch.bfloat16).contiguous()
-    scale = torch.tensor([0.5], device=device, dtype=torch.float32)
-
-    out = flashrt_ops.gelu_quantize_fp8_static_bf16(input, scale)
+    input = torch.randn((8, 64), device="cuda", dtype=torch.bfloat16).contiguous()
+    scale = torch.tensor([0.5], device="cuda", dtype=torch.float32)
+    out = ops.gelu_quantize(input, scale)
     expected = _reference(input, None, scale)
-
     torch.testing.assert_close(out.float(), expected.float(), rtol=0, atol=0)
+    print("PASS gelu_quantize_fp8 (no bias): exact")
 
 
-def test_out_tensor_is_reused():
-    device = torch.device("cuda")
-    input = torch.randn((2, 16), device=device, dtype=torch.bfloat16).contiguous()
-    bias = torch.randn((16,), device=device, dtype=torch.bfloat16).contiguous()
-    scale = torch.tensor([1.0], device=device, dtype=torch.float32)
-    out = torch.empty(input.shape, device=device, dtype=_fp8_dtype())
-
-    returned = flashrt_ops.bias_gelu_quantize_fp8_static_bf16(
-        input, bias, scale, out=out
-    )
-
-    assert returned is out
-
-
-def test_rejects_wrong_bias_shape():
-    device = torch.device("cuda")
-    input = torch.randn((2, 16), device=device, dtype=torch.bfloat16).contiguous()
-    bias = torch.randn((15,), device=device, dtype=torch.bfloat16).contiguous()
-    scale = torch.tensor([1.0], device=device, dtype=torch.float32)
-
-    with pytest.raises(RuntimeError, match="bias length"):
-        flashrt_ops.bias_gelu_quantize_fp8_static_bf16(input, bias, scale)
-
-
-@pytest.mark.parametrize("shape", [(4, 16), (2, 3, 32)])
-def test_channel_scale_quantize_fp8_static_bf16(shape):
+def _check_channel_scale(ops, shape) -> None:
     torch.manual_seed(2)
-    device = torch.device("cuda")
-    input = torch.randn(shape, device=device, dtype=torch.bfloat16).contiguous()
-    channel_scale = torch.randn(
-        (shape[-1],), device=device, dtype=torch.bfloat16
-    ).contiguous()
-    scale = torch.tensor([0.25], device=device, dtype=torch.float32)
-
-    out = flashrt_ops.channel_scale_quantize_fp8_static_bf16(
-        input, channel_scale, scale
-    )
-    expected = _channel_scale_reference(input, channel_scale, scale)
-
+    input = torch.randn(shape, device="cuda", dtype=torch.bfloat16).contiguous()
+    channel_scale = torch.randn((shape[-1],), device="cuda", dtype=torch.bfloat16).contiguous()
+    scale = torch.tensor([0.25], device="cuda", dtype=torch.float32)
+    out = ops.channel_scale_quantize(input, channel_scale, scale)
+    expected = _channel_reference(input, channel_scale, scale)
     torch.testing.assert_close(out.float(), expected.float(), rtol=0, atol=0)
+    print(f"PASS channel_scale_quantize_fp8 shape={shape}: exact")
+
+
+def _check_reuse_and_reject(ops) -> None:
+    input = torch.randn((2, 16), device="cuda", dtype=torch.bfloat16).contiguous()
+    bias = torch.randn((16,), device="cuda", dtype=torch.bfloat16).contiguous()
+    scale = torch.tensor([1.0], device="cuda", dtype=torch.float32)
+    out = torch.empty(input.shape, device="cuda", dtype=_fp8_dtype())
+    returned = ops.bias_gelu_quantize(input, bias, scale, out=out)
+    if returned is not out:
+        raise AssertionError("out tensor must be reused in place")
+    for fn_args, match in [
+        ((input, torch.randn((15,), device="cuda", dtype=torch.bfloat16).contiguous(), scale), "bias length"),
+        ((input, torch.randn((15,), device="cuda", dtype=torch.bfloat16).contiguous(), scale), "channel_scale length"),
+    ]:
+        if match == "bias length":
+            try:
+                ops.bias_gelu_quantize(*fn_args)
+            except RuntimeError as exc:
+                if match not in str(exc):
+                    raise
+            else:
+                raise AssertionError("wrong bias shape must be rejected")
+        else:
+            try:
+                ops.channel_scale_quantize(input, fn_args[1], scale)
+            except RuntimeError as exc:
+                if match not in str(exc):
+                    raise
+            else:
+                raise AssertionError("wrong channel_scale shape must be rejected")
+    print("PASS fp8 epilogue reuse + rejections")
+
+
+def run(args) -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    if not (hasattr(torch, "float8_e4m3fn") or hasattr(torch, "float8_e4m3fnuz")):
+        raise RuntimeError("FP8 support is required")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    shapes = [(4, 16)] if args.mode == "smoke" else [(4, 16), (2, 3, 32)]
+    for shape in shapes:
+        _check_bias_gelu(ops, shape)
+        _check_channel_scale(ops, shape)
+    _check_gelu(ops)
+    _check_reuse_and_reject(ops)
+    print(f"PASS flashrt-gemm-epilogues/fp8-quant {args.backend} mode={args.mode}: "
+          f"{2 * len(shapes) + 2} checks")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    try:
+        run(args)
+    except Exception:
+        import traceback
+        traceback.print_exc()
+        return 1
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(
+            json.dumps({"passed": 1, "total": 1, "backend": args.backend}) + "\n"
+        )
+    return 0
 
 
-def test_rejects_wrong_channel_scale_shape():
-    device = torch.device("cuda")
-    input = torch.randn((2, 16), device=device, dtype=torch.bfloat16).contiguous()
-    channel_scale = torch.randn((15,), device=device, dtype=torch.bfloat16)
-    scale = torch.tensor([1.0], device=device, dtype=torch.float32)
-
-    with pytest.raises(RuntimeError, match="channel_scale length"):
-        flashrt_ops.channel_scale_quantize_fp8_static_bf16(
-            input, channel_scale, scale
-        )
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-nvfp4/VALIDATION.md b/flashrt-nvfp4/VALIDATION.md
index f2cf180..559e654 100644
--- a/flashrt-nvfp4/VALIDATION.md
+++ b/flashrt-nvfp4/VALIDATION.md
@@ -34,15 +34,15 @@ Runtime smoke environment:
 Builder tooling:
 
 - `kernel-builder` 0.16.0-dev0
-- Docker/Nix build wrapper available under
-  `/home/heima/suliang/PI/.hf-kernel-env/bin`
+- Docker/Nix build wrapper available via `KERNEL_BUILDER_DOCKER` or
+  `kernel-builder-docker` on `PATH`
 
 ## Commands
 
 From this package directory:
 
 ```bash
-/home/heima/suliang/PI/.hf-kernel-env/bin/kernel-builder-docker check-config .
+kernel-builder-docker check-config .
 ```
 
 Host-side source-extension correctness was validated with:
diff --git a/flashrt-nvfp4/tests/test_nvfp4_sf_reshape.py b/flashrt-nvfp4/tests/test_nvfp4_sf_reshape.py
index c056cd0..f6fa937 100644
--- a/flashrt-nvfp4/tests/test_nvfp4_sf_reshape.py
+++ b/flashrt-nvfp4/tests/test_nvfp4_sf_reshape.py
@@ -1,19 +1,117 @@
-import pytest
+#!/usr/bin/env python3
+"""Correctness tests for flashrt-nvfp4 (NVFP4 scale-factor layout reshape).
+
+Runs the same checks against the source build and the installed Hub artifact.
+"""
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import os
+import sys
+from pathlib import Path
+
 import torch
 
-from flashrt_nvfp4 import (
-    nvfp4_sf_linear_to_swizzled,
-    nvfp4_sf_swizzled_bytes,
+
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "flashrt-nvfp4"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
 )
 
 
+def _arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if major >= 12:
+        return "12.0a"
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    return f"{major}.{minor}"
+
+
+def _swizzled_bytes(rows: int, D: int) -> int:
+    if rows <= 0:
+        raise ValueError("rows must be positive")
+    if D <= 0 or D % 16 != 0:
+        raise ValueError("D must be positive and divisible by 16")
+    n_blocks = D // 16
+    n_row_super = (rows + 127) // 128
+    n_col_super = (n_blocks + 3) // 4
+    return n_row_super * n_col_super * 512
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+
+    def swizzled_bytes(self, rows: int, D: int) -> int:
+        return _swizzled_bytes(rows, D)
+
+    def linear_to_swizzled(self, scales, *, out=None, is_sfb=False):
+        if out is None:
+            out = torch.zeros(
+                (_swizzled_bytes(scales.shape[0], scales.shape[1] * 16),),
+                device=scales.device,
+                dtype=torch.uint8,
+            )
+        self._ops.nvfp4_sf_linear_to_swizzled(scales, out, scales.shape[1] * 16, bool(is_sfb))
+        return out
+
+
+class InstalledOps:
+    def __init__(self, module) -> None:
+        self._module = module
+
+    def swizzled_bytes(self, rows: int, D: int) -> int:
+        return int(self._module.nvfp4_sf_swizzled_bytes(rows, D))
+
+    def linear_to_swizzled(self, scales, *, out=None, is_sfb=False):
+        return self._module.nvfp4_sf_linear_to_swizzled(scales, out=out, is_sfb=is_sfb)
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
+    namespace = "flashrt_nvfp4_test"
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            str(PACKAGE / "csrc" / "nvfp4_sf_reshape_sm120.cu"),
+        ],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+
+
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return InstalledOps(importlib.import_module("flashrt_nvfp4"))
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
 def _reference_swizzle(scales: torch.Tensor) -> torch.Tensor:
     rows, n_blocks = scales.shape
     n_col_super = (n_blocks + 3) // 4
-    out = torch.zeros(
-        ((rows + 127) // 128) * n_col_super * 512,
-        dtype=torch.uint8,
-    )
+    out = torch.zeros(_swizzled_bytes(rows, n_blocks * 16), dtype=torch.uint8)
     src = scales.cpu()
     for row in range(rows):
         rb = row // 128
@@ -27,67 +125,90 @@ def _reference_swizzle(scales: torch.Tensor) -> torch.Tensor:
     return out
 
 
-@pytest.mark.parametrize(
-    ("rows", "D"),
-    [
-        (1, 1024),
-        (2, 4096),
-        (31, 4096),
-        (32, 4096),
-        (33, 4096),
-        (127, 4096),
-        (128, 4096),
-        (129, 4096),
-        (16, 12288),
-        (64, 16384),
-    ],
-)
-def test_nvfp4_sf_swizzled_bytes(rows, D):
-    n_blocks = D // 16
-    expected = ((rows + 127) // 128) * ((n_blocks + 3) // 4) * 512
-    assert nvfp4_sf_swizzled_bytes(rows, D) == expected
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-@pytest.mark.parametrize(
-    ("rows", "D"),
-    [
-        (1, 1024),
-        (4, 4096),
-        (33, 4096),
-        (128, 4096),
-        (129, 4096),
-        (16, 12288),
-    ],
-)
-def test_nvfp4_sf_linear_to_swizzled(rows, D):
+def _check_bytes(ops) -> None:
+    cases = [
+        (1, 1024), (2, 4096), (31, 4096), (32, 4096), (33, 4096),
+        (127, 4096), (128, 4096), (129, 4096), (16, 12288), (64, 16384),
+    ]
+    for rows, D in cases:
+        got = ops.swizzled_bytes(rows, D)
+        expected = _swizzled_bytes(rows, D)
+        if got != expected:
+            raise AssertionError(f"swizzled_bytes({rows}, {D}) = {got}, expected {expected}")
+    print(f"PASS nvfp4 sf swizzled-bytes: {len(cases)} checks")
+
+
+def _check_reshape(ops, rows: int, D: int) -> None:
     torch.manual_seed(0)
     n_blocks = D // 16
-    scales_cpu = torch.randint(0, 256, (rows, n_blocks), dtype=torch.uint8)
-    scales = scales_cpu.cuda()
+    scales = torch.randint(0, 256, (rows, n_blocks), dtype=torch.uint8)
+    out = ops.linear_to_swizzled(scales.cuda())
+    expected = _reference_swizzle(scales).cuda()
+    torch.testing.assert_close(out, expected)
+    print(f"PASS linear_to_swizzled rows={rows} D={D}: exact")
 
-    out = nvfp4_sf_linear_to_swizzled(scales)
-    expected = _reference_swizzle(scales_cpu).cuda()
 
-    torch.testing.assert_close(out, expected)
+def _check_reuse(ops) -> None:
+    scales = torch.arange(4 * 64, device="cuda", dtype=torch.uint8).reshape(4, 64)
+    out = torch.zeros(_swizzled_bytes(4, 1024), device="cuda", dtype=torch.uint8)
+    returned = ops.linear_to_swizzled(scales, out=out)
+    if returned is not out:
+        raise AssertionError("out tensor must be reused in place")
+    print("PASS linear_to_swizzled reuses out")
 
 
-def test_nvfp4_sf_swizzled_bytes_rejects_invalid_shape():
-    with pytest.raises(ValueError, match="rows"):
-        nvfp4_sf_swizzled_bytes(0, 4096)
-    with pytest.raises(ValueError, match="D"):
-        nvfp4_sf_swizzled_bytes(1, 15)
+def _check_invalid(ops) -> None:
+    for fn, args, match in [
+        (ops.swizzled_bytes, (0, 4096), "rows"),
+        (ops.swizzled_bytes, (1, 15), "D"),
+    ]:
+        try:
+            fn(*args)
+        except (ValueError, RuntimeError) as exc:
+            if match not in str(exc):
+                raise AssertionError(f"expected error containing {match!r}, got {exc}")
+        else:
+            raise AssertionError(f"expected error containing {match!r}")
+    print("PASS swizzled-bytes rejects invalid shape")
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_nvfp4_sf_linear_to_swizzled_reuses_out():
-    scales = torch.arange(4 * 64, device="cuda", dtype=torch.uint8).reshape(4, 64)
-    out = torch.zeros(
-        (nvfp4_sf_swizzled_bytes(4, 1024),),
-        device="cuda",
-        dtype=torch.uint8,
-    )
+def run(args) -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    checks = 0
+    _check_bytes(ops); checks += 1
+    shapes = [(1, 1024), (4, 4096), (33, 4096), (128, 4096), (129, 4096), (16, 12288)]
+    if args.mode == "full":
+        shapes += [(257, 4096), (64, 16384)]
+    for rows, D in shapes:
+        _check_reshape(ops, rows, D)
+        checks += 1
+    _check_reuse(ops); checks += 1
+    _check_invalid(ops); checks += 1
+    print(f"PASS flashrt-nvfp4 {args.backend} mode={args.mode}: {checks} checks")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    try:
+        run(args)
+    except Exception:
+        import traceback
+        traceback.print_exc()
+        return 1
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(
+            json.dumps({"passed": 1, "total": 1, "backend": args.backend}) + "\n"
+        )
+    return 0
 
-    returned = nvfp4_sf_linear_to_swizzled(scales, out=out)
 
-    assert returned is out
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-qkv-cache-rope/benchmarks/benchmark.py b/flashrt-qkv-cache-rope/benchmarks/benchmark.py
index 7346291..61781a7 100644
--- a/flashrt-qkv-cache-rope/benchmarks/benchmark.py
+++ b/flashrt-qkv-cache-rope/benchmarks/benchmark.py
@@ -17,6 +17,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 PACKAGE = ROOT / "flashrt-qkv-cache-rope"
 REGISTRATION_INCLUDE = (
@@ -884,7 +894,9 @@ def main() -> None:
     parser.add_argument("--p99-abs-limit", type=float, default=0.015625)
     parser.add_argument("--output", default=None)
     parser.add_argument("--markdown", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
diff --git a/flashrt-qkv-epilogue-train/benchmarks/benchmark.py b/flashrt-qkv-epilogue-train/benchmarks/benchmark.py
new file mode 100644
index 0000000..1ecd58f
--- /dev/null
+++ b/flashrt-qkv-epilogue-train/benchmarks/benchmark.py
@@ -0,0 +1,98 @@
+#!/usr/bin/env python3
+"""QKV + RoPE epilogue reference benchmark (reference-only package, no CUDA kernel yet).
+
+This package publishes a reference/autograd API only. No speedup is claimed;
+the table records reference latency vs an equivalent torch.compile region.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from pathlib import Path
+
+import torch
+
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
+def elapsed_us(fn, warmup: int = 10, repeats: int = 50) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(repeats):
+        fn()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) * 1000.0 / repeats
+
+
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        root = Path(__file__).resolve().parents[1]
+        sys.path.insert(0, str(root / "torch-ext"))
+        return importlib.import_module("flashrt_qkv_epilogue_train")
+    if artifact:
+        sys.path.insert(0, artifact)
+    return importlib.import_module("flashrt_qkv_epilogue_train")
+
+
+def run_case(ops, label: str, b: int, t: int, hid: int, qh: int, kvh: int, d: int) -> dict:
+    torch.manual_seed(0)
+    x = torch.randn(b, t, hid, device="cuda", dtype=torch.bfloat16)
+    wq = torch.randn(qh * d, hid, device="cuda", dtype=torch.bfloat16)
+    wk = torch.randn(kvh * d, hid, device="cuda", dtype=torch.bfloat16)
+    wv = torch.randn(kvh * d, hid, device="cuda", dtype=torch.bfloat16)
+    cos = torch.randn(b, t, d, device="cuda", dtype=torch.bfloat16)
+    sin = torch.randn(b, t, d, device="cuda", dtype=torch.bfloat16)
+
+    def ref():
+        return ops.qkv_rope_reference(x, wq, wk, wv, cos, sin, qh, kvh, d)
+
+    ref_us = elapsed_us(ref)
+    compiled = torch.compile(ref, mode="reduce-overhead")
+    compiled()
+    compile_us = elapsed_us(lambda: compiled())
+    return {
+        "label": label, "B": b, "T": t, "H": hid, "qh": qh, "kvh": kvh, "D": d,
+        "reference_us": ref_us, "compile_us": compile_us,
+        "ratio": compile_us / ref_us,
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="installed")
+    parser.add_argument("--artifact")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
+    args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
+    ops = load_ops(args.backend, args.artifact)
+    print("label,B,T,H,qh,kvh,D,reference_us,compile_us,compile_over_reference")
+    for label, b, t, hid, qh, kvh, d in [
+        ("qkv_rope_2k", 1, 2048, 4096, 32, 8, 128),
+        ("qkv_rope_4k", 1, 4096, 4096, 32, 8, 128),
+    ]:
+        r = run_case(ops, label, b, t, hid, qh, kvh, d)
+        print(
+            f"{r['label']},{r['B']},{r['T']},{r['H']},{r['qh']},{r['kvh']},"
+            f"{r['D']},{r['reference_us']:.3f},{r['compile_us']:.3f},"
+            f"{r['ratio']:.2f}x"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-residual-norm-quant/benchmarks/benchmark.py b/flashrt-residual-norm-quant/benchmarks/benchmark.py
index beb5667..c024d42 100644
--- a/flashrt-residual-norm-quant/benchmarks/benchmark.py
+++ b/flashrt-residual-norm-quant/benchmarks/benchmark.py
@@ -17,6 +17,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 PACKAGE = ROOT / "flashrt-residual-norm-quant"
 REGISTRATION_INCLUDE = (
@@ -281,7 +291,9 @@ def main() -> None:
     parser.add_argument("--p99-abs-limit", type=float, default=0.5)
     parser.add_argument("--output", default=None)
     parser.add_argument("--markdown", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
diff --git a/flashrt-rope-train/benchmarks/benchmark.py b/flashrt-rope-train/benchmarks/benchmark.py
new file mode 100755
index 0000000..8620abd
--- /dev/null
+++ b/flashrt-rope-train/benchmarks/benchmark.py
@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+"""RoPE training smoke test (reference-only package, no CUDA kernel yet).
+
+Not a benchmark: runs each shape once and checks the op produces
+same-shaped, finite outputs. Exits non-zero on any failure.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from pathlib import Path
+
+import torch
+
+
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        root = Path(__file__).resolve().parents[1]
+        sys.path.insert(0, str(root / "torch-ext"))
+        return importlib.import_module("flashrt_rope_train")
+    if artifact:
+        sys.path.insert(0, artifact)
+    return importlib.import_module("flashrt_rope_train")
+
+
+def apply_max_mem_cap(max_mem_gb: float) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
+def run_case(ops, label: str, b: int, h: int, t: int, d: int) -> None:
+    torch.manual_seed(0)
+    with torch.no_grad():
+        q = torch.randn(b, h, t, d, device="cuda", dtype=torch.bfloat16)
+        k = torch.randn(b, h, t, d, device="cuda", dtype=torch.bfloat16)
+        cos = torch.randn(b, t, d, device="cuda", dtype=torch.bfloat16)
+        sin = torch.randn(b, t, d, device="cuda", dtype=torch.bfloat16)
+
+        q_out, k_out = ops.apply_rope_train(q, k, cos, sin)
+    del q, k, cos, sin
+
+    assert q_out.shape == (b, h, t, d) and q_out.dtype == torch.bfloat16, (
+        f"{label}: bad q_out {tuple(q_out.shape)}/{q_out.dtype}"
+    )
+    assert k_out.shape == (b, h, t, d) and k_out.dtype == torch.bfloat16, (
+        f"{label}: bad k_out {tuple(k_out.shape)}/{k_out.dtype}"
+    )
+    assert torch.isfinite(q_out).all() and torch.isfinite(k_out).all(), (
+        f"{label}: non-finite output"
+    )
+    print(f"ok {label} B={b} H={h} T={t} D={d} -> {tuple(q_out.shape)} {q_out.dtype}")
+
+    del q_out, k_out
+    torch.cuda.empty_cache()
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="installed")
+    parser.add_argument("--artifact")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
+    args = parser.parse_args()
+    apply_max_mem_cap(args.max_mem_gb)
+    ops = load_ops(args.backend, args.artifact)
+    for label, b, h, t, d in [
+        ("rope_2k", 1, 32, 2048, 128),
+        ("rope_4k", 1, 32, 4096, 128),
+    ]:
+        run_case(ops, label, b, h, t, d)
+    print("all smoke cases passed")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-rope-train/torch-ext/flashrt_rope_train/__init__.py b/flashrt-rope-train/torch-ext/flashrt_rope_train/__init__.py
index 6110650..c79083c 100644
--- a/flashrt-rope-train/torch-ext/flashrt_rope_train/__init__.py
+++ b/flashrt-rope-train/torch-ext/flashrt_rope_train/__init__.py
@@ -12,13 +12,17 @@ def _flashrt_training_package_marker(self, x):
 def rotate_half(x: torch.Tensor) -> torch.Tensor:
     half=x.shape[-1]//2; return torch.cat((-x[...,half:], x[...,:half]), dim=-1)
 def _align(freq: torch.Tensor, x: torch.Tensor, unsqueeze_dim: int) -> torch.Tensor:
+    if freq.dim() == 2:
+        freq = freq.reshape((1,) * (x.dim() - 2) + freq.shape)
+        return freq
     while freq.dim() < x.dim(): freq = freq.unsqueeze(unsqueeze_dim)
     return freq
+def _rope_one(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, qd: torch.dtype, unsqueeze_dim: int) -> torch.Tensor:
+    xf = x.to(qd)
+    return (xf * _align(cos, x, unsqueeze_dim).to(qd) + rotate_half(xf) * _align(sin, x, unsqueeze_dim).to(qd)).to(x.dtype)
 def apply_rope_train(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int = 1):
-    c_q=_align(cos,q,int(unsqueeze_dim)); s_q=_align(sin,q,int(unsqueeze_dim)); c_k=_align(cos,k,int(unsqueeze_dim)); s_k=_align(sin,k,int(unsqueeze_dim))
     qd = torch.float64 if q.dtype == torch.float64 else torch.float32
-    kd = torch.float64 if k.dtype == torch.float64 else torch.float32
-    return (q.to(qd)*c_q.to(qd)+rotate_half(q.to(qd))*s_q.to(qd)).to(q.dtype), (k.to(kd)*c_k.to(kd)+rotate_half(k.to(kd))*s_k.to(kd)).to(k.dtype)
+    return _rope_one(q, cos, sin, qd, unsqueeze_dim), _rope_one(k, cos, sin, qd, unsqueeze_dim)
 def apply_rope_backward_reference(dq: torch.Tensor, dk: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, unsqueeze_dim: int = 1):
     return apply_rope_train(dq, dk, cos, -sin, unsqueeze_dim)
 def backend_marker(x: torch.Tensor) -> torch.Tensor:
diff --git a/flashrt-siglip-fwd-fusion/benchmarks/benchmark.py b/flashrt-siglip-fwd-fusion/benchmarks/benchmark.py
new file mode 100644
index 0000000..2adc72a
--- /dev/null
+++ b/flashrt-siglip-fwd-fusion/benchmarks/benchmark.py
@@ -0,0 +1,105 @@
+#!/usr/bin/env python3
+"""SigLIP forward-fusion reference benchmark (reference-only package, no CUDA kernel yet).
+
+No speedup is claimed; the table records reference latency vs an equivalent
+torch.compile region.
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from pathlib import Path
+
+import torch
+
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
+def elapsed_us(fn, warmup: int = 10, repeats: int = 50) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(repeats):
+        fn()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) * 1000.0 / repeats
+
+
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        root = Path(__file__).resolve().parents[1]
+        sys.path.insert(0, str(root / "torch-ext"))
+        return importlib.import_module("flashrt_siglip_fwd_fusion")
+    if artifact:
+        sys.path.insert(0, artifact)
+    return importlib.import_module("flashrt_siglip_fwd_fusion")
+
+
+def run_case(ops, label: str, b: int, t: int, d: int) -> dict:
+    torch.manual_seed(0)
+    x = torch.randn(b, t, d, device="cuda", dtype=torch.bfloat16)
+    residual = torch.randn(b, t, d, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(d, device="cuda", dtype=torch.bfloat16) + 1.0
+    bias = torch.randn(d, device="cuda", dtype=torch.bfloat16)
+
+    def ln_ref():
+        return ops.siglip_residual_layernorm_fwd(x, residual, weight, bias)
+
+    def gelu_ref():
+        return ops.siglip_gelu_fwd(x, bias)
+
+    ln_us = elapsed_us(ln_ref)
+    gelu_us = elapsed_us(gelu_ref)
+    ln_c = torch.compile(ln_ref, mode="reduce-overhead")
+    ln_c()
+    ln_c_us = elapsed_us(lambda: ln_c())
+    gelu_c = torch.compile(gelu_ref, mode="reduce-overhead")
+    gelu_c()
+    gelu_c_us = elapsed_us(lambda: gelu_c())
+    return {
+        "label": label, "B": b, "T": t, "D": d,
+        "ln_us": ln_us, "ln_compile_us": ln_c_us,
+        "gelu_us": gelu_us, "gelu_compile_us": gelu_c_us,
+        "ln_ratio": ln_c_us / ln_us,
+        "gelu_ratio": gelu_c_us / gelu_us,
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="installed")
+    parser.add_argument("--artifact")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
+    args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
+    ops = load_ops(args.backend, args.artifact)
+    print("label,B,T,D,ln_us,ln_compile_us,gelu_us,gelu_compile_us,ln_compile_over_ref,gelu_compile_over_ref")
+    for label, b, t, d in [
+        ("siglip_1k", 1, 1024, 768),
+        ("siglip_2k", 2, 1024, 1024),
+    ]:
+        r = run_case(ops, label, b, t, d)
+        print(
+            f"{r['label']},{r['B']},{r['T']},{r['D']},{r['ln_us']:.3f},"
+            f"{r['ln_compile_us']:.3f},{r['gelu_us']:.3f},{r['gelu_compile_us']:.3f},"
+            f"{r['ln_ratio']:.2f}x,{r['gelu_ratio']:.2f}x"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-smallm-gemm/VALIDATION.md b/flashrt-smallm-gemm/VALIDATION.md
index dd02310..bf8b2f9 100644
--- a/flashrt-smallm-gemm/VALIDATION.md
+++ b/flashrt-smallm-gemm/VALIDATION.md
@@ -104,8 +104,7 @@ Worst built-artifact case:
 ## Known Gaps
 
 - `build.toml`, `flake.nix`, and `flake.lock` are present.
-- `/home/heima/suliang/PI/.hf-kernel-env/bin/kernel-builder-docker
-  check-config .` passed for this package.
+- `kernel-builder-docker check-config .` passed for this package.
 - `kernel-builder build --variant torch211-cxx11-cu128-x86_64-linux` passed
   for this package.
 - Full `kernel-builder build-and-copy` matrix has not been run for this
diff --git a/flashrt-smallm-gemm/tests/test_nvfp4_w4a4_decode_matvec.py b/flashrt-smallm-gemm/tests/test_nvfp4_w4a4_decode_matvec.py
index 3b7d45b..d05a09f 100644
--- a/flashrt-smallm-gemm/tests/test_nvfp4_w4a4_decode_matvec.py
+++ b/flashrt-smallm-gemm/tests/test_nvfp4_w4a4_decode_matvec.py
@@ -1,7 +1,28 @@
-import pytest
+#!/usr/bin/env python3
+"""Correctness tests for flashrt-smallm-gemm (NVFP4 W4A4 M=1 decode matvec)."""
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import os
+import sys
+from pathlib import Path
+
 import torch
 
-from flashrt_smallm_gemm import nvfp4_w4a4_decode_matvec_bf16out
+
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "flashrt-smallm-gemm"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
+)
 
 
 def _swizzled_bytes(rows: int, D: int) -> int:
@@ -12,10 +33,7 @@ def _swizzled_bytes(rows: int, D: int) -> int:
 def _swizzle_scales(scales: torch.Tensor) -> torch.Tensor:
     rows, n_blocks = scales.shape
     n_col_super = (n_blocks + 3) // 4
-    out = torch.zeros(
-        (_swizzled_bytes(rows, n_blocks * 16),),
-        dtype=torch.uint8,
-    )
+    out = torch.zeros(_swizzled_bytes(rows, n_blocks * 16), dtype=torch.uint8)
     src = scales.cpu()
     for row in range(rows):
         rb = row // 128
@@ -29,9 +47,68 @@ def _swizzle_scales(scales: torch.Tensor) -> torch.Tensor:
     return out
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-@pytest.mark.parametrize("K", [4096, 12288])
-def test_nvfp4_w4a4_decode_matvec_constant_inputs(K):
+def _arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if major >= 12:
+        return "12.0a"
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    return f"{major}.{minor}"
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+
+    def decode_matvec(self, a_packed, b_packed, sfa, sfb, *, alpha=1.0, out=None):
+        if out is None:
+            out = torch.empty((b_packed.shape[0],), device=b_packed.device, dtype=torch.bfloat16)
+        self._ops.nvfp4_w4a4_decode_matvec_bf16out(a_packed, b_packed, sfa, sfb, out, float(alpha))
+        return out
+
+
+class InstalledOps:
+    def __init__(self, module) -> None:
+        self._module = module
+
+    def decode_matvec(self, a_packed, b_packed, sfa, sfb, *, alpha=1.0, out=None):
+        return self._module.nvfp4_w4a4_decode_matvec_bf16out(
+            a_packed, b_packed, sfa, sfb, alpha=alpha, out=out
+        )
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
+    namespace = "flashrt_smallm_gemm_test"
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            str(PACKAGE / "csrc" / "fp4_w4a4_matvec_sm120.cu"),
+        ],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"],
+        verbose=False,
+    )
+    return SourceOps(namespace)
+
+
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return InstalledOps(importlib.import_module("flashrt_smallm_gemm"))
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
+
+
+def _check_constant_inputs(ops, K: int) -> None:
     torch.manual_seed(0)
     N = 16
     alpha = 0.5
@@ -40,20 +117,13 @@ def test_nvfp4_w4a4_decode_matvec_constant_inputs(K):
     sfa = _swizzle_scales(torch.full((1, K // 16), 0x38, dtype=torch.uint8)).cuda()
     sfb = _swizzle_scales(torch.full((N, K // 16), 0x38, dtype=torch.uint8)).cuda()
 
-    out = nvfp4_w4a4_decode_matvec_bf16out(
-        a_packed,
-        b_packed,
-        sfa,
-        sfb,
-        alpha=alpha,
-    )
-
+    out = ops.decode_matvec(a_packed, b_packed, sfa, sfb, alpha=alpha)
     expected = torch.full((N,), K * 0.25 * alpha, device="cuda", dtype=torch.bfloat16)
     torch.testing.assert_close(out, expected)
+    print(f"PASS w4a4 decode constant K={K}: exact")
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_nvfp4_w4a4_decode_matvec_reuses_out():
+def _check_reuse(ops) -> None:
     K = 4096
     N = 8
     a_packed = torch.full((K // 2,), 0x11, device="cuda", dtype=torch.uint8)
@@ -61,26 +131,61 @@ def test_nvfp4_w4a4_decode_matvec_reuses_out():
     sfa = _swizzle_scales(torch.full((1, K // 16), 0x38, dtype=torch.uint8)).cuda()
     sfb = _swizzle_scales(torch.full((N, K // 16), 0x38, dtype=torch.uint8)).cuda()
     out = torch.empty((N,), device="cuda", dtype=torch.bfloat16)
+    returned = ops.decode_matvec(a_packed, b_packed, sfa, sfb, out=out)
+    if returned is not out:
+        raise AssertionError("out tensor must be reused in place")
+    print("PASS w4a4 decode reuses out")
 
-    returned = nvfp4_w4a4_decode_matvec_bf16out(
-        a_packed,
-        b_packed,
-        sfa,
-        sfb,
-        out=out,
-    )
 
-    assert returned is out
-
-
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-def test_nvfp4_w4a4_decode_matvec_rejects_unsupported_k():
+def _check_rejects_unsupported_k(ops) -> None:
     K = 8192
     N = 8
     a_packed = torch.full((K // 2,), 0x11, device="cuda", dtype=torch.uint8)
     b_packed = torch.full((N, K // 2), 0x11, device="cuda", dtype=torch.uint8)
-    sfa = torch.zeros((_swizzled_bytes(1, K),), device="cuda", dtype=torch.uint8)
-    sfb = torch.zeros((_swizzled_bytes(N, K),), device="cuda", dtype=torch.uint8)
-
-    with pytest.raises(RuntimeError, match="K=4096 and K=12288"):
-        nvfp4_w4a4_decode_matvec_bf16out(a_packed, b_packed, sfa, sfb)
+    sfa = torch.zeros(_swizzled_bytes(1, K), device="cuda", dtype=torch.uint8)
+    sfb = torch.zeros(_swizzled_bytes(N, K), device="cuda", dtype=torch.uint8)
+    try:
+        ops.decode_matvec(a_packed, b_packed, sfa, sfb)
+    except RuntimeError as exc:
+        if "K=4096 and K=12288" not in str(exc):
+            raise
+    else:
+        raise AssertionError("K=8192 must be rejected")
+    print("PASS w4a4 decode rejects unsupported K")
+
+
+def run(args) -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    for K in ([4096, 12288] if args.mode == "full" else [4096]):
+        _check_constant_inputs(ops, K)
+    _check_reuse(ops)
+    _check_rejects_unsupported_k(ops)
+    print(f"PASS flashrt-smallm-gemm {args.backend} mode={args.mode}: "
+          f"{2 + len([4096, 12288] if args.mode == 'full' else [4096])} checks")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    try:
+        run(args)
+    except Exception:
+        import traceback
+        traceback.print_exc()
+        return 1
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(
+            json.dumps({"passed": 1, "total": 1, "backend": args.backend}) + "\n"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-spatiotemporal-layout/benchmarks/benchmark.py b/flashrt-spatiotemporal-layout/benchmarks/benchmark.py
index f49ad85..713f254 100644
--- a/flashrt-spatiotemporal-layout/benchmarks/benchmark.py
+++ b/flashrt-spatiotemporal-layout/benchmarks/benchmark.py
@@ -6,6 +6,7 @@
 import argparse
 import ctypes
 import ctypes.util
+import importlib
 import json
 import os
 import sys
@@ -15,6 +16,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 PACKAGE = ROOT / "flashrt-spatiotemporal-layout"
 REGISTRATION_INCLUDE = (
@@ -221,7 +232,9 @@ def main():
     parser.add_argument("--iters", type=int, default=20)
     parser.add_argument("--output", default=None)
     parser.add_argument("--markdown", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
     torch.manual_seed(61)
diff --git a/flashrt-vla-residual-gates/benchmarks/benchmark.py b/flashrt-vla-residual-gates/benchmarks/benchmark.py
index 096308e..f728635 100644
--- a/flashrt-vla-residual-gates/benchmarks/benchmark.py
+++ b/flashrt-vla-residual-gates/benchmarks/benchmark.py
@@ -6,6 +6,7 @@
 import argparse
 import ctypes
 import ctypes.util
+import importlib
 import json
 import math
 import os
@@ -16,6 +17,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 PACKAGE = ROOT / "flashrt-vla-residual-gates"
 REGISTRATION_INCLUDE = (
@@ -167,17 +178,19 @@ def metrics(got_parts, expected_parts):
 def run_one(ops, name: str, rows: tuple[int, int, int], dim: int, args) -> Result:
     v, a, u = make_case(rows, dim)
     ops.joint3_bias_gate_residual_action_nobias_bf16(
-        v[0], v[1], v[2], v[3], v[4],
-        a[0], a[1], a[3], a[4],
-        u[0], u[1], u[2],
+        v[0], v[1], v[2], v[3],
+        a[0], a[1], a[3],
+        u[0], u[1],
+        v[4], a[4], u[2],
     )
     expected = torch_ref(v, a, u)
     p99_abs, cosine = metrics((v[4], a[4], u[2]), expected)
     flashrt_us = time_us(
         lambda: ops.joint3_bias_gate_residual_action_nobias_bf16(
-            v[0], v[1], v[2], v[3], v[4],
-            a[0], a[1], a[3], a[4],
-            u[0], u[1], u[2],
+            v[0], v[1], v[2], v[3],
+            a[0], a[1], a[3],
+            u[0], u[1],
+            v[4], a[4], u[2],
         ),
         args.warmup,
         args.iters,
@@ -226,7 +239,9 @@ def main() -> None:
     parser.add_argument("--p99-abs-limit", type=float, default=0.0)
     parser.add_argument("--output", default=None)
     parser.add_argument("--markdown", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
diff --git a/flashrt-vla-video/VALIDATION.md b/flashrt-vla-video/VALIDATION.md
index 6fa2e84..94ebb2d 100644
--- a/flashrt-vla-video/VALIDATION.md
+++ b/flashrt-vla-video/VALIDATION.md
@@ -14,7 +14,7 @@ Environment:
 Config check:
 
 ```bash
-/home/heima/suliang/PI/.hf-kernel-env/bin/kernel-builder-docker check-config .
+kernel-builder-docker check-config .
 ```
 
 Result: passed.
diff --git a/flashrt-vla-video/tests/test_q_norm_rope_bf16.py b/flashrt-vla-video/tests/test_q_norm_rope_bf16.py
index 780af56..b07188f 100644
--- a/flashrt-vla-video/tests/test_q_norm_rope_bf16.py
+++ b/flashrt-vla-video/tests/test_q_norm_rope_bf16.py
@@ -1,69 +1,150 @@
-import pytest
+#!/usr/bin/env python3
+"""Correctness tests for flashrt-vla-video (Q/K RMSNorm + RoPE, packed split)."""
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import os
+import sys
+from pathlib import Path
+
 import torch
 
-from flashrt_vla_video import (
-    k_norm_rope_v_cache_bf16,
-    q_norm_rope_bf16,
-    qkv_split_norm_rope_bf16,
+
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "flashrt-vla-video"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
 )
 
 
-def _reference_norm_rope(x, weight, cos, sin, eps=1e-6):
-    half = x.shape[-1] // 2
-    rstd = torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps)
-    normed = x.float() * rstd * weight.float()
-    lo = normed[..., :half]
-    hi = normed[..., half:]
-    out_lo = lo * cos.float() - hi * sin.float()
-    out_hi = hi * cos.float() + lo * sin.float()
-    return torch.cat([out_lo, out_hi], dim=-1).to(torch.bfloat16)
+def _arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if major >= 12:
+        return "12.0a"
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    return f"{major}.{minor}"
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-@pytest.mark.parametrize("shape", [(1, 128), (8, 128), (2, 4, 128)])
-def test_q_norm_rope_bf16(shape):
-    torch.manual_seed(0)
-    q = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.2).contiguous()
-    weight = (torch.randn(128, device="cuda", dtype=torch.bfloat16) * 0.1 + 1).contiguous()
-    cos = torch.randn(64, device="cuda", dtype=torch.bfloat16).contiguous()
-    sin = torch.randn(64, device="cuda", dtype=torch.bfloat16).contiguous()
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
 
-    out = q_norm_rope_bf16(q, weight, cos, sin)
-    ref = _reference_norm_rope(q, weight, cos, sin)
+    def q_norm_rope(self, q, weight, cos, sin, *, eps=1e-6, out=None):
+        if out is None:
+            out = torch.empty_like(q)
+        self._ops.q_norm_rope_bf16(q, weight, cos, sin, out, float(eps))
+        return out
 
-    torch.testing.assert_close(out.float(), ref.float(), atol=0.03125, rtol=0)
+    def k_norm_rope_v_cache(self, k, v, weight, cos, sin, *, eps=1e-6, k_out=None, v_out=None):
+        if k_out is None:
+            k_out = torch.empty_like(k)
+        if v_out is None:
+            v_out = torch.empty_like(v)
+        self._ops.k_norm_rope_v_cache_bf16(k, v, weight, cos, sin, k_out, v_out, float(eps))
+        return k_out, v_out
+
+    def qkv_split_norm_rope(
+        self, packed, qw, kw, freqs_re, freqs_im, *, heads, head_dim, seq_len=None,
+        q_out=None, k_out=None, eps=1e-6,
+    ):
+        if seq_len is None:
+            seq_len = packed.shape[1]
+        if q_out is None:
+            q_out = torch.empty((packed.shape[0], packed.shape[1], heads, head_dim),
+                                device=packed.device, dtype=packed.dtype)
+        if k_out is None:
+            k_out = torch.empty_like(q_out)
+        self._ops.qkv_split_norm_rope_bf16(
+            packed, qw, kw, freqs_re, freqs_im, q_out, k_out,
+            int(heads), int(head_dim), int(seq_len), float(eps),
+        )
+        return q_out, k_out
+
+
+class InstalledOps:
+    def __init__(self, module) -> None:
+        self._module = module
+
+    def q_norm_rope(self, q, weight, cos, sin, *, eps=1e-6, out=None):
+        return self._module.q_norm_rope_bf16(q, weight, cos, sin, out=out, eps=eps)
+
+    def k_norm_rope_v_cache(self, k, v, weight, cos, sin, *, eps=1e-6, k_out=None, v_out=None):
+        return self._module.k_norm_rope_v_cache_bf16(
+            k, v, weight, cos, sin, k_out=k_out, v_out=v_out, eps=eps
+        )
+
+    def qkv_split_norm_rope(
+        self, packed, qw, kw, freqs_re, freqs_im, *, heads, head_dim, seq_len=None,
+        q_out=None, k_out=None, eps=1e-6,
+    ):
+        return self._module.qkv_split_norm_rope_bf16(
+            packed, qw, kw, freqs_re, freqs_im, heads=heads, head_dim=head_dim,
+            seq_len=seq_len, q_out=q_out, k_out=k_out, eps=eps,
+        )
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
+    namespace = "flashrt_vla_video_test"
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            str(PACKAGE / "csrc" / "q_norm_rope_bf16.cu"),
+        ],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"],
+        verbose=False,
+    )
+    return SourceOps(namespace)
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-@pytest.mark.parametrize("shape", [(1, 128), (8, 128), (2, 4, 128)])
-def test_k_norm_rope_v_cache_bf16(shape):
-    torch.manual_seed(1)
-    k = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.2).contiguous()
-    v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.2).contiguous()
-    weight = (torch.randn(128, device="cuda", dtype=torch.bfloat16) * 0.1 + 1).contiguous()
-    cos = torch.randn(64, device="cuda", dtype=torch.bfloat16).contiguous()
-    sin = torch.randn(64, device="cuda", dtype=torch.bfloat16).contiguous()
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return InstalledOps(importlib.import_module("flashrt_vla_video"))
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
 
-    k_out, v_out = k_norm_rope_v_cache_bf16(k, v, weight, cos, sin)
-    k_ref = _reference_norm_rope(k, weight, cos, sin)
 
-    torch.testing.assert_close(k_out.float(), k_ref.float(), atol=0.03125, rtol=0)
-    torch.testing.assert_close(v_out, v)
+def _ref_norm_rope(x, weight, cos, sin, eps=1e-6):
+    half = x.shape[-1] // 2
+    rstd = torch.rsqrt(x.float().square().mean(dim=-1, keepdim=True) + eps)
+    normed = x.float() * rstd * weight.float()
+    lo = normed[..., :half]
+    hi = normed[..., half:]
+    out_lo = lo * cos.float() - hi * sin.float()
+    out_hi = hi * cos.float() + lo * sin.float()
+    return torch.cat([out_lo, out_hi], dim=-1).to(torch.bfloat16)
 
 
-def _reference_qkv_split_norm_rope(
-    packed_qkv, norm_q_weight, norm_k_weight, freqs_re, freqs_im, heads, head_dim, eps=1e-6
-):
-    batch, tokens, _ = packed_qkv.shape
+def _ref_qkv_split(packed, qw, kw, freqs_re, freqs_im, heads, head_dim, eps=1e-6):
+    batch, tokens, _ = packed.shape
     dim = heads * head_dim
-    q = packed_qkv[..., :dim].reshape(batch, tokens, heads, head_dim)
-    k = packed_qkv[..., dim : 2 * dim].reshape(batch, tokens, heads, head_dim)
+    q = packed[..., :dim].reshape(batch, tokens, heads, head_dim)
+    k = packed[..., dim : 2 * dim].reshape(batch, tokens, heads, head_dim)
     qf = q.float()
     kf = k.float()
     qn = qf * torch.rsqrt((qf * qf).mean(dim=(-2, -1), keepdim=True) + eps)
     kn = kf * torch.rsqrt((kf * kf).mean(dim=(-2, -1), keepdim=True) + eps)
-    qn = qn * norm_q_weight.reshape(1, 1, heads, head_dim).float()
-    kn = kn * norm_k_weight.reshape(1, 1, heads, head_dim).float()
+    qn = qn * qw.reshape(1, 1, heads, head_dim).float()
+    kn = kn * kw.reshape(1, 1, heads, head_dim).float()
 
     def rope(x):
         xr = x[..., 0::2].float()
@@ -73,44 +154,89 @@ def rope(x):
         out = torch.empty_like(x, dtype=torch.float32)
         out[..., 0::2] = xr * fr - xi * fi
         out[..., 1::2] = xr * fi + xi * fr
-        out = out.to(torch.bfloat16)
-        return out
+        return out.to(torch.bfloat16)
 
     return rope(qn), rope(kn)
 
 
-@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required")
-@pytest.mark.parametrize("tokens", [4, 64])
-def test_qkv_split_norm_rope_bf16(tokens):
+def _check_q_norm(ops, shape) -> None:
+    torch.manual_seed(0)
+    q = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.2).contiguous()
+    weight = (torch.randn(128, device="cuda", dtype=torch.bfloat16) * 0.1 + 1).contiguous()
+    cos = torch.randn(64, device="cuda", dtype=torch.bfloat16).contiguous()
+    sin = torch.randn(64, device="cuda", dtype=torch.bfloat16).contiguous()
+    out = ops.q_norm_rope(q, weight, cos, sin)
+    ref = _ref_norm_rope(q, weight, cos, sin)
+    torch.testing.assert_close(out.float(), ref.float(), atol=0.03125, rtol=0)
+    print(f"PASS q_norm_rope shape={shape}")
+
+
+def _check_k_norm(ops, shape) -> None:
+    torch.manual_seed(1)
+    k = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.2).contiguous()
+    v = (torch.randn(shape, device="cuda", dtype=torch.bfloat16) * 0.2).contiguous()
+    weight = (torch.randn(128, device="cuda", dtype=torch.bfloat16) * 0.1 + 1).contiguous()
+    cos = torch.randn(64, device="cuda", dtype=torch.bfloat16).contiguous()
+    sin = torch.randn(64, device="cuda", dtype=torch.bfloat16).contiguous()
+    k_out, v_out = ops.k_norm_rope_v_cache(k, v, weight, cos, sin)
+    k_ref = _ref_norm_rope(k, weight, cos, sin)
+    torch.testing.assert_close(k_out.float(), k_ref.float(), atol=0.03125, rtol=0)
+    torch.testing.assert_close(v_out, v)
+    print(f"PASS k_norm_rope_v_cache shape={shape}")
+
+
+def _check_qkv_split(ops, tokens) -> None:
     torch.manual_seed(2)
     heads = 24
     head_dim = 128
     dim = heads * head_dim
-    packed_qkv = (
-        torch.randn((1, tokens, 3 * dim), device="cuda", dtype=torch.bfloat16) * 0.2
-    ).contiguous()
-    norm_q_weight = (
-        torch.randn(dim, device="cuda", dtype=torch.bfloat16) * 0.1 + 1
-    ).contiguous()
-    norm_k_weight = (
-        torch.randn(dim, device="cuda", dtype=torch.bfloat16) * 0.1 + 1
-    ).contiguous()
+    packed = (torch.randn((1, tokens, 3 * dim), device="cuda", dtype=torch.bfloat16) * 0.2).contiguous()
+    qw = (torch.randn(dim, device="cuda", dtype=torch.bfloat16) * 0.1 + 1).contiguous()
+    kw = (torch.randn(dim, device="cuda", dtype=torch.bfloat16) * 0.1 + 1).contiguous()
     freqs_re = torch.randn((128, head_dim // 2), device="cuda", dtype=torch.float32).contiguous()
     freqs_im = torch.randn((128, head_dim // 2), device="cuda", dtype=torch.float32).contiguous()
-
-    q_out, k_out = qkv_split_norm_rope_bf16(
-        packed_qkv,
-        norm_q_weight,
-        norm_k_weight,
-        freqs_re,
-        freqs_im,
-        heads=heads,
-        head_dim=head_dim,
-        seq_len=tokens,
-    )
-    q_ref, k_ref = _reference_qkv_split_norm_rope(
-        packed_qkv, norm_q_weight, norm_k_weight, freqs_re, freqs_im, heads, head_dim
-    )
-
+    q_out, k_out = ops.qkv_split_norm_rope(packed, qw, kw, freqs_re, freqs_im, heads=heads, head_dim=head_dim)
+    q_ref, k_ref = _ref_qkv_split(packed, qw, kw, freqs_re, freqs_im, heads, head_dim)
     torch.testing.assert_close(q_out.float(), q_ref.float(), atol=0.03125, rtol=0)
     torch.testing.assert_close(k_out.float(), k_ref.float(), atol=0.03125, rtol=0)
+    print(f"PASS qkv_split_norm_rope tokens={tokens}")
+
+
+def run(args) -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+    shapes = [(1, 128), (8, 128)]
+    if args.mode == "full":
+        shapes += [(2, 4, 128)]
+    for shape in shapes:
+        _check_q_norm(ops, shape)
+        _check_k_norm(ops, shape)
+    for tokens in ([4, 64] if args.mode == "full" else [4]):
+        _check_qkv_split(ops, tokens)
+    print(f"PASS flashrt-vla-video {args.backend} mode={args.mode}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    try:
+        run(args)
+    except Exception:
+        import traceback
+        traceback.print_exc()
+        return 1
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(
+            json.dumps({"passed": 1, "total": 1, "backend": args.backend}) + "\n"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-vocab-ce-train/SYNC.md b/flashrt-vocab-ce-train/SYNC.md
new file mode 100644
index 0000000..dad7e19
--- /dev/null
+++ b/flashrt-vocab-ce-train/SYNC.md
@@ -0,0 +1,31 @@
+# Source Sync
+
+- Package: `flashrt-vocab-ce-train` (Vocabulary/classification head cross-entropy training kernels).
+- Upstream FlashRT source: `../official/FlashRT`
+- Upstream revision: pending confirmation; packaged source is maintained in the
+  flashrt-project FlashRT-HF-kernels repository
+  (https://github.com/flashrt-project/FlashRT-HF-kernels).
+
+Copied source files:
+
+- `csrc/README.md`
+- `csrc/vocab_ce_train.cu`
+- `csrc/vocab_ce_train.cuh`
+
+Local packaging edits:
+
+- Added Tensor-facing PyTorch custom ops in `torch-ext/torch_binding.cpp`.
+- Added Python wrappers and fake registrations in `torch-ext/flashrt_vocab_ce_train`.
+- Kept public APIs Tensor-facing; no raw pointer or stream arguments.
+- Includes rewritten to be package-local; serving-runtime dependencies removed.
+- CUDA launchers kept graph-safe: no dynamic allocation inside hot kernels.
+
+Architecture assumptions:
+
+- CUDA 12.8+ / 13.0+ (CUDA 13.2 validated on NVIDIA Thor, sm_110a).
+- NVIDIA Blackwell-family targets; Thor sm_110a validated on real hardware.
+
+Runtime constraints:
+
+- Inputs and outputs are `torch.Tensor`; shapes and dtypes are validated in the binding.
+- Benchmarks cap CUDA memory at 30 GB per process via `set_per_process_memory_fraction`.
diff --git a/flashrt-vocab-ce-train/benchmarks/benchmark.py b/flashrt-vocab-ce-train/benchmarks/benchmark.py
new file mode 100644
index 0000000..7430bc2
--- /dev/null
+++ b/flashrt-vocab-ce-train/benchmarks/benchmark.py
@@ -0,0 +1,101 @@
+#!/usr/bin/env python3
+"""Fused linear-CE streaming-kernel benchmark (kernel vs reference vs torch.compile)."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from pathlib import Path
+
+import torch
+
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
+def elapsed_us(fn, warmup: int = 20, repeats: int = 100) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(repeats):
+        fn()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) * 1000.0 / repeats
+
+
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        root = Path(__file__).resolve().parents[1]
+        sys.path.insert(0, str(root / "torch-ext"))
+        return importlib.import_module("flashrt_vocab_ce_train")
+    if artifact:
+        sys.path.insert(0, artifact)
+    return importlib.import_module("flashrt_vocab_ce_train")
+
+
+def run_case(ops, label: str, n: int, h: int, v: int) -> dict:
+    torch.manual_seed(0)
+    x = torch.randn(n, h, device="cuda", dtype=torch.bfloat16)
+    w = torch.randn(v, h, device="cuda", dtype=torch.float32) * 0.05
+    labels = torch.randint(0, v, (n,), device="cuda", dtype=torch.int64)
+    labels[0] = -100
+
+    def kernel():
+        return ops.vocab_ce(x, w, labels, 0.01)
+
+    def reference():
+        return ops.reference_vocab_ce(x, w, labels, 0.01)
+
+    kernel_us = elapsed_us(kernel)
+    ref_us = elapsed_us(reference)
+
+    compiled = torch.compile(reference, mode="reduce-overhead")
+    compiled()
+    compile_us = elapsed_us(lambda: compiled())
+
+    return {
+        "label": label, "N": n, "H": h, "V": v,
+        "kernel_us": kernel_us, "reference_us": ref_us,
+        "compile_us": compile_us,
+        "vs_reference": ref_us / kernel_us,
+        "vs_compile": compile_us / kernel_us,
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="installed")
+    parser.add_argument("--artifact")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
+    args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
+    ops = load_ops(args.backend, args.artifact)
+    print("label,N,H,V,kernel_us,reference_us,compile_us,vs_reference,vs_compile")
+    for label, n, h, v in [
+        ("small_n", 16, 2048, 65536),
+        ("mid_n", 64, 2048, 257152),
+        ("max_n", 128, 2048, 257152),
+    ]:
+        r = run_case(ops, label, n, h, v)
+        print(
+            f"{r['label']},{r['N']},{r['H']},{r['V']},{r['kernel_us']:.3f},"
+            f"{r['reference_us']:.3f},{r['compile_us']:.3f},"
+            f"{r['vs_reference']:.2f}x,{r['vs_compile']:.2f}x"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/flashrt-vocab-ce-train/tests/test_flashrt_vocab_ce_train.py b/flashrt-vocab-ce-train/tests/test_flashrt_vocab_ce_train.py
old mode 100644
new mode 100755
index 14824fe..6cca2ba
--- a/flashrt-vocab-ce-train/tests/test_flashrt_vocab_ce_train.py
+++ b/flashrt-vocab-ce-train/tests/test_flashrt_vocab_ce_train.py
@@ -1,26 +1,90 @@
 #!/usr/bin/env python3
 from __future__ import annotations
-import argparse, importlib, sys, torch, torch.nn.functional as F
+import argparse, importlib, importlib.util, os, sys, types
+from pathlib import Path
+import torch, torch.nn.functional as F
 from torch.utils.checkpoint import checkpoint
 
-def load_ops(artifact=None):
-    if artifact: sys.path.insert(0, artifact)
-    try: return importlib.import_module("flashrt_vocab_ce_train")
-    finally:
-        if artifact: sys.path.remove(artifact)
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "flashrt-vocab-ce-train"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
+)
+
+
+def load_ops(backend, artifact=None):
+    if backend == "installed":
+        if artifact:
+            sys.path.insert(0, artifact)
+        try:
+            return importlib.import_module("flashrt_vocab_ce_train")
+        finally:
+            if artifact:
+                sys.path.remove(artifact)
+    return load_source_ops()
+
+
+def load_source_ops():
+    from torch.utils.cpp_extension import load
+
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(
+            f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}"
+        )
+    major, minor = torch.cuda.get_device_capability()
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", f"{major}.{minor}a")
+    namespace = "flashrt_vocab_ce_train_source_test"
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            str(PACKAGE / "csrc" / "vocab_ce_train.cu"),
+        ],
+        extra_include_paths=[str(PACKAGE / "csrc"), str(REGISTRATION_INCLUDE)],
+        extra_cflags=["-O3", "-DCUDA_KERNEL"],
+        extra_cuda_cflags=["-O3", "--expt-relaxed-constexpr", "-DCUDA_KERNEL"],
+        verbose=False,
+    )
+    # Exercise the shipped Python wrapper, including its custom-op fake
+    # registration, against the source-built low-level namespace.
+    package_name = "flashrt_vocab_ce_train_source"
+    ops_module = types.ModuleType(f"{package_name}._ops")
+    ops_module.ops = getattr(torch.ops, namespace)
+    ops_module.add_op_namespace_prefix = lambda name: f"{namespace}::{name}"
+    sys.modules[ops_module.__name__] = ops_module
+
+    package_dir = PACKAGE / "torch-ext" / "flashrt_vocab_ce_train"
+    spec = importlib.util.spec_from_file_location(
+        package_name,
+        package_dir / "__init__.py",
+        submodule_search_locations=[str(package_dir)],
+    )
+    if spec is None or spec.loader is None:
+        raise RuntimeError("failed to load packaged flashrt-vocab-ce-train wrapper")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[package_name] = module
+    spec.loader.exec_module(module)
+    return module
+
 
 def run(ops, mode):
     torch.manual_seed(11); count=0
     shapes=[(7,16,31),(23,32,257)] if mode=="full" else [(7,16,31)]
     for n,h,v in shapes:
         x=torch.randn(n,h,device="cuda",dtype=torch.float64,requires_grad=True); w=torch.randn(v,h,device="cuda",dtype=torch.float64,requires_grad=True); labels=torch.randint(0,v,(n,),device="cuda"); labels[0]=-100
-        torch.autograd.gradcheck(lambda a,b: ops.vocab_ce_loss(a,b,labels,0.01), (x,w), eps=1e-6, atol=1e-4, rtol=1e-3)
-        got=ops.vocab_ce_loss(x.float(),w.float(),labels,0.01); logits=x.float()@w.float().t(); valid=labels!=-100; nv=valid.sum().clamp(min=1)
+        torch.autograd.gradcheck(lambda a,b: ops.vocab_ce(a,b,labels,0.01), (x,w), eps=1e-6, atol=1e-4, rtol=1e-3)
+        got=ops.vocab_ce(x.float(),w.float(),labels,0.01); logits=x.float()@w.float().t(); valid=labels!=-100; nv=valid.sum().clamp(min=1)
         ref=F.cross_entropy(logits,labels,ignore_index=-100,reduction="sum")/nv + 0.01*(torch.logsumexp(logits,-1).square()*valid).sum()/nv
         torch.testing.assert_close(got,ref)
-        with torch.autograd.set_detect_anomaly(True): checkpoint(lambda a,b: ops.vocab_ce_loss(a,b,labels,0.0), x.float(),w.float(), use_reentrant=False).backward()
-        torch.compile(lambda a,b: ops.vocab_ce_loss(a,b,labels,0.0), fullgraph=False)(x.float().detach(),w.float().detach()); count+=1
+        with torch.autograd.set_detect_anomaly(True): checkpoint(lambda a,b: ops.vocab_ce(a,b,labels,0.0), x.float(),w.float(), use_reentrant=False).backward()
+        torch.compile(lambda a,b: ops.vocab_ce(a,b,labels,0.0), fullgraph=False)(x.float().detach(),w.float().detach()); count+=1
     print(f"flashrt-vocab-ce-train {mode}: passed {count}/{count}")
 if __name__=="__main__":
-    p=argparse.ArgumentParser(); p.add_argument("--backend",choices=["installed"],default="installed"); p.add_argument("--artifact"); p.add_argument("--mode",choices=["smoke","full"],default="smoke")
-    a=p.parse_args(); run(load_ops(a.artifact),a.mode)
+    p=argparse.ArgumentParser(); p.add_argument("--backend",choices=["source","installed"],default="installed"); p.add_argument("--artifact"); p.add_argument("--mode",choices=["smoke","full"],default="smoke")
+    a=p.parse_args(); run(load_ops(a.backend, a.artifact), a.mode)
diff --git a/fp4-fused-ops/benchmarks/benchmark.py b/fp4-fused-ops/benchmarks/benchmark.py
index 93186e1..538cdf6 100644
--- a/fp4-fused-ops/benchmarks/benchmark.py
+++ b/fp4-fused-ops/benchmarks/benchmark.py
@@ -14,6 +14,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 TEST_FILE = ROOT / "fp4-fused-ops" / "tests" / "test_fp4_fused_ops.py"
 
@@ -93,7 +103,7 @@ def bench_case(helpers, ops, native, rows: int, dim: int, warmup: int, iters: in
         iters,
     )
     stream = torch.cuda.current_stream().cuda_stream
-    native_f3_us = measure(
+    native_f3_us = float("nan") if native is None else measure(
         lambda: native.residual_add_rms_norm_fp4_sfa_v2_fp16(
             residual_v2.copy_(residual).data_ptr(), x.data_ptr(),
             packed.data_ptr(), sfa.data_ptr(), rows, dim, stream
@@ -106,7 +116,7 @@ def bench_case(helpers, ops, native, rows: int, dim: int, warmup: int, iters: in
         warmup,
         iters,
     )
-    native_graph_f3_us = measure_graph(
+    native_graph_f3_us = float("nan") if native is None else measure_graph(
         lambda: native.residual_add_rms_norm_fp4_sfa_v2_fp16(
             residual_v2.copy_(residual).data_ptr(), x.data_ptr(),
             packed.data_ptr(), sfa.data_ptr(), rows, dim,
@@ -138,7 +148,7 @@ def bench_case(helpers, ops, native, rows: int, dim: int, warmup: int, iters: in
     packed_v2, sfa_v2 = ops.alloc(rows, dim)
     ref_us = measure(lambda: ops.silu_mul_fp4_sfa_fp16(merged, packed_v1, sfa_v1), warmup, iters)
     f4_us = measure(lambda: ops.silu_mul_fp4_sfa_v2_fp16(merged, packed_v2, sfa_v2), warmup, iters)
-    native_f4_us = measure(
+    native_f4_us = float("nan") if native is None else measure(
         lambda: native.gate_geglu_fp4_sfa_v2_fp16(
             merged.data_ptr(), packed_v2.data_ptr(), sfa_v2.data_ptr(),
             rows, dim, stream
@@ -149,7 +159,7 @@ def bench_case(helpers, ops, native, rows: int, dim: int, warmup: int, iters: in
     graph_f4_us = measure_graph(
         lambda: ops.silu_mul_fp4_sfa_v2_fp16(merged, packed_v2, sfa_v2), warmup, iters
     )
-    native_graph_f4_us = measure_graph(
+    native_graph_f4_us = float("nan") if native is None else measure_graph(
         lambda: native.gate_geglu_fp4_sfa_v2_fp16(
             merged.data_ptr(), packed_v2.data_ptr(), sfa_v2.data_ptr(), rows, dim,
             torch.cuda.current_stream().cuda_stream
@@ -177,7 +187,7 @@ def bench_case(helpers, ops, native, rows: int, dim: int, warmup: int, iters: in
 
     inv_s = (torch.rand((dim,), device="cuda") * 0.25 + 0.875).to(torch.float16).contiguous()
     awq_us = measure(lambda: ops.silu_mul_mul_fp4_sfa_v2_fp16(merged, inv_s, packed_v2, sfa_v2), warmup, iters)
-    native_awq_us = measure(
+    native_awq_us = float("nan") if native is None else measure(
         lambda: native.gate_geglu_mul_fp4_sfa_v2_fp16(
             merged.data_ptr(), inv_s.data_ptr(), packed_v2.data_ptr(),
             sfa_v2.data_ptr(), rows, dim, stream
@@ -190,7 +200,7 @@ def bench_case(helpers, ops, native, rows: int, dim: int, warmup: int, iters: in
         warmup,
         iters,
     )
-    native_graph_awq_us = measure_graph(
+    native_graph_awq_us = float("nan") if native is None else measure_graph(
         lambda: native.gate_geglu_mul_fp4_sfa_v2_fp16(
             merged.data_ptr(), inv_s.data_ptr(), packed_v2.data_ptr(),
             sfa_v2.data_ptr(), rows, dim, torch.cuda.current_stream().cuda_stream
@@ -226,7 +236,7 @@ def bench_case(helpers, ops, native, rows: int, dim: int, warmup: int, iters: in
         warmup,
         iters,
     )
-    native_two_us = measure(
+    native_two_us = float("nan") if native is None else measure(
         lambda: native.geglu_two_fp4_to_fp4(
             gate_packed.data_ptr(), gate_sfa.data_ptr(), up_packed.data_ptr(),
             up_sfa.data_ptr(), out_packed.data_ptr(), out_sfa.data_ptr(),
@@ -240,7 +250,7 @@ def bench_case(helpers, ops, native, rows: int, dim: int, warmup: int, iters: in
         warmup,
         iters,
     )
-    native_graph_two_us = measure_graph(
+    native_graph_two_us = float("nan") if native is None else measure_graph(
         lambda: native.geglu_two_fp4_to_fp4(
             gate_packed.data_ptr(), gate_sfa.data_ptr(), up_packed.data_ptr(),
             up_sfa.data_ptr(), out_packed.data_ptr(), out_sfa.data_ptr(), rows, dim,
@@ -254,7 +264,7 @@ def bench_case(helpers, ops, native, rows: int, dim: int, warmup: int, iters: in
         warmup,
         iters,
     )
-    native_two_mul_us = measure(
+    native_two_mul_us = float("nan") if native is None else measure(
         lambda: native.geglu_two_mul_fp4_to_fp4(
             gate_packed.data_ptr(), gate_sfa.data_ptr(), up_packed.data_ptr(),
             up_sfa.data_ptr(), inv_s.data_ptr(), out_packed.data_ptr(),
@@ -268,7 +278,7 @@ def bench_case(helpers, ops, native, rows: int, dim: int, warmup: int, iters: in
         warmup,
         iters,
     )
-    native_graph_two_mul_us = measure_graph(
+    native_graph_two_mul_us = float("nan") if native is None else measure_graph(
         lambda: native.geglu_two_mul_fp4_to_fp4(
             gate_packed.data_ptr(), gate_sfa.data_ptr(), up_packed.data_ptr(),
             up_sfa.data_ptr(), inv_s.data_ptr(), out_packed.data_ptr(),
@@ -324,7 +334,9 @@ def main() -> int:
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iterations", type=int, default=100)
     parser.add_argument("--json-out", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required")
@@ -337,11 +349,15 @@ def main() -> int:
     if not hasattr(ops, "alloc"):
         ops.alloc = lambda rows, dim: helpers.alloc_fp4(ops, rows, dim)
     native_root = Path(os.environ.get("FLASHRT_NATIVE_ROOT", str(ROOT.parent / "official" / "FlashRT")))
-    sys.path.insert(0, str(native_root))
-    try:
-        import flash_rt.flash_rt_fp4 as native
-    finally:
-        sys.path.pop(0)
+    native = None
+    if native_root.is_dir():
+        sys.path.insert(0, str(native_root))
+        try:
+            import flash_rt.flash_rt_fp4 as native
+        except Exception:
+            native = None
+        finally:
+            sys.path.pop(0)
 
     shapes = [(1, 1024), (10, 2048)] if args.mode == "smoke" else [(1, 1024), (10, 2048), (64, 2048), (128, 4096)]
     if args.mode == "thor-models":
diff --git a/fp4-gemm/benchmarks/benchmark.py b/fp4-gemm/benchmarks/benchmark.py
index f709637..18ea369 100644
--- a/fp4-gemm/benchmarks/benchmark.py
+++ b/fp4-gemm/benchmarks/benchmark.py
@@ -14,6 +14,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 TEST_FILE = ROOT / "fp4-gemm" / "tests" / "test_fp4_gemm.py"
 
@@ -92,31 +102,33 @@ def torch_ref():
             iters,
         )
         native_variant = variant
-        if native_variant < 0:
-            native_variant = helpers.select_sm110_variant(shape)
-        native_function = (
-            native.fp4_w4a16_gemm_sm120_bf16out
-            if native_variant == 0
-            else native.fp4_w4a16_gemm_sm120_bf16out_widen
-            if native_variant == 1
-            else native.fp4_w4a16_gemm_sm120_bf16out_pingpong
-        )
-        native_us = measure(
-            lambda: native_function(
-                a_packed.data_ptr(),
-                b_packed.data_ptr(),
-                out.data_ptr(),
-                m,
-                n,
-                k,
-                sfa.data_ptr(),
-                sfb.data_ptr(),
-                1.0,
-                stream,
-            ),
-            warmup,
-            iters,
-        )
+        native_us = float("nan")
+        if native is not None:
+            if native_variant < 0:
+                native_variant = helpers.select_sm110_variant(shape)
+            native_function = (
+                native.fp4_w4a16_gemm_sm120_bf16out
+                if native_variant == 0
+                else native.fp4_w4a16_gemm_sm120_bf16out_widen
+                if native_variant == 1
+                else native.fp4_w4a16_gemm_sm120_bf16out_pingpong
+            )
+            native_us = measure(
+                lambda: native_function(
+                    a_packed.data_ptr(),
+                    b_packed.data_ptr(),
+                    out.data_ptr(),
+                    m,
+                    n,
+                    k,
+                    sfa.data_ptr(),
+                    sfb.data_ptr(),
+                    1.0,
+                    stream,
+                ),
+                warmup,
+                iters,
+            )
         results.append(
             BenchResult(
                 shape=name,
@@ -169,7 +181,7 @@ def native_direct():
     torch.cuda.synchronize()
     direct_us = measure(direct, warmup, iters)
     compat_us = measure(compat, warmup, iters)
-    native_us = measure(native_direct, warmup, iters)
+    native_us = float("nan") if native is None else measure(native_direct, warmup, iters)
     return {
         "M": 1,
         "K": k,
@@ -195,17 +207,23 @@ def main() -> int:
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iterations", type=int, default=100)
     parser.add_argument("--json-out", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     helpers = load_helpers()
     native_root = Path(
         os.environ.get("FLASHRT_NATIVE_ROOT", str(ROOT.parent / "official" / "FlashRT"))
     )
-    sys.path.insert(0, str(native_root))
-    try:
-        import flash_rt.flash_rt_kernels as native
-    finally:
-        sys.path.pop(0)
+    native = None
+    if native_root.is_dir():
+        sys.path.insert(0, str(native_root))
+        try:
+            import flash_rt.flash_rt_kernels as native
+        except Exception:
+            native = None
+        finally:
+            sys.path.pop(0)
     ops = (
         helpers.load_source_ops()
         if args.backend == "source"
diff --git a/fp8-cross-attention-blackwell/benchmarks/benchmark.py b/fp8-cross-attention-blackwell/benchmarks/benchmark.py
index 30d9160..be35507 100644
--- a/fp8-cross-attention-blackwell/benchmarks/benchmark.py
+++ b/fp8-cross-attention-blackwell/benchmarks/benchmark.py
@@ -11,6 +11,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def elapsed_us(fn, warmup, iterations):
     for _ in range(warmup):
         fn()
@@ -40,7 +50,9 @@ def main():
     parser.add_argument("--hkv", type=int, default=4)
     parser.add_argument("--warmup", type=int, default=10)
     parser.add_argument("--iterations", type=int, default=50)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     tests = Path(__file__).resolve().parents[1] / "tests"
     sys.path.insert(0, str(tests))
diff --git a/fp8-gemm/benchmarks/benchmark.py b/fp8-gemm/benchmarks/benchmark.py
index ccd1a45..39eccb9 100644
--- a/fp8-gemm/benchmarks/benchmark.py
+++ b/fp8-gemm/benchmarks/benchmark.py
@@ -16,6 +16,16 @@
 import torch
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 PACKAGE = ROOT / "fp8-gemm"
 REGISTRATION_INCLUDE = (
@@ -545,7 +555,9 @@ def main() -> None:
     parser.add_argument("--iterations", type=int, default=100)
     parser.add_argument("--compile-ref", action="store_true")
     parser.add_argument("--json-out", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
diff --git a/fp8-kv-attention/benchmarks/benchmark.py b/fp8-kv-attention/benchmarks/benchmark.py
index 21ff781..e44abeb 100644
--- a/fp8-kv-attention/benchmarks/benchmark.py
+++ b/fp8-kv-attention/benchmarks/benchmark.py
@@ -14,6 +14,16 @@
 import torch.nn.functional as F
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 PACKAGE = Path(__file__).resolve().parents[1]
 ROOT = PACKAGE.parent
 REGISTRATION = (
@@ -199,7 +209,9 @@ def main() -> int:
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iters", type=int, default=100)
     parser.add_argument("--json-out")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     wrapper = load_wrapper(args.backend, args.artifact)
     native = build_native()
diff --git a/fp8-prefill-attention-blackwell/benchmarks/benchmark.py b/fp8-prefill-attention-blackwell/benchmarks/benchmark.py
new file mode 100644
index 0000000..0e57132
--- /dev/null
+++ b/fp8-prefill-attention-blackwell/benchmarks/benchmark.py
@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+"""FP8 causal GQA prefill attention benchmark (kernel vs eager SDPA vs torch.compile)."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import math
+import sys
+from pathlib import Path
+
+import torch
+import torch.nn.functional as F
+
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
+def elapsed_us(fn, warmup: int = 20, repeats: int = 50) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(repeats):
+        fn()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) * 1000.0 / repeats
+
+
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        root = Path(__file__).resolve().parents[1]
+        sys.path.insert(0, str(root / "torch-ext"))
+        return importlib.import_module("fp8_prefill_attention_blackwell")
+    if artifact:
+        sys.path.insert(0, artifact)
+    return importlib.import_module("fp8_prefill_attention_blackwell")
+
+
+def run_case(ops, label: str, s: int) -> dict:
+    q = (torch.randn(s, 32, 128, device="cuda") * 0.5).to(torch.float8_e4m3fn)
+    k = (torch.randn(s, 8, 128, device="cuda") * 0.5).to(torch.float8_e4m3fn)
+    v = (torch.randn(s, 8, 128, device="cuda") * 0.5).to(torch.float8_e4m3fn)
+    scale = 1.0 / math.sqrt(128)
+
+    def kernel():
+        return ops.fp8_causal_gqa_attention_bf16(q, k, v, softmax_scale=scale)
+
+    qb = q.float().to(torch.bfloat16).transpose(0, 1)
+    kb = k.float().to(torch.bfloat16).transpose(0, 1)
+    vb = v.float().to(torch.bfloat16).transpose(0, 1)
+    # GQA: expand KV heads 8 -> 32 to match Hq for SDPA.
+    kb_e = kb.repeat_interleave(4, dim=0)
+    vb_e = vb.repeat_interleave(4, dim=0)
+
+    def eager():
+        return F.scaled_dot_product_attention(qb, kb_e, vb_e, is_causal=True)
+
+    eager_us = elapsed_us(eager)
+    compiled = torch.compile(eager, mode="reduce-overhead")
+    compiled()
+    compile_us = elapsed_us(lambda: compiled())
+    kernel_us = elapsed_us(kernel)
+    return {
+        "label": label, "S": s, "Hq": 32, "Hkv": 8, "D": 128,
+        "kernel_us": kernel_us, "eager_us": eager_us, "compile_us": compile_us,
+        "vs_eager": eager_us / kernel_us,
+        "vs_compile": compile_us / kernel_us,
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="installed")
+    parser.add_argument("--artifact")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
+    args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
+    ops = load_ops(args.backend, args.artifact)
+    print("label,S,Hq,Hkv,D,kernel_us,eager_us,compile_us,vs_eager,vs_compile")
+    for label, s in [
+        ("prefill_1k", 1024),
+        ("prefill_2k", 2048),
+        ("prefill_4k", 4096),
+    ]:
+        r = run_case(ops, label, s)
+        print(
+            f"{r['label']},{r['S']},{r['Hq']},{r['Hkv']},{r['D']},"
+            f"{r['kernel_us']:.3f},{r['eager_us']:.3f},{r['compile_us']:.3f},"
+            f"{r['vs_eager']:.2f}x,{r['vs_compile']:.2f}x"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/fused-mlp-megakernels-blackwell/VALIDATION.md b/fused-mlp-megakernels-blackwell/VALIDATION.md
index 1308e28..cb34f25 100644
--- a/fused-mlp-megakernels-blackwell/VALIDATION.md
+++ b/fused-mlp-megakernels-blackwell/VALIDATION.md
@@ -17,3 +17,42 @@ through `768x16384x2048`, including M-tail rows 127 and 129, plus
 recorded `max=0.00146484`, `p99=0.00012207`, `mean=0.00001121`, and
 `cosine=0.99999988`. Installed-artifact execution remains a separate release
 gate.
+
+## Installed artifact on NVIDIA Thor (SM110) via portable SIMT fallback
+
+The SM100-family CUTLASS 4.0 megakernel uses tcgen05/TMA descriptor paths that
+assert at runtime on `sm_110a` (Thor). `portable_geglu_simt.cu` provides a pure
+SIMT FMA reference of the same fusion:
+
+- `gate_scratch[m,n] = fp16( gelu_tanh( X@W_gate^T ) )`
+- `hidden[m,n]     = fp16( gate_scratch[m,n] * ( X@W_up^T ) )`
+
+Dispatch in `torch_binding.cpp`: `sm_100`/`sm_103` keep the CUTLASS megakernel;
+`sm_110a` and the `FLASHRT_FORCE_SIMT` override route to the SIMT reference
+(mirrors the fp8-gemm / world-model-conv / grouped-moe-gemm fallbacks).
+
+Command:
+
+```bash
+python fused-mlp-megakernels-blackwell/tests/test_fused_mlp_megakernels_blackwell.py \
+    --backend installed --artifact <installed-dir> --mode full
+```
+
+Result: 7/7 numeric rows passed on Thor, plus `torch.compile(fullgraph=True)`
+and prewarmed CUDA Graph replay. Rows and metrics (vs the PyTorch FP16 formula
+`gelu_tanh(x@Wg.t()) * (x@Wu.t())`):
+
+| M | N | K | max_abs | p99 | mean | cosine |
+|---|---|---|---|---|---|---|
+| 128 | 128 | 128 | 0.0000305 | 0.0000076 | 0.0000007 | 0.99999988 |
+| 64 | 256 | 256 | 0.0000610 | 0.0000153 | 0.0000014 | 0.99999976 |
+| 127 | 256 | 256 | 0.0000610 | 0.0000153 | 0.0000014 | 0.99999982 |
+| 129 | 256 | 256 | 0.0000610 | 0.0000153 | 0.0000014 | 0.99999982 |
+| 256 | 1024 | 1024 | 0.0004883 | 0.0000610 | 0.0000056 | 0.99999988 |
+| 768 | 2048 | 2048 | 0.0009766 | 0.0001221 | 0.0000113 | 0.99999994 |
+| 768 | 16384 | 2048 | 0.0014648 | 0.0001221 | 0.0000112 | 0.99999988 |
+
+The SIMT path is a correctness/compatibility fallback for non-SM100-family
+devices; `sm_100`/`sm_103` continue to use the fused CUTLASS megakernel.
+Performance qualification against PyTorch eager / `torch.compile` remains an
+open release gate (see item 5 above).
diff --git a/fused-mlp-megakernels-blackwell/benchmarks/RESULTS.md b/fused-mlp-megakernels-blackwell/benchmarks/RESULTS.md
index 902ef07..663072c 100644
--- a/fused-mlp-megakernels-blackwell/benchmarks/RESULTS.md
+++ b/fused-mlp-megakernels-blackwell/benchmarks/RESULTS.md
@@ -3,6 +3,8 @@
 Source-extension triage on NVIDIA Thor SM110, CUDA 13, PyTorch 2.11.0+cu130.
 The release table will be regenerated from the installed Hub artifact.
 
+## SM100-family megakernel (production path)
+
 | Shape | Eager us | Compile us | Wrapper us | Raw native us | vs compile |
 |---|---:|---:|---:|---:|---:|
 | `M768 N16384 K2048` | 1465.035 | 1296.318 | 1005.803 | 998.419 | 1.29x |
@@ -10,3 +12,22 @@ The release table will be regenerated from the installed Hub artifact.
 Wrapper overhead versus the raw native entry is 0.74%. The production row
 passed with `max=0.00146484`, `p99=0.00012207`, `mean=0.00001121`, and
 `cosine=0.99999988`.
+
+## Portable SIMT fallback on NVIDIA Thor (SM110 installed artifact)
+
+The SM100-family CUTLASS megakernel asserts at runtime on `sm_110a`, so the
+installed artifact routes SM110 to the portable SIMT fallback
+(`portable_geglu_simt.cu`). Measured against the installed
+`<installed-dir>/fused-mlp-megakernels-blackwell` artifact on Thor:
+
+| Shape | Eager us | Compile us | Wrapper us | Raw native us | vs eager |
+|---|---:|---:|---:|---:|---:|
+| `M768 N16384 K2048` | 1633.251 | 1358.933 | 85558.905 | 85538.818 | 0.019x |
+
+Block-tiled SIMT (32x32 tile, shared-memory K chunks) runs ~22x faster than the
+initial one-thread-per-output reference (1.89s -> 85.6ms). It remains a
+correctness/compatibility path, not a production kernel: on `sm_110a` it is the
+only available implementation (the CUTLASS megakernel cannot launch), and it is
+~52x slower than eager. `sm_100`/`sm_103` continue to use the fused CUTLASS
+megakernel; the gap is documented and the fallback is not used where a native
+path exists.
diff --git a/fused-mlp-megakernels-blackwell/benchmarks/benchmark.py b/fused-mlp-megakernels-blackwell/benchmarks/benchmark.py
index 7fbb740..ad3418d 100644
--- a/fused-mlp-megakernels-blackwell/benchmarks/benchmark.py
+++ b/fused-mlp-megakernels-blackwell/benchmarks/benchmark.py
@@ -12,6 +12,16 @@
 import torch.nn.functional as F
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def elapsed_us(fn, warmup, iterations):
     for _ in range(warmup):
         fn()
@@ -35,7 +45,9 @@ def main():
     parser.add_argument("--k", type=int, default=2048)
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iterations", type=int, default=100)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if args.backend == "source":
         tests = Path(__file__).resolve().parents[1] / "tests"
diff --git a/fused-mlp-megakernels-blackwell/build.toml b/fused-mlp-megakernels-blackwell/build.toml
index 6554dec..f8a85a2 100644
--- a/fused-mlp-megakernels-blackwell/build.toml
+++ b/fused-mlp-megakernels-blackwell/build.toml
@@ -30,4 +30,6 @@ src = [
   "csrc/mega/cutlass/epilogue/collective/collective_epilogue.hpp",
   "csrc/mega/cutlass/epilogue/collective/sm100_epilogue_tma_warpspecialized.hpp",
   "csrc/mega/cutlass/gemm/collective/sm100_mma_warpspecialized.hpp",
+  "csrc/mega/portable_geglu_simt.cu",
+  "csrc/mega/portable_geglu_simt.cuh",
 ]
diff --git a/fused-mlp-megakernels-blackwell/csrc/mega/portable_geglu_simt.cu b/fused-mlp-megakernels-blackwell/csrc/mega/portable_geglu_simt.cu
new file mode 100644
index 0000000..a070707
--- /dev/null
+++ b/fused-mlp-megakernels-blackwell/csrc/mega/portable_geglu_simt.cu
@@ -0,0 +1,139 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Portable SIMT implementation of the fused FP16 GeGLU megakernel for sm_110a
+// (Thor). Block-tiled over shared memory: a 32x32 output tile reuses each X
+// row / weight column across the whole tile, so global traffic drops ~32x
+// versus a one-thread-per-output reference. sm_100/sm_103 keep the CUTLASS
+// megakernel path; this is a compatibility path, not a performance kernel.
+
+#include "portable_geglu_simt.cuh"
+
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+
+namespace flashrt {
+namespace megakernel {
+
+namespace {
+
+constexpr int THREADS = 256;
+constexpr int BM = 32;
+constexpr int BN = 32;
+constexpr int BK = 16;               // per-chunk K depth (fp16 tile fits SMEM)
+constexpr int ELEMS_PER_THREAD = (BM * BN) / THREADS;  // 4
+
+__device__ __forceinline__ float gelu_tanh_f32(float x) {
+  // Matches cutlass::epilogue::thread::GELU_taylor (the tanh approximation).
+  const float t = 0.5f * x;
+  return t * (1.0f + tanhf(0.7978845608028654f *
+                           (x + 0.044715f * x * x * x)));
+}
+
+__global__ void geglu_fused_tiled_kernel(
+    const __half* __restrict__ X,      // (M, K) row-major
+    const __half* __restrict__ W_gate, // (N, K) row-major
+    const __half* __restrict__ W_up,   // (N, K) row-major
+    __half* __restrict__ gate_scratch, // (M, N) row-major
+    __half* __restrict__ hidden,       // (M, N) row-major
+    int M, int N, int K) {
+  __shared__ __half xs[BM][BK];
+  __shared__ __half wgs[BN][BK];
+  __shared__ __half wus[BN][BK];
+
+  const int tile_m = blockIdx.y * BM;
+  const int tile_n = blockIdx.x * BN;
+  const int nk = (K + BK - 1) / BK;
+
+  float gate_acc[ELEMS_PER_THREAD];
+  float up_acc[ELEMS_PER_THREAD];
+  int rows[ELEMS_PER_THREAD];
+  int cols[ELEMS_PER_THREAD];
+  #pragma unroll
+  for (int i = 0; i < ELEMS_PER_THREAD; ++i) {
+    const int lin = threadIdx.x + i * THREADS;
+    const int r = lin >> 5;           // / BN
+    const int c = lin & 31;           // % BN
+    rows[i] = r;
+    cols[i] = c;
+    gate_acc[i] = 0.0f;
+    up_acc[i] = 0.0f;
+  }
+
+  for (int kb = 0; kb < nk; ++kb) {
+    const int k0 = kb * BK;
+    // cooperative tile load with edge masking
+    const int tid = threadIdx.x;
+    const int tiles_per_chunk = BM + 2 * BN;  // X tile + gate/up weight tiles
+    const int load_total = tiles_per_chunk * BK;
+    for (int e = tid; e < load_total; e += THREADS) {
+      const int row_sel = e / BK;
+      const int j = e % BK;
+      const int kk = k0 + j;
+      if (row_sel < BM) {
+        const int m = tile_m + row_sel;
+        xs[row_sel][j] = (m < M && kk < K) ? X[(size_t)m * K + kk] : __float2half(0.0f);
+      } else if (row_sel < BM + BN) {
+        const int c = row_sel - BM;
+        const int n = tile_n + c;
+        const bool ok = (n < N && kk < K);
+        wgs[c][j] = ok ? W_gate[(size_t)n * K + kk] : __float2half(0.0f);
+      } else {
+        const int c = row_sel - BM - BN;
+        const int n = tile_n + c;
+        const bool ok = (n < N && kk < K);
+        wus[c][j] = ok ? W_up[(size_t)n * K + kk] : __float2half(0.0f);
+      }
+    }
+    __syncthreads();
+    #pragma unroll
+    for (int i = 0; i < ELEMS_PER_THREAD; ++i) {
+      const int r = rows[i];
+      const int c = cols[i];
+      float g = gate_acc[i];
+      float u = up_acc[i];
+      #pragma unroll
+      for (int j = 0; j < BK; ++j) {
+        const float xv = __half2float(xs[r][j]);
+        g += xv * __half2float(wgs[c][j]);
+        u += xv * __half2float(wus[c][j]);
+      }
+      gate_acc[i] = g;
+      up_acc[i] = u;
+    }
+    __syncthreads();
+  }
+
+  #pragma unroll
+  for (int i = 0; i < ELEMS_PER_THREAD; ++i) {
+    const int m = tile_m + rows[i];
+    const int n = tile_n + cols[i];
+    if (m >= M || n >= N) continue;
+    const __half g_h = __float2half(gelu_tanh_f32(gate_acc[i]));
+    const size_t idx = (size_t)m * N + n;
+    gate_scratch[idx] = g_h;
+    hidden[idx] = __float2half(__half2float(g_h) * up_acc[i]);
+  }
+}
+
+}  // namespace
+
+int geglu_fused_fp16_simt(
+    const void* X, const void* W_gate, const void* W_up,
+    void* gate_scratch, void* hidden,
+    int M, int N, int K, cudaStream_t stream) {
+  if (M <= 0 || N <= 0 || K <= 0) return 1;
+  dim3 block(THREADS);
+  dim3 grid((N + BN - 1) / BN, (M + BM - 1) / BM);
+  geglu_fused_tiled_kernel<<<grid, block, 0, stream>>>(
+      reinterpret_cast<const __half*>(X),
+      reinterpret_cast<const __half*>(W_gate),
+      reinterpret_cast<const __half*>(W_up),
+      reinterpret_cast<__half*>(gate_scratch),
+      reinterpret_cast<__half*>(hidden),
+      M, N, K);
+  return (cudaGetLastError() == cudaSuccess) ? 0 : 1;
+}
+
+}  // namespace megakernel
+}  // namespace flashrt
diff --git a/fused-mlp-megakernels-blackwell/csrc/mega/portable_geglu_simt.cuh b/fused-mlp-megakernels-blackwell/csrc/mega/portable_geglu_simt.cuh
new file mode 100644
index 0000000..4e9dcc2
--- /dev/null
+++ b/fused-mlp-megakernels-blackwell/csrc/mega/portable_geglu_simt.cuh
@@ -0,0 +1,30 @@
+// SPDX-License-Identifier: Apache-2.0
+#pragma once
+
+#include <cuda_runtime.h>
+
+// Portable SIMT implementation of the fused FP16 GeGLU megakernel. The
+// SM100-family CUTLASS 4.0 megakernel (flashrt_megakernel_geglu_fp16) uses
+// tcgen05/TMA descriptor paths that assert on sm_110a (Thor); this reference
+// computes the same fusion in pure SIMT FMA so the op stays usable there:
+//
+//   gate[m, n]         = sum_k  X[m,k] * W_gate[n,k]          (fp32 acc)
+//   gate_scratch[m, n] = fp16( gelu_tanh(gate[m, n]) )
+//   hidden[m, n]       = fp16( gate_scratch[m, n] * sum_k X[m,k]*W_up[n,k] )
+//
+// sm_100/sm_103 keep the CUTLASS megakernel path; this is a compatibility
+// path only (matches the fp8-gemm / world-model-conv portable fallbacks).
+
+namespace flashrt {
+namespace megakernel {
+
+// hidden[M,N] fp16, gate_scratch[M,N] fp16 =
+//   gelu_tanh(X@W_gate^T) * (X@W_up^T), with gate_scratch = gelu_tanh(gate).
+// Returns 0 on success, non-zero on invalid dimensions.
+int geglu_fused_fp16_simt(
+    const void* X, const void* W_gate, const void* W_up,
+    void* gate_scratch, void* hidden,
+    int M, int N, int K, cudaStream_t stream);
+
+}  // namespace megakernel
+}  // namespace flashrt
diff --git a/fused-mlp-megakernels-blackwell/torch-ext/torch_binding.cpp b/fused-mlp-megakernels-blackwell/torch-ext/torch_binding.cpp
index f2fa6e7..4e8f4ba 100644
--- a/fused-mlp-megakernels-blackwell/torch-ext/torch_binding.cpp
+++ b/fused-mlp-megakernels-blackwell/torch-ext/torch_binding.cpp
@@ -2,6 +2,8 @@
 #include <torch/all.h>
 #include <torch/library.h>
 
+#include <cstdlib>
+
 #if defined(CUDA_KERNEL)
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
@@ -14,6 +16,7 @@
 #if defined(CUDA_KERNEL)
 extern "C" int flashrt_megakernel_geglu_fp16(
     void*, void*, void*, void*, void*, int, int, int, cudaStream_t);
+#include "portable_geglu_simt.cuh"
 #endif
 
 namespace {
@@ -63,12 +66,23 @@ void fp16_geglu_fused_out(
   TORCH_CHECK(capability == 100 || capability == 103 || capability == 110,
               "fp16_geglu_fused_out requires SM100, SM103, or SM110");
   auto stream = at::cuda::getCurrentCUDAStream(input.get_device()).stream();
-  int rc = flashrt_megakernel_geglu_fp16(
-      input.data_ptr(), gate_weight.data_ptr(), up_weight.data_ptr(),
-      gate_scratch.data_ptr(), output.data_ptr(),
-      static_cast<int>(input.size(0)), static_cast<int>(gate_weight.size(0)),
-      static_cast<int>(input.size(1)), stream);
-  TORCH_CHECK(rc == 0, "FlashRT FP16 GeGLU megakernel failed with status ", rc);
+  const int M = static_cast<int>(input.size(0));
+  const int N = static_cast<int>(gate_weight.size(0));
+  const int K = static_cast<int>(input.size(1));
+  const bool force_simt = std::getenv("FLASHRT_FORCE_SIMT") != nullptr;
+  if (force_simt || (properties->major == 11 && properties->minor == 0)) {
+    // sm_110a (Thor): the SM100 CUTLASS megakernel's tcgen05/TMA descriptor
+    // paths assert at runtime, so route to the portable SIMT reference.
+    int rc = flashrt::megakernel::geglu_fused_fp16_simt(
+        input.data_ptr(), gate_weight.data_ptr(), up_weight.data_ptr(),
+        gate_scratch.data_ptr(), output.data_ptr(), M, N, K, stream);
+    TORCH_CHECK(rc == 0, "FP16 GeGLU SIMT fallback failed with status ", rc);
+  } else {
+    int rc = flashrt_megakernel_geglu_fp16(
+        input.data_ptr(), gate_weight.data_ptr(), up_weight.data_ptr(),
+        gate_scratch.data_ptr(), output.data_ptr(), M, N, K, stream);
+    TORCH_CHECK(rc == 0, "FlashRT FP16 GeGLU megakernel failed with status ", rc);
+  }
 #else
   TORCH_CHECK(false, "CUDA support was not built");
 #endif
diff --git a/grouped-moe-gemm/CARD.md b/grouped-moe-gemm/CARD.md
index ae3bdc4..76d3730 100644
--- a/grouped-moe-gemm/CARD.md
+++ b/grouped-moe-gemm/CARD.md
@@ -22,3 +22,16 @@ weighted unpermute. Those scheduling tensors are intentionally outside this
 compute-only ABI.
 
 No hidden dequantization or eager fallback occurs in this package.
+
+## Architecture support
+
+- `sm_120a`/`sm_121`: native CUTLASS block-scaled MMA tiles (`M16`, `M64`,
+  `64x64` block tile). This is the only dispatch target on SM120.
+- `sm_110a` (Jetson AGX Thor): a portable pure-SIMT reference
+  (`portable_moe_simt.cu`) computes the same grouped FP4 x FP4 -> BF16 GEMM.
+  It is a compatibility path, not a performance kernel.
+
+The dispatcher selects the SIMT reference only on non-SM120 devices where no
+tensor-core backend exists. Set `FLASHRT_FORCE_SIMT=1` to route any device
+through the SIMT reference (used by the correctness test to validate parity
+against the native path).
diff --git a/grouped-moe-gemm/VALIDATION.md b/grouped-moe-gemm/VALIDATION.md
index b04eb66..cae5166 100644
--- a/grouped-moe-gemm/VALIDATION.md
+++ b/grouped-moe-gemm/VALIDATION.md
@@ -9,6 +9,26 @@ are mandatory rows. Performance is compared to the original FlashRT
 entry points and relevant eager/compiled references without counting Python
 packing loops as a kernel speedup baseline.
 
+## SM110 portable SIMT fallback
+
+`grouped_nvfp4_gemm_bf16` also ships a pure-SIMT reference
+(`portable_moe_simt.cu`) compiled for `sm_110a`. The correctness test runs the
+full case matrix a second time with `FLASHRT_FORCE_SIMT=1`, validating the SIMT
+path against the same FP32 reference on any device. Native SM120 tile behavior
+is unchanged.
+
+### SM110 (NVIDIA Thor) real-hardware validation
+
+Validated on an NVIDIA Thor (`sm_110a`) device with PyTorch 2.11+cu130 and
+CUDA 13.2 against the FP32 dequantize-then-matmul reference:
+
+- tile=16 K=64: rel L2 0.0, max rel 0.0
+- tile=16 K=2048: rel L2 1.3e-5, max rel 3.7e-4
+- tile=64 K=2048: rel L2 3.6e-5, max rel 1.5e-3
+
+All within the package gate (`rel_l2 <= 2.5e-3`, `max_rel <= 0.01`,
+`cosine >= 0.999`).
+
 ```bash
 python grouped-moe-gemm/tests/test_grouped_moe_gemm.py --backend source
 python grouped-moe-gemm/tests/test_grouped_moe_gemm.py \
diff --git a/grouped-moe-gemm/benchmarks/benchmark.py b/grouped-moe-gemm/benchmarks/benchmark.py
index e6eee21..ae2a72e 100644
--- a/grouped-moe-gemm/benchmarks/benchmark.py
+++ b/grouped-moe-gemm/benchmarks/benchmark.py
@@ -24,6 +24,16 @@
 }
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def time_us(fn, warmup: int, iters: int) -> float:
     for _ in range(warmup):
         fn()
@@ -71,8 +81,10 @@ def main() -> int:
     parser.add_argument("--artifact")
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iters", type=int, default=100)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     parser.add_argument("--json-out")
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     torch.manual_seed(9102)
     ops = load_ops(args.backend, args.artifact)
diff --git a/grouped-moe-gemm/build.toml b/grouped-moe-gemm/build.toml
index e2b84a8..9e87d8c 100644
--- a/grouped-moe-gemm/build.toml
+++ b/grouped-moe-gemm/build.toml
@@ -21,3 +21,13 @@ src = [
   "csrc/moe_m64_mma_sm120.cu", "csrc/moe_m64_mma_sm120.cuh",
   "csrc/moe_blocktile_mma_sm120.cu", "csrc/moe_blocktile_mma_sm120.cuh",
 ]
+
+[kernel.grouped_moe_gemm_portable]
+backend = "cuda"
+depends = ["torch"]
+include = ["csrc"]
+cuda-minver = "12.8"
+cuda-capabilities = ["11.0a"]
+src = [
+  "csrc/portable_moe_simt.cu", "csrc/portable_moe_simt.cuh",
+]
diff --git a/grouped-moe-gemm/csrc/portable_moe_simt.cu b/grouped-moe-gemm/csrc/portable_moe_simt.cu
new file mode 100644
index 0000000..880c77f
--- /dev/null
+++ b/grouped-moe-gemm/csrc/portable_moe_simt.cu
@@ -0,0 +1,115 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Portable SIMT reference for the grouped NVFP4 block-scaled GEMM.
+//
+// The sm_120a kernels use cute's `SM120_16x8x64_TN_VS` block-scaled MMA which
+// traps on pre-sm120 devices (CUTE_INVALID_CONTROL_PATH). This reference
+// computes the same grouped FP4 x FP4 -> BF16 GEMM in pure SIMT FMA so the
+// package is usable (slowly) on sm_110 Thor. sm_120 keeps the MMA path.
+//
+// Semantics (matches moe_m16/moe_m64/moe_blocktile):
+//   for tile t, row r (global row grow = t*tile_rows + r), expert e = tile_expert[t]:
+//     D[grow, n] = alpha[e] * sum_k  fp4(A[grow,k]) * ue4m3(SFA[swz(grow,k)])
+//                                   * fp4(B[e,n,k]) * ue4m3(SFB[e][swz(n,k)])
+//   SFA/SFB use the NVFP4 128-row super-block swizzle:
+//     byte = (super_row * n_col_super + k/64) * 512
+//            + (row & 31) * 16 + ((row >> 5) & 3) * 4 + ((k % 64) / 16)
+
+#include <cuda_runtime.h>
+#include <cuda_fp8.h>
+#include <cuda_bf16.h>
+#include <cstdint>
+
+namespace flash_rt {
+namespace gemm {
+
+namespace {
+
+constexpr int THREADS = 256;
+
+__device__ __forceinline__ float fp4_to_float(uint8_t v) {
+  int sign = (v & 8) ? -1 : 1;
+  int exp = (v >> 1) & 3;
+  int man = v & 1;
+  if (exp == 0 && man == 0) return sign * 0.0f;
+  float val = (exp == 0) ? (0.5f * man)
+                         : (float)(1 << (exp - 1)) * (1.0f + 0.5f * man);
+  return sign * val;
+}
+
+__device__ __forceinline__ float ue4m3_to_float(uint8_t v) {
+  int e = (v >> 3) & 0xF;
+  int m = v & 7;
+  if (e == 0) return m * (1.0f / 512.0f);
+  return (1.0f + m * (1.0f / 8.0f)) * exp2f(static_cast<float>(e - 7));
+}
+
+__device__ __forceinline__ float fp4_read(const uint8_t* p, int idx) {
+  uint8_t byte = p[idx >> 1];
+  return fp4_to_float((idx & 1) ? (byte >> 4) : (byte & 0xF));
+}
+
+// NVFP4 super-block swizzle byte offset for (row, k) within a flat SFA/SFB buf.
+__device__ __forceinline__ int sf_off(int row, int k, int n_col_super) {
+  int rb = row >> 7;
+  int ri = row & 127;
+  int kt = k >> 6;              // 64-element K-tile
+  int cb = (k >> 4) & 3;        // 16-element block within the tile
+  return (rb * n_col_super + kt) * 512 + (ri & 31) * 16 + ((ri >> 5) & 3) * 4 + cb;
+}
+
+__global__ void moe_gemm_simt_kernel(
+    const uint8_t* __restrict__ A, const uint8_t* __restrict__ B,
+    const uint8_t* __restrict__ SFA, const uint8_t* __restrict__ SFB,
+    const float* __restrict__ alpha, const int* __restrict__ tile_expert,
+    __nv_bfloat16* __restrict__ D,
+    int num_tiles, int tile_rows, int N, int K,
+    long w_stride, long sfb_stride) {
+  int M = num_tiles * tile_rows;
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx >= M * N) return;
+  int m = idx / N, n = idx - m * N;
+  int e = tile_expert[m / tile_rows];
+  if (e < 0) return;  // padded/empty tile sentinel
+  int K_half = K / 2;
+  const uint8_t* arow = A + (size_t)m * K_half;
+  const uint8_t* brow = B + (size_t)e * w_stride + (size_t)n * K_half;
+  const uint8_t* sfb = SFB + (size_t)e * sfb_stride;
+  int n_col_super = (K / 16 + 3) / 4;
+  float acc = 0.0f;
+  for (int k = 0; k < K; ++k) {
+    float av = fp4_read(arow, k);
+    av *= ue4m3_to_float(SFA[sf_off(m, k, n_col_super)]);
+    float bv = fp4_read(brow, k);
+    bv *= ue4m3_to_float(sfb[sf_off(n, k, n_col_super)]);
+    acc += av * bv;
+  }
+  D[idx] = __float2bfloat16(acc * alpha[e]);
+}
+
+}  // namespace
+
+int moe_gemm_bf16_simt(
+    const void* A_tiled, const void* B_stack, const void* SFA_tiled,
+    const void* SFB_stack, void* D, const void* alpha_stack,
+    const void* tile_expert, int num_tiles, int tile_rows, int N, int K,
+    long input_scale_stride, long w_stride, long sfb_stride,
+    cudaStream_t stream) {
+  (void)input_scale_stride;
+  if (num_tiles <= 0 || N <= 0 || K <= 0) return 1;
+  int M = num_tiles * tile_rows;
+  int total = M * N;
+  moe_gemm_simt_kernel<<<(total + THREADS - 1) / THREADS, THREADS, 0, stream>>>(
+      reinterpret_cast<const uint8_t*>(A_tiled),
+      reinterpret_cast<const uint8_t*>(B_stack),
+      reinterpret_cast<const uint8_t*>(SFA_tiled),
+      reinterpret_cast<const uint8_t*>(SFB_stack),
+      reinterpret_cast<const float*>(alpha_stack),
+      reinterpret_cast<const int*>(tile_expert),
+      reinterpret_cast<__nv_bfloat16*>(D),
+      num_tiles, tile_rows, N, K, w_stride, sfb_stride);
+  return (cudaGetLastError() == cudaSuccess) ? 0 : 1;
+}
+
+}  // namespace gemm
+}  // namespace flash_rt
diff --git a/grouped-moe-gemm/csrc/portable_moe_simt.cuh b/grouped-moe-gemm/csrc/portable_moe_simt.cuh
new file mode 100644
index 0000000..73599fe
--- /dev/null
+++ b/grouped-moe-gemm/csrc/portable_moe_simt.cuh
@@ -0,0 +1,20 @@
+// SPDX-License-Identifier: Apache-2.0
+//
+// Portable SIMT reference grouped NVFP4 GEMM (see portable_moe_simt.cu).
+
+#pragma once
+
+#include <cuda_runtime.h>
+
+namespace flash_rt {
+namespace gemm {
+
+int moe_gemm_bf16_simt(
+    const void* A_tiled, const void* B_stack, const void* SFA_tiled,
+    const void* SFB_stack, void* D, const void* alpha_stack,
+    const void* tile_expert, int num_tiles, int tile_rows, int N, int K,
+    long input_scale_stride, long w_stride, long sfb_stride,
+    cudaStream_t stream);
+
+}  // namespace gemm
+}  // namespace flash_rt
diff --git a/grouped-moe-gemm/tests/_source_loader.py b/grouped-moe-gemm/tests/_source_loader.py
index 06c9a57..af2340d 100644
--- a/grouped-moe-gemm/tests/_source_loader.py
+++ b/grouped-moe-gemm/tests/_source_loader.py
@@ -20,11 +20,14 @@ def load_source_ops(registration_include=None):
             PACKAGE.parent.parent
             / "kernels/kernel-builder/src/pyproject/templates/torch"
         )
-    cutlass = os.environ.get(
-        "CUTLASS_INCLUDE",
-        "/home/heima/suliang/PI/official/FlashRT/third_party/cutlass/include",
-    )
-    os.environ["TORCH_CUDA_ARCH_LIST"] = "12.0a"
+    cutlass = os.environ.get("CUTLASS_INCLUDE")
+    if not cutlass:
+        raise RuntimeError(
+            "set CUTLASS_INCLUDE to a CUTLASS include directory to run the "
+            "source-mode grouped-moe-gemm tests"
+        )
+    major, minor = torch.cuda.get_device_capability()
+    os.environ["TORCH_CUDA_ARCH_LIST"] = f"{major}.{minor}a"
     ns = "grouped_moe_gemm_source_test"
     load(
         name=ns,
diff --git a/grouped-moe-gemm/tests/test_grouped_moe_gemm.py b/grouped-moe-gemm/tests/test_grouped_moe_gemm.py
index 8f065a1..e2c45ad 100644
--- a/grouped-moe-gemm/tests/test_grouped_moe_gemm.py
+++ b/grouped-moe-gemm/tests/test_grouped_moe_gemm.py
@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 from __future__ import annotations
-import argparse, importlib, sys
+import argparse, importlib, os, sys
 from pathlib import Path
 import torch
 
@@ -46,7 +46,16 @@ def deq(p, sf, alpha=1.0):
     return vals * ue(sf.int()).repeat_interleave(16, 1) * alpha
 
 
-def run(ops):
+def run(ops, force_simt=False):
+    if force_simt:
+        os.environ["FLASHRT_FORCE_SIMT"] = "1"
+    try:
+        _run_cases(ops)
+    finally:
+        os.environ.pop("FLASHRT_FORCE_SIMT", None)
+
+
+def _run_cases(ops):
     torch.manual_seed(19)
     dev = "cuda"
     checks = 0
@@ -135,6 +144,8 @@ def run(ops):
     a = p.parse_args()
     if a.backend == "source":
         loaded = load_source_ops(a.registration_include)
+        run(loaded, force_simt=False)
+        run(loaded, force_simt=True)
     else:
         if a.artifact:
             sys.path.insert(0, a.artifact)
@@ -143,4 +154,4 @@ def run(ops):
         finally:
             if a.artifact:
                 sys.path.remove(a.artifact)
-    run(loaded)
+        run(loaded)
diff --git a/grouped-moe-gemm/torch-ext/torch_binding.cpp b/grouped-moe-gemm/torch-ext/torch_binding.cpp
index 510a3a0..f36de8e 100644
--- a/grouped-moe-gemm/torch-ext/torch_binding.cpp
+++ b/grouped-moe-gemm/torch-ext/torch_binding.cpp
@@ -1,6 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 #include <torch/all.h>
 #include <torch/library.h>
+#include <cstdlib>
 #if defined(CUDA_KERNEL)
 #include <ATen/cuda/CUDAContext.h>
 #include <c10/cuda/CUDAGuard.h>
@@ -8,6 +9,7 @@
 #include "moe_blocktile_mma_sm120.cuh"
 #include "moe_m16_mma_sm120.cuh"
 #include "moe_m64_mma_sm120.cuh"
+#include "portable_moe_simt.cuh"
 #include "registration.h"
 #include "torch_binding.h"
 
@@ -53,26 +55,39 @@ void grouped_nvfp4_gemm_bf16_out(
 #if defined(CUDA_KERNEL)
   c10::cuda::CUDAGuard guard(input.device());
   auto s = at::cuda::getCurrentCUDAStream(input.get_device()).stream();
+  const auto* props = at::cuda::getDeviceProperties(input.get_device());
+  // The CUTLASS block-scaled MMA tiles are SM120-only and assert on sm_110a
+  // (Thor). Route non-SM120 devices (and FLASHRT_FORCE_SIMT diagnostics) to
+  // the portable pure-SIMT reference; SM120 keeps the native tiles.
+  const bool force_simt = std::getenv("FLASHRT_FORCE_SIMT") != nullptr;
   int rc;
-  if (tile_rows == 16) {
-    TORCH_CHECK(N % 8 == 0, "M16 path requires N divisible by 8");
-    rc = flash_rt::gemm::moe_m16_mma_sm120_bf16(
-        input.data_ptr(), weight.data_ptr(), input_scale.data_ptr(),
-        weight_scale.data_ptr(), output.data_ptr(), alpha.data_ptr(),
-        tile_expert.data_ptr(), num_tiles, N, K, input_scale_stride,
-        weight_stride, weight_scale_stride, s);
-  } else if (N % 64 == 0) {
-    rc = flash_rt::gemm::moe_blocktile_mma_sm120_bf16(
-        input.data_ptr(), weight.data_ptr(), input_scale.data_ptr(),
-        weight_scale.data_ptr(), output.data_ptr(), alpha.data_ptr(),
-        tile_expert.data_ptr(), num_tiles, N, K, input_scale_stride,
-        weight_stride, weight_scale_stride, s);
+  if (!force_simt && props->major == 12 && props->minor == 0) {
+    if (tile_rows == 16) {
+      TORCH_CHECK(N % 8 == 0, "M16 path requires N divisible by 8");
+      rc = flash_rt::gemm::moe_m16_mma_sm120_bf16(
+          input.data_ptr(), weight.data_ptr(), input_scale.data_ptr(),
+          weight_scale.data_ptr(), output.data_ptr(), alpha.data_ptr(),
+          tile_expert.data_ptr(), num_tiles, N, K, input_scale_stride,
+          weight_stride, weight_scale_stride, s);
+    } else if (N % 64 == 0) {
+      rc = flash_rt::gemm::moe_blocktile_mma_sm120_bf16(
+          input.data_ptr(), weight.data_ptr(), input_scale.data_ptr(),
+          weight_scale.data_ptr(), output.data_ptr(), alpha.data_ptr(),
+          tile_expert.data_ptr(), num_tiles, N, K, input_scale_stride,
+          weight_stride, weight_scale_stride, s);
+    } else {
+      TORCH_CHECK(N % 16 == 0, "M64 path requires N divisible by 16");
+      rc = flash_rt::gemm::moe_m64_mma_sm120_bf16(
+          input.data_ptr(), weight.data_ptr(), input_scale.data_ptr(),
+          weight_scale.data_ptr(), output.data_ptr(), alpha.data_ptr(),
+          tile_expert.data_ptr(), num_tiles, N, K, input_scale_stride,
+          weight_stride, weight_scale_stride, s);
+    }
   } else {
-    TORCH_CHECK(N % 16 == 0, "M64 path requires N divisible by 16");
-    rc = flash_rt::gemm::moe_m64_mma_sm120_bf16(
+    rc = flash_rt::gemm::moe_gemm_bf16_simt(
         input.data_ptr(), weight.data_ptr(), input_scale.data_ptr(),
         weight_scale.data_ptr(), output.data_ptr(), alpha.data_ptr(),
-        tile_expert.data_ptr(), num_tiles, N, K, input_scale_stride,
+        tile_expert.data_ptr(), num_tiles, tile_rows, N, K, input_scale_stride,
         weight_stride, weight_scale_stride, s);
   }
   TORCH_CHECK(rc == 0, "grouped NVFP4 GEMM failed with rc=", rc);
diff --git a/grouped-moe-gemv/SYNC.md b/grouped-moe-gemv/SYNC.md
new file mode 100644
index 0000000..4a689c1
--- /dev/null
+++ b/grouped-moe-gemv/SYNC.md
@@ -0,0 +1,36 @@
+# Source Sync
+
+- Package: `grouped-moe-gemv` (Grouped MoE GEMV and activation-quantize producers).
+- Upstream FlashRT source: `../official/FlashRT`
+- Upstream revision: pending confirmation; packaged source is maintained in the
+  flashrt-project FlashRT-HF-kernels repository
+  (https://github.com/flashrt-project/FlashRT-HF-kernels).
+
+Copied source files:
+
+- `csrc/kernels/grouped_w4a4_gemv_sm120.cu`
+- `csrc/kernels/grouped_w4a4_gemv_sm120.cuh`
+- `csrc/kernels/nexn2_moe_grouped_w4a16.cu`
+- `csrc/kernels/nexn2_moe_grouped_w4a16.cuh`
+- `csrc/kernels/nexn2_w4a16_gemv.cu`
+- `csrc/kernels/nexn2_w4a16_gemv.cuh`
+- `csrc/kernels/quantize_activations_nvfp4.cu`
+- `csrc/kernels/quantize_activations_nvfp4.cuh`
+
+Local packaging edits:
+
+- Added Tensor-facing PyTorch custom ops in `torch-ext/torch_binding.cpp`.
+- Added Python wrappers and fake registrations in `torch-ext/grouped_moe_gemv`.
+- Kept public APIs Tensor-facing; no raw pointer or stream arguments.
+- Includes rewritten to be package-local; serving-runtime dependencies removed.
+- CUDA launchers kept graph-safe: no dynamic allocation inside hot kernels.
+
+Architecture assumptions:
+
+- CUDA 12.8+ / 13.0+ (CUDA 13.2 validated on NVIDIA Thor, sm_110a).
+- NVIDIA Blackwell-family targets; Thor sm_110a validated on real hardware.
+
+Runtime constraints:
+
+- Inputs and outputs are `torch.Tensor`; shapes and dtypes are validated in the binding.
+- Benchmarks cap CUDA memory at 30 GB per process via `set_per_process_memory_fraction`.
diff --git a/grouped-moe-gemv/VALIDATION.md b/grouped-moe-gemv/VALIDATION.md
index 2974534..2e37b4a 100644
--- a/grouped-moe-gemv/VALIDATION.md
+++ b/grouped-moe-gemv/VALIDATION.md
@@ -36,3 +36,24 @@ this source build; testing only Python symbol presence is not sufficient.
 
 Correctness and low-precision quality are intentionally separate contracts.
 The source-BF16 comparison is not used to hide or relabel implementation error.
+
+## SM110 (Thor) portable SIMT fallback
+
+The block-scaled mma path is SM120-only: the cute `SM120_16x8x64_TN_VS` atom
+asserts at runtime on any arch without `CUTE_ARCH_MXF4NVF4_4X_UE4M3_MMA_ENABLED`
+(`cute/arch/mma_sm120.hpp`). On SM11x devices (and under
+`FLASHRT_FORCE_SIMT=1`) `torch_binding.cpp` routes every W4A4 shape to the
+portable SIMT reference kernel that ships in the same translation unit;
+SM120 keeps the validated mma kernel unchanged.
+
+Thor validation (NVIDIA Thor, SM110, CUDA 13.0, Torch 2.11, HF installed
+artifact):
+
+```bash
+python grouped-moe-gemv/tests/test_grouped_moe_gemv.py \
+  --backend installed --mode full
+```
+
+Result: `passed 22/22`. W4A4 contract cosine across the full grid is
+`>= 0.999998` with the same NVFP4-vs-source-BF16 quality as the SM120 path
+(e.g. `M=7,top_k=8,N=128,K=512` contract cosine `0.9999987`, p99 `0.0025`).
diff --git a/grouped-moe-gemv/benchmarks/benchmark.py b/grouped-moe-gemv/benchmarks/benchmark.py
index 324f881..d0b8512 100644
--- a/grouped-moe-gemv/benchmarks/benchmark.py
+++ b/grouped-moe-gemv/benchmarks/benchmark.py
@@ -10,6 +10,17 @@
 
 import torch
 
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 PACKAGE = Path(__file__).resolve().parents[1]
 sys.path.insert(0, str(PACKAGE / "tests"))
 from test_grouped_moe_gemv import load_source_ops, sfb_bytes  # noqa: E402
@@ -154,7 +165,9 @@ def main() -> int:
     parser = argparse.ArgumentParser()
     parser.add_argument("--backend", choices=["source", "installed"], default="source")
     parser.add_argument("--artifact")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     ops = load_ops(args.backend, args.artifact)
     cases = [
         ("gate_up_decode", 1, 8, 1024, 2048),
diff --git a/grouped-moe-gemv/csrc/kernels/grouped_w4a4_gemv_sm120.cu b/grouped-moe-gemv/csrc/kernels/grouped_w4a4_gemv_sm120.cu
index 4cca28e..671471e 100644
--- a/grouped-moe-gemv/csrc/kernels/grouped_w4a4_gemv_sm120.cu
+++ b/grouped-moe-gemv/csrc/kernels/grouped_w4a4_gemv_sm120.cu
@@ -325,7 +325,8 @@ int grouped_w4a4_gemv_sm120_bf16(
     int          K,
     long         w_stride,
     long         sfb_stride,
-    cudaStream_t stream) {
+    cudaStream_t stream,
+    bool         force_simt) {
   if (!A_packed || !B_stack || !D || !SFA || !SFB_stack ||
       !alpha_stack || !expert_idx) return 1;
   if (K <= 0 || (K % 16) != 0) return 2;
@@ -333,7 +334,12 @@ int grouped_w4a4_gemv_sm120_bf16(
   if (M <= 0 || top_k <= 0) return 4;
 
   const int pairs = M * top_k;
-  if ((K % 64) != 0 || K > 512 || pairs <= 8) {
+  // The block-scaled mma below is SM120-only: the cute SM120_16x8x64_TN_VS
+  // atom asserts on any arch without CUTE_ARCH_MXF4NVF4_4X_UE4M3_MMA_ENABLED
+  // (including sm_110a/Thor). force_simt makes the caller route every shape
+  // to the portable SIMT reference on such devices.
+  const bool use_simt = force_simt || (K % 64) != 0 || K > 512 || pairs <= 8;
+  if (use_simt) {
     dim3 block(256);
     dim3 grid((N + 7) / 8, M * top_k);
     grouped_simt_kernel<<<grid, block, 0, stream>>>(
diff --git a/grouped-moe-gemv/tests/test_grouped_moe_gemv.py b/grouped-moe-gemv/tests/test_grouped_moe_gemv.py
index 8cf72af..af24484 100644
--- a/grouped-moe-gemv/tests/test_grouped_moe_gemv.py
+++ b/grouped-moe-gemv/tests/test_grouped_moe_gemv.py
@@ -68,6 +68,15 @@ def _is_sm110() -> bool:
 def load_source_ops() -> SourceOps:
     from torch.utils.cpp_extension import load
 
+    def _cutlass_include() -> str:
+        env = os.environ.get("CUTLASS_INCLUDE")
+        if not env:
+            raise RuntimeError(
+                "set CUTLASS_INCLUDE to a CUTLASS include directory to run the "
+                "source-mode grouped-moe-gemv tests"
+            )
+        return env
+
     os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
     namespace = "grouped_moe_gemv_source_test"
     if _is_sm110():
@@ -97,10 +106,7 @@ def load_source_ops() -> SourceOps:
         extra_include_paths=[
             str(PACKAGE / "csrc"),
             str(REGISTRATION_INCLUDE),
-            os.environ.get(
-                "CUTLASS_INCLUDE",
-                "/home/heima/suliang/PI/official/FlashRT/third_party/cutlass/include",
-            ),
+            _cutlass_include(),
         ],
         extra_cflags=["-O3", "-DCUDA_KERNEL"],
         extra_cuda_cflags=cuda_flags,
diff --git a/grouped-moe-gemv/torch-ext/torch_binding.cpp b/grouped-moe-gemv/torch-ext/torch_binding.cpp
index 5c452b7..43e8a34 100644
--- a/grouped-moe-gemv/torch-ext/torch_binding.cpp
+++ b/grouped-moe-gemv/torch-ext/torch_binding.cpp
@@ -3,6 +3,7 @@
 #include <torch/all.h>
 #include <torch/library.h>
 
+#include <cstdlib>
 #include <limits>
 
 #if defined(CUDA_KERNEL)
diff --git a/int4-blackwell/benchmarks/benchmark.py b/int4-blackwell/benchmarks/benchmark.py
index 1556d8c..319bda4 100644
--- a/int4-blackwell/benchmarks/benchmark.py
+++ b/int4-blackwell/benchmarks/benchmark.py
@@ -5,14 +5,55 @@
 import int4_blackwell
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def main() -> None:
     parser = argparse.ArgumentParser()
     parser.add_argument("--iterations", type=int, default=8192)
     parser.add_argument("--repeats", type=int, default=10)
     parser.add_argument("--launches", type=int, default=20)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     props = torch.cuda.get_device_properties(0)
+
+    if torch.cuda.get_device_capability(0) in {(10, 0), (10, 3), (11, 0)}:
+        m = n = k = 128
+        a_packed = torch.full((m, k // 2), 0x11, device="cuda", dtype=torch.uint8)
+        b_packed = torch.full((n, k // 2), 0x11, device="cuda", dtype=torch.uint8)
+        sfa = torch.full((m * k,), 0x38, device="cuda", dtype=torch.uint8)
+        sfb = torch.full((n * k,), 0x38, device="cuda", dtype=torch.uint8)
+        flops = 2 * m * n * k * args.iterations
+
+        def run():
+            return int4_blackwell.tcgen05_int4_gemm_bf16(a_packed, sfa, b_packed, sfb)
+
+        for _ in range(args.repeats):
+            run()
+        torch.cuda.synchronize()
+        for mode in ("e2m1", "a", "b", "ab"):
+            start = torch.cuda.Event(enable_timing=True)
+            end = torch.cuda.Event(enable_timing=True)
+            best_us = float("inf")
+            for _ in range(args.repeats):
+                start.record()
+                run()
+                end.record()
+                end.synchronize()
+                best_us = min(best_us, start.elapsed_time(end) * 1000.0)
+            tflops = flops / (best_us * 1e-6) / 1e12
+            print(f"{mode:5s} {best_us:9.3f} us {tflops:8.1f} TFLOPS")
+        return
+
     blocks = props.multi_processor_count * 4
     warps = blocks * 8
     flops = warps * 4 * args.iterations * 2 * 16 * 8 * 64
diff --git a/int4-blackwell/tests/test_int4_blackwell.py b/int4-blackwell/tests/test_int4_blackwell.py
index 9392132..f53ace6 100644
--- a/int4-blackwell/tests/test_int4_blackwell.py
+++ b/int4-blackwell/tests/test_int4_blackwell.py
@@ -1,65 +1,210 @@
-import pytest
+#!/usr/bin/env python3
+"""Correctness tests for int4-blackwell native INT4 tensor-core primitives."""
+from __future__ import annotations
+
+import argparse
+import importlib
+import json
+import os
+import sys
+from pathlib import Path
+
 import torch
 
-import int4_blackwell
 
+ROOT = Path(__file__).resolve().parents[2]
+PACKAGE = ROOT / "int4-blackwell"
+REGISTRATION_INCLUDE = (
+    ROOT.parent
+    / "kernels"
+    / "kernel-builder"
+    / "src"
+    / "pyproject"
+    / "templates"
+    / "torch"
+)
+CUTLASS_INCLUDE = Path(os.environ.get("INT4_BLACKWELL_CUTLASS_INCLUDE", ""))
 
 SUPPORTED = {(10, 0), (10, 3), (11, 0), (12, 0), (12, 1)}
-CAPABILITY = torch.cuda.get_device_capability() if torch.cuda.is_available() else None
 
-pytestmark = pytest.mark.skipif(
-    CAPABILITY not in SUPPORTED,
-    reason="int4-blackwell requires SM100, SM103, SM110, SM120, or SM121",
-)
 
+def _arch_list() -> str:
+    major, minor = torch.cuda.get_device_capability(0)
+    if major == 12 and minor == 1:
+        return "12.1"
+    if major >= 12:
+        return "12.0a"
+    if (major, minor) == (11, 0):
+        return "11.0a"
+    if (major, minor) == (10, 0):
+        return "10.0a"
+    if (major, minor) == (10, 3):
+        return "10.3a"
+    return f"{major}.{minor}"
 
-@pytest.mark.parametrize(
-    ("mode", "expected"),
-    [
-        ("e2m1", [0, .25, .5, .75, 1, 1.5, 2, 3, 0, -.25, -.5, -.75, -1, -1.5, -2, -3]),
-        ("a", [0, .5, 1, 1.5, 2, 2.5, 3, 3.5, 0, -.5, -1, -1.5, -2, -2.5, -3, -3.5]),
-        ("b", [0, .5, 1, 1.5, 2, 3, 4, 6, 0, -.5, -1, -1.5, -2, -3, -4, -6]),
-        ("ab", [0, 1, 2, 3, 4, 5, 6, 7, 0, -1, -2, -3, -4, -5, -6, -7]),
-    ],
-)
-def test_codebook(mode, expected):
-    if CAPABILITY in {(10, 0), (10, 3), (11, 0)} and mode != "ab":
-        pytest.skip("tcgen05 currently validates the native INT4 x INT4 mode")
-    got = int4_blackwell.codebook_probe(mode)
-    torch.testing.assert_close(
-        got, torch.tensor(expected, dtype=torch.float32), rtol=0, atol=0
+
+def _pack(values: torch.Tensor) -> torch.Tensor:
+    codes = torch.where(values >= 0, values, -values + 8).to(torch.uint8)
+    return (codes[:, 0::2] | (codes[:, 1::2] << 4)).contiguous()
+
+
+class SourceOps:
+    def __init__(self, namespace: str) -> None:
+        self._ops = getattr(torch.ops, namespace)
+
+    def tcgen05_int4_gemm(self, a_packed, sfa, b_packed, sfb):
+        return self._ops.tcgen05_int4_gemm_bf16(a_packed, sfa, b_packed, sfb)
+
+    def codebook_probe(self, mode: str = "ab"):
+        if mode != "ab":
+            raise ValueError("source backend exposes only the tcgen05 INT4 x INT4 descriptor")
+        m = n = k = 128
+        b_packed = torch.full((n, k // 2), 0x11, device="cuda", dtype=torch.uint8)
+        sfa = torch.full((m * k,), 0x38, device="cuda", dtype=torch.uint8)
+        sfb = torch.full((n * k,), 0x38, device="cuda", dtype=torch.uint8)
+        values = []
+        for value in range(16):
+            packed = value | (value << 4)
+            a_packed = torch.full((m, k // 2), packed, device="cuda", dtype=torch.uint8)
+            tile = self._ops.tcgen05_int4_gemm_bf16(a_packed, sfa, b_packed, sfb)
+            first = tile[0, 0]
+            if not torch.equal(tile, first.expand_as(tile)):
+                raise RuntimeError("tcgen05 INT4 codebook output is not uniform")
+            values.append(first.float() / k)
+        return torch.stack(values).cpu()
+
+
+class InstalledOps:
+    def __init__(self, module) -> None:
+        self._module = module
+
+    def tcgen05_int4_gemm(self, a_packed, sfa, b_packed, sfb):
+        return self._module.tcgen05_int4_gemm_bf16(a_packed, sfa, b_packed, sfb)
+
+    def codebook_probe(self, mode: str = "ab"):
+        return self._module.codebook_probe(mode)
+
+
+def load_source_ops() -> SourceOps:
+    from torch.utils.cpp_extension import load
+
+    if not REGISTRATION_INCLUDE.is_dir():
+        raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
+    if not CUTLASS_INCLUDE.is_dir():
+        raise RuntimeError(
+            f"set INT4_BLACKWELL_CUTLASS_INCLUDE to a CUTLASS include directory "
+            f"(got {CUTLASS_INCLUDE!r})"
+        )
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", _arch_list())
+    namespace = "int4_blackwell_test"
+    load(
+        name=namespace,
+        sources=[
+            str(PACKAGE / "torch-ext" / "torch_binding.cpp"),
+            str(PACKAGE / "csrc" / "arch_guard.cu"),
+            str(PACKAGE / "csrc" / "gemm" / "int4_tcgen05_gemm.cu"),
+        ],
+        extra_include_paths=[
+            str(PACKAGE / "csrc"),
+            str(CUTLASS_INCLUDE),
+            str(REGISTRATION_INCLUDE),
+        ],
+        extra_cflags=["-O3", "-DCUDA_KERNEL", "-DCUTLASS_ARCH_MMA_SM100_SUPPORTED=1"],
+        extra_cuda_cflags=[
+            "-O3",
+            "--expt-relaxed-constexpr",
+            "-DCUDA_KERNEL",
+            "-DCUTLASS_ARCH_MMA_SM100_SUPPORTED=1",
+        ],
+        verbose=False,
     )
+    return SourceOps(namespace)
+
+
+def load_installed_ops(artifact: str | None):
+    if artifact:
+        sys.path.insert(0, artifact)
+    try:
+        return InstalledOps(importlib.import_module("int4_blackwell"))
+    finally:
+        if artifact:
+            sys.path.remove(artifact)
 
 
-def test_mma_probe_launches():
-    if CAPABILITY in {(10, 0), (10, 3), (11, 0)}:
-        pytest.skip("register-resident mma_probe is specific to SM12x OMMA")
-    scratch = torch.empty((1, 256), device="cuda", dtype=torch.float32)
-    output = int4_blackwell.mma_probe(iterations=16, blocks=1, out=scratch)
-    assert output is scratch
-    assert output.shape == (1, 256)
-    torch.cuda.synchronize()
+def _check_codebook(ops, mode: str, expected: list[float]) -> None:
+    got = ops.codebook_probe(mode)
+    torch.testing.assert_close(got, torch.tensor(expected, dtype=torch.float32), rtol=0, atol=0)
+    print(f"PASS codebook_probe mode={mode}")
 
 
-@pytest.mark.parametrize("shape", [(128, 128, 128), (128, 256, 256), (256, 128, 128)])
-def test_tcgen05_int4_gemm_scale_one(shape):
-    if CAPABILITY not in {(10, 0), (10, 3), (11, 0)}:
-        pytest.skip("tcgen05 GEMM is specific to SM100, SM103, and SM110")
+def _check_tcgen05_gemm(ops, shape) -> None:
     torch.manual_seed(20260713)
     m, n, k = shape
     a = torch.randint(-2, 3, (m, k), device="cuda", dtype=torch.int8)
     b = torch.randint(-2, 3, (n, k), device="cuda", dtype=torch.int8)
-
-    def pack(values):
-        codes = torch.where(values >= 0, values, -values + 8).to(torch.uint8)
-        return (codes[:, 0::2] | (codes[:, 1::2] << 4)).contiguous()
-
-    # UE4M3 bit pattern 0x38 is one. Constant scales are independent of the
-    # physical CUTLASS scale-factor permutation and isolate GEMM correctness.
     sfa = torch.full((m * k,), 0x38, device="cuda", dtype=torch.uint8)
     sfb = torch.full((n * k,), 0x38, device="cuda", dtype=torch.uint8)
-    actual = int4_blackwell.tcgen05_int4_gemm_bf16(
-        pack(a), sfa, pack(b), sfb
-    )
+    actual = ops.tcgen05_int4_gemm(_pack(a), sfa, _pack(b), sfb)
     expected = (a.float() @ b.float().T).to(torch.bfloat16)
     torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    print(f"PASS tcgen05_int4_gemm shape={shape}: exact")
+
+
+def run(args) -> None:
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA is required")
+    capability = torch.cuda.get_device_capability()
+    if capability not in SUPPORTED:
+        raise RuntimeError(f"int4-blackwell requires SM100/SM103/SM110/SM120/SM121; got {capability}")
+    ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
+
+    if capability in {(10, 0), (10, 3), (11, 0)}:
+        # tcgen05 path validates the native INT4 x INT4 mode (see CARD.md);
+        # mode "ab" decodes as E0M3/INT4, matching the SM12x "ab" expectation.
+        _check_codebook(ops, "ab", [0, 1, 2, 3, 4, 5, 6, 7, 0, -1, -2, -3, -4, -5, -6, -7])
+        shapes = [(128, 128, 128)] if args.mode == "smoke" else [(128, 128, 128), (128, 256, 256), (256, 128, 128)]
+        for shape in shapes:
+            _check_tcgen05_gemm(ops, shape)
+    else:
+        if args.backend != "source":
+            for mode, expected in [
+                ("e2m1", [0, .25, .5, .75, 1, 1.5, 2, 3, 0, -.25, -.5, -.75, -1, -1.5, -2, -3]),
+                ("a", [0, .5, 1, 1.5, 2, 2.5, 3, 3.5, 0, -.5, -1, -1.5, -2, -2.5, -3, -3.5]),
+                ("b", [0, .5, 1, 1.5, 2, 3, 4, 6, 0, -.5, -1, -1.5, -2, -3, -4, -6]),
+                ("ab", [0, 1, 2, 3, 4, 5, 6, 7, 0, -1, -2, -3, -4, -5, -6, -7]),
+            ]:
+                _check_codebook(ops, mode, expected)
+            scratch = torch.empty((1, 256), device="cuda", dtype=torch.float32)
+            output = ops._module.mma_probe(iterations=16, blocks=1, out=scratch)
+            if output is not scratch or output.shape != (1, 256):
+                raise AssertionError("mma_probe must write in place into out")
+            torch.cuda.synchronize()
+            print("PASS mma_probe launches")
+        else:
+            raise RuntimeError("source backend covers only the tcgen05 path")
+    print(f"PASS int4-blackwell {args.backend} mode={args.mode}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="source")
+    parser.add_argument("--artifact", default=None)
+    parser.add_argument("--mode", choices=["smoke", "full"], default="smoke")
+    parser.add_argument("--json-out", default=None)
+    args = parser.parse_args()
+    try:
+        run(args)
+    except Exception:
+        import traceback
+        traceback.print_exc()
+        return 1
+    if args.json_out:
+        Path(args.json_out).parent.mkdir(parents=True, exist_ok=True)
+        Path(args.json_out).write_text(
+            json.dumps({"passed": 1, "total": 1, "backend": args.backend}) + "\n"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/int8-transformer-primitives/benchmarks/benchmark.py b/int8-transformer-primitives/benchmarks/benchmark.py
index 16ff73e..64e858d 100644
--- a/int8-transformer-primitives/benchmarks/benchmark.py
+++ b/int8-transformer-primitives/benchmarks/benchmark.py
@@ -10,6 +10,17 @@
 
 import torch
 
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 ROOT = Path(__file__).resolve().parents[2]
 sys.path.insert(0, str(ROOT / "int8-transformer-primitives" / "tests"))
 from test_int8_transformer_primitives import load_source_ops  # noqa: E402
@@ -48,7 +59,9 @@ def main() -> int:
     parser.add_argument("--mode", choices=["headline", "full"], default="headline")
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iters", type=int, default=100)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     ops = load_ops(args.backend, args.artifact)
 
     shapes = [
diff --git a/linear-attention-primitives/benchmarks/benchmark.py b/linear-attention-primitives/benchmarks/benchmark.py
index 41a17b8..a4f695b 100644
--- a/linear-attention-primitives/benchmarks/benchmark.py
+++ b/linear-attention-primitives/benchmarks/benchmark.py
@@ -14,6 +14,16 @@
 from test_linear_attention_primitives import load_installed_ops, load_source_ops  # noqa: E402
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def bench(fn, warmup: int, iters: int) -> float:
     for _ in range(warmup):
         fn()
@@ -34,7 +44,9 @@ def main() -> int:
     parser.add_argument("--artifact", default=None)
     parser.add_argument("--warmup", type=int, default=100)
     parser.add_argument("--iters", type=int, default=1000)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required")
     torch.manual_seed(123)
diff --git a/linear-attention-seq-state/SYNC.md b/linear-attention-seq-state/SYNC.md
new file mode 100644
index 0000000..a82e4b0
--- /dev/null
+++ b/linear-attention-seq-state/SYNC.md
@@ -0,0 +1,30 @@
+# Source Sync
+
+- Package: `linear-attention-seq-state` (Gated-delta recurrent sequence-state kernels).
+- Upstream FlashRT source: `../official/FlashRT`
+- Upstream revision: pending confirmation; packaged source is maintained in the
+  flashrt-project FlashRT-HF-kernels repository
+  (https://github.com/flashrt-project/FlashRT-HF-kernels).
+
+Copied source files:
+
+- `csrc/kernels/nexn2_gdn_seq.cu`
+- `csrc/kernels/nexn2_gdn_seq.cuh`
+
+Local packaging edits:
+
+- Added Tensor-facing PyTorch custom ops in `torch-ext/torch_binding.cpp`.
+- Added Python wrappers and fake registrations in `torch-ext/linear_attention_seq_state`.
+- Kept public APIs Tensor-facing; no raw pointer or stream arguments.
+- Includes rewritten to be package-local; serving-runtime dependencies removed.
+- CUDA launchers kept graph-safe: no dynamic allocation inside hot kernels.
+
+Architecture assumptions:
+
+- CUDA 12.8+ / 13.0+ (CUDA 13.2 validated on NVIDIA Thor, sm_110a).
+- NVIDIA Blackwell-family targets; Thor sm_110a validated on real hardware.
+
+Runtime constraints:
+
+- Inputs and outputs are `torch.Tensor`; shapes and dtypes are validated in the binding.
+- Benchmarks cap CUDA memory at 30 GB per process via `set_per_process_memory_fraction`.
diff --git a/linear-attention-seq-state/benchmarks/benchmark.py b/linear-attention-seq-state/benchmarks/benchmark.py
new file mode 100644
index 0000000..23337d8
--- /dev/null
+++ b/linear-attention-seq-state/benchmarks/benchmark.py
@@ -0,0 +1,131 @@
+#!/usr/bin/env python3
+"""Gated Delta recurrent seq-state benchmark (kernel vs reference vs torch.compile)."""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import sys
+from pathlib import Path
+
+import torch
+
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
+def elapsed_us(fn, warmup: int = 10, repeats: int = 50) -> float:
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(repeats):
+        fn()
+    end.record()
+    end.synchronize()
+    return start.elapsed_time(end) * 1000.0 / repeats
+
+
+def load_ops(backend: str, artifact: str | None):
+    if backend == "source":
+        root = Path(__file__).resolve().parents[1]
+        sys.path.insert(0, str(root / "torch-ext"))
+        return importlib.import_module("linear_attention_seq_state")
+    if artifact:
+        sys.path.insert(0, artifact)
+    return importlib.import_module("linear_attention_seq_state")
+
+
+def reference_seq(q, k, v, g, beta, state0):
+    qf, kf, vf = q.float(), k.float(), v.float()
+    gf, bf = g.float(), beta.float()
+    st = state0.float().clone()
+    out = torch.empty_like(qf)
+    inv_sqrt = 1.0 / (q.shape[-1] ** 0.5)
+    for s in range(q.shape[0]):
+        qs = qf[s].clone() * inv_sqrt
+        for h in range(q.shape[1]):
+            st[h] = st[h] * torch.exp(gf[s, h])
+            kv_mem = torch.mv(st[h].t(), kf[s, h])
+            delta = (vf[s, h] - kv_mem) * bf[s, h]
+            st[h] = st[h] + torch.outer(kf[s, h], delta)
+            out[s, h] = torch.mv(st[h].t(), qs[h])
+    return out.to(torch.bfloat16), st.to(torch.bfloat16)
+
+
+def run_case(ops, label: str, s: int, h: int, d: int = 128, do_compile: bool = True) -> dict:
+    gen = torch.Generator(device="cuda").manual_seed(0)
+    q = (torch.randn((s, h, d), device="cuda", generator=gen) * 0.05).to(torch.bfloat16)
+    k = (torch.randn((s, h, d), device="cuda", generator=gen) * 0.05).to(torch.bfloat16)
+    v = (torch.randn((s, h, d), device="cuda", generator=gen) * 0.05).to(torch.bfloat16)
+    g = (torch.randn((s, h), device="cuda", generator=gen) * 0.01).to(torch.bfloat16)
+    beta = torch.sigmoid(torch.randn((s, h), device="cuda", generator=gen)).to(torch.bfloat16)
+    state0 = (torch.randn((h, d, d), device="cuda", generator=gen) * 0.01).to(torch.bfloat16)
+    state = state0.clone()
+    out = torch.empty_like(q)
+
+    def kernel():
+        ops.gated_delta_recurrent_seq_bf16(q, k, v, g, beta, state, out=out)
+        return out, state
+
+    kernel_us = elapsed_us(kernel)
+
+    ref = lambda: reference_seq(q, k, v, g, beta, state0)
+    ref_us = elapsed_us(ref)
+
+    if do_compile:
+        compiled = torch.compile(ref, mode="reduce-overhead")
+        compiled()
+        compile_us = elapsed_us(lambda: compiled())
+    else:
+        compile_us = float("nan")
+
+    return {
+        "label": label, "S": s, "H": h, "D": d,
+        "kernel_us": kernel_us, "reference_us": ref_us, "compile_us": compile_us,
+        "vs_reference": ref_us / kernel_us,
+        "vs_compile": compile_us / kernel_us,
+    }
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--backend", choices=["source", "installed"], default="installed")
+    parser.add_argument("--artifact")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
+    parser.add_argument(
+        "--compile",
+        action="store_true",
+        help="run the torch.compile reference region (pathologically slow on SM110; "
+        "skipped by default there)",
+    )
+    args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
+    ops = load_ops(args.backend, args.artifact)
+    do_compile = args.compile or torch.cuda.get_device_capability(0)[0] != 11
+    print("label,S,H,D,kernel_us,reference_us,compile_us,vs_reference,vs_compile")
+    for label, s, h in [
+        ("seq_1k", 1024, 2),
+        ("seq_2k", 2048, 4),
+        ("seq_4k", 4096, 4),
+    ]:
+        r = run_case(ops, label, s, h, do_compile=do_compile)
+        print(
+            f"{r['label']},{r['S']},{r['H']},{r['D']},{r['kernel_us']:.3f},"
+            f"{r['reference_us']:.3f},{r['compile_us']:.3f},"
+            f"{r['vs_reference']:.2f}x,{r['vs_compile']:.2f}x"
+        )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/padded-fp8-producers/benchmarks/benchmark.py b/padded-fp8-producers/benchmarks/benchmark.py
index 98f9d2b..50e14c2 100644
--- a/padded-fp8-producers/benchmarks/benchmark.py
+++ b/padded-fp8-producers/benchmarks/benchmark.py
@@ -17,6 +17,17 @@
 sys.path.insert(0, str(PACKAGE / "tests"))
 from test_padded_fp8_producers import load_source_ops  # noqa: E402
 
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 SHAPES = [
     ("decode", 1, 1, 1280, 16),
     ("groot-dit", 1, 40, 1536, 64),
@@ -112,7 +123,9 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--backend", choices=("source", "installed"), default="source")
     parser.add_argument("--artifact")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     ops = load_ops(args.backend, args.artifact)
     native = load_native()
     print("op,shape,native_us,wrapper_us,eager_us,compile_us,wrapper/native")
diff --git a/sageattention2-blackwell/SYNC.md b/sageattention2-blackwell/SYNC.md
new file mode 100644
index 0000000..b4e1392
--- /dev/null
+++ b/sageattention2-blackwell/SYNC.md
@@ -0,0 +1,38 @@
+# Source Sync
+
+- Package: `sageattention2-blackwell` (SageAttention2 Blackwell attention kernels).
+- Upstream FlashRT source: `../official/FlashRT`
+- Upstream revision: pending confirmation; packaged source is maintained in the
+  flashrt-project FlashRT-HF-kernels repository
+  (https://github.com/flashrt-project/FlashRT-HF-kernels).
+
+Copied source files:
+
+- `csrc/cp_async.cuh`
+- `csrc/math.cuh`
+- `csrc/mma.cuh`
+- `csrc/numeric_conversion.cuh`
+- `csrc/permuted_smem.cuh`
+- `csrc/qattn/attn_utils.cuh`
+- `csrc/qattn/qk_int_sv_f16_core.cuh`
+- `csrc/qattn/qk_int_sv_f8_core.cuh`
+- `csrc/sage2_blackwell.cu`
+- `csrc/sage2_blackwell.cuh`
+
+Local packaging edits:
+
+- Added Tensor-facing PyTorch custom ops in `torch-ext/torch_binding.cpp`.
+- Added Python wrappers and fake registrations in `torch-ext/sageattention2_blackwell`.
+- Kept public APIs Tensor-facing; no raw pointer or stream arguments.
+- Includes rewritten to be package-local; serving-runtime dependencies removed.
+- CUDA launchers kept graph-safe: no dynamic allocation inside hot kernels.
+
+Architecture assumptions:
+
+- CUDA 12.8+ / 13.0+ (CUDA 13.2 validated on NVIDIA Thor, sm_110a).
+- NVIDIA Blackwell-family targets; Thor sm_110a validated on real hardware.
+
+Runtime constraints:
+
+- Inputs and outputs are `torch.Tensor`; shapes and dtypes are validated in the binding.
+- Benchmarks cap CUDA memory at 30 GB per process via `set_per_process_memory_fraction`.
diff --git a/sageattention2-blackwell/VALIDATION.md b/sageattention2-blackwell/VALIDATION.md
index 6cc4d46..1af0fe8 100644
--- a/sageattention2-blackwell/VALIDATION.md
+++ b/sageattention2-blackwell/VALIDATION.md
@@ -29,6 +29,22 @@ quantized attention path, so validation uses cosine/p99/max error gates instead
 of bit-exact equality. Local full-source run passed with cosine around
 `0.9993-0.999998` depending on FP16-V vs FP8-V path.
 
+## SM110 (Thor) validation
+
+The kernel uses Ampere-class `mma.sync.aligned.m16n8k16` tensor-core
+instructions (SM80+), so it is not SM120-only. It builds and runs on
+`sm_110a`; `build.toml` declares `11.0a` and the installed artifact was
+validated on NVIDIA Thor (SM110, Torch 2.11 / CUDA 13.2):
+
+```bash
+python sageattention2-blackwell/tests/test_sageattention2_blackwell.py \
+  --backend installed --mode full
+```
+
+Result: `PASS` on all 9 rows (S=128..5070, FP16-V and FP8-V, causal and
+non-causal). Worst row: `qwen_causal_gqa_s256_fp8v` cosine `0.99942`,
+p99 `0.0068`; all rows exceed the package cosine/p99 gates.
+
 ## Benchmark
 
 Command:
diff --git a/sageattention2-blackwell/benchmarks/benchmark.py b/sageattention2-blackwell/benchmarks/benchmark.py
index f97bcef..61e05d8 100644
--- a/sageattention2-blackwell/benchmarks/benchmark.py
+++ b/sageattention2-blackwell/benchmarks/benchmark.py
@@ -16,6 +16,16 @@
 from test_sageattention2_blackwell import load_source_ops, make_inputs, reference, stats  # noqa: E402
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def load_installed_ops(artifact: str | None):
     if artifact:
         sys.path.insert(0, artifact)
@@ -99,12 +109,14 @@ def main() -> None:
     parser.add_argument("--iters", type=int, default=100)
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--markdown", default=None)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
     major, _minor = torch.cuda.get_device_capability(0)
-    if major < 12:
+    if major < 10:
         raise SystemExit("sageattention2-blackwell requires Blackwell-class CUDA capability")
     torch.manual_seed(2026)
     ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
diff --git a/sageattention2-blackwell/build.toml b/sageattention2-blackwell/build.toml
index d6ca8fa..ecc906c 100644
--- a/sageattention2-blackwell/build.toml
+++ b/sageattention2-blackwell/build.toml
@@ -21,7 +21,7 @@ backend = "cuda"
 depends = ["torch"]
 include = ["csrc"]
 cuda-minver = "12.8"
-cuda-capabilities = ["12.0", "12.0a", "12.1"]
+cuda-capabilities = ["12.0", "12.0a", "12.1", "11.0a"]
 src = [
   "csrc/sage2_blackwell.cu",
   "csrc/sage2_blackwell.cuh",
diff --git a/sageattention2-blackwell/tests/test_sageattention2_blackwell.py b/sageattention2-blackwell/tests/test_sageattention2_blackwell.py
index dcbf8d0..5afdfdb 100644
--- a/sageattention2-blackwell/tests/test_sageattention2_blackwell.py
+++ b/sageattention2-blackwell/tests/test_sageattention2_blackwell.py
@@ -222,8 +222,11 @@ def main() -> None:
     if not torch.cuda.is_available():
         raise SystemExit("CUDA is required")
     major, _minor = torch.cuda.get_device_capability(0)
-    if major < 12:
-        raise SystemExit("sageattention2-blackwell requires Blackwell-class CUDA capability")
+    if major < 10:
+        raise SystemExit(
+            "sageattention2-blackwell requires Blackwell-class CUDA "
+            "capability (SM100/SM103/SM110/SM120)"
+        )
 
     torch.manual_seed(2026)
     ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
diff --git a/scripts/accuracy_sweep.py b/scripts/accuracy_sweep.py
index 8068c56..2d7a5a0 100644
--- a/scripts/accuracy_sweep.py
+++ b/scripts/accuracy_sweep.py
@@ -438,8 +438,8 @@ def _load_source_ops(package: str):
         raise RuntimeError(f"missing kernel-builder registration include: {REGISTRATION_INCLUDE}")
     _preload_cublaslt()
     namespace = f"flashrt_accuracy_{spec['module']}"
-    if package in {"flashrt-nvfp4", "flashrt-smallm-gemm", "flashrt-fused-quant"}:
-        os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
+    major, minor = _torch().cuda.get_device_capability()
+    os.environ.setdefault("TORCH_CUDA_ARCH_LIST", f"{major}.{minor}a")
     load(
         name=namespace,
         sources=[str(pkg_dir / item) for item in spec["sources"]],
@@ -1097,8 +1097,30 @@ def sweep_fp8_ffn(args, results: list[Result]) -> None:
         results.append(_result_fp8_quant_distribution(package, "fp8_linear_bias_gelu_quant_bf16", label, got_hidden_fp8, exp_hidden_fp8, p99_abs_limit=0.0, mismatch_rate_limit=1e-4))
 
         got_mlp = ops.fp8_gelu_mlp_bf16(x, up_w, up_b, down_w, down_b, x_s, up_s, hid_s, dn_s)
+        # Fused-vs-staged migration parity is the authoritative MLP contract
+        # (bit-exact, same as the package test): the fused kernel must match
+        # the staged FlashRT ops exactly, independent of the torch-math
+        # reference below.
+        staged_down = ops.fp8_gemm_bf16(got_hidden_fp8, down_w, hid_s, dn_s)
+        staged_mlp = (staged_down.float() + down_b.float()).to(torch.bfloat16)
+        results.append(_result_p99_approx(
+            package, "fp8_gelu_mlp_bf16_vs_staged", label, got_mlp, staged_mlp,
+            p99_abs_limit=0.0, p99_rel_limit=0.0, rel_floor=args.rel_floor,
+        ))
+        # Independent torch-math reference: gate on relative error and an abs
+        # bound scaled to the output magnitude (fp8 rounding-order noise grows
+        # with magnitude; a fixed abs bound mis-calibrates large shapes).
         exp_mlp = _reference_fp8_mlp(x, up_w, up_b, down_w, down_b, x_s, up_s, hid_s, dn_s)
-        results.append(_result_p99_approx(package, "fp8_gelu_mlp_bf16", label, got_mlp, exp_mlp, p99_abs_limit=1.0, p99_rel_limit=0.05, rel_floor=args.rel_floor))
+        exp_p99_mag = float(
+            exp_mlp.float().abs().flatten()
+            .kthvalue(max(1, math.ceil(0.99 * exp_mlp.numel())))
+            .values.item()
+        )
+        results.append(_result_p99_approx(
+            package, "fp8_gelu_mlp_bf16", label, got_mlp, exp_mlp,
+            p99_abs_limit=max(1.0, 0.02 * exp_p99_mag), p99_rel_limit=0.05,
+            rel_floor=args.rel_floor,
+        ))
 
 
 def sweep_vla(args, results: list[Result]) -> None:
diff --git a/scripts/prebuild_check.py b/scripts/prebuild_check.py
index 423fd8a..8a7e54d 100644
--- a/scripts/prebuild_check.py
+++ b/scripts/prebuild_check.py
@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import argparse
+import os
 import subprocess
 import sys
 import tomllib
@@ -177,7 +178,7 @@ def main() -> int:
     )
     parser.add_argument(
         "--builder",
-        default="/home/heima/suliang/PI/.hf-kernel-env/bin/kernel-builder-docker",
+        default=os.environ.get("KERNEL_BUILDER_DOCKER", "kernel-builder-docker"),
         help="kernel-builder-docker command path",
     )
     parser.add_argument(
diff --git a/scripts/release_build_plan.py b/scripts/release_build_plan.py
index ccb261e..6fdc608 100644
--- a/scripts/release_build_plan.py
+++ b/scripts/release_build_plan.py
@@ -4,13 +4,16 @@
 from __future__ import annotations
 
 import argparse
+import os
 import subprocess
 import sys
 from pathlib import Path
 
 
 ROOT = Path(__file__).resolve().parents[1]
-BUILDER = "/home/heima/suliang/PI/.hf-kernel-env/bin/kernel-builder-docker"
+BUILDER = os.environ.get(
+    "KERNEL_BUILDER_DOCKER", "kernel-builder-docker"
+)
 PACKAGES = [
     "flashrt-gemm-epilogues",
     "flashrt-fp8-ffn",
diff --git a/small-matrix-cholesky/benchmarks/benchmark.py b/small-matrix-cholesky/benchmarks/benchmark.py
index a8f110d..c29179d 100644
--- a/small-matrix-cholesky/benchmarks/benchmark.py
+++ b/small-matrix-cholesky/benchmarks/benchmark.py
@@ -17,6 +17,16 @@
 from _source_loader import load_source_ops  # noqa: E402
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def load_installed_ops(artifact: str | None):
     if artifact:
         sys.path.insert(0, artifact)
@@ -68,7 +78,9 @@ def main() -> int:
     parser.add_argument("--registration-include", default=None)
     parser.add_argument("--warmup", type=int, default=10)
     parser.add_argument("--iterations", type=int, default=50)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required")
diff --git a/smallm-ffn-megakernels-blackwell/benchmarks/benchmark.py b/smallm-ffn-megakernels-blackwell/benchmarks/benchmark.py
index 1465db7..1fe3702 100644
--- a/smallm-ffn-megakernels-blackwell/benchmarks/benchmark.py
+++ b/smallm-ffn-megakernels-blackwell/benchmarks/benchmark.py
@@ -11,6 +11,17 @@
 sys.path.insert(0, str(PACKAGE / "tests"))
 from _source_loader import load_installed_ops, load_source_ops  # noqa: E402
 
+
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 F8 = torch.float8_e4m3fn
 
 
@@ -38,7 +49,9 @@ def main() -> int:
     parser.add_argument("--artifact")
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iterations", type=int, default=100)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     if args.backend == "installed":
         if not args.artifact:
             parser.error("--artifact is required for --backend installed")
diff --git a/speculative-draft-primitives/benchmarks/benchmark.py b/speculative-draft-primitives/benchmarks/benchmark.py
index 5b8408a..d192d16 100644
--- a/speculative-draft-primitives/benchmarks/benchmark.py
+++ b/speculative-draft-primitives/benchmarks/benchmark.py
@@ -15,6 +15,16 @@
 from test_speculative_draft_primitives import load_installed_ops, load_source_ops  # noqa: E402
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def time_us(fn, warmup: int, iters: int) -> float:
     for _ in range(warmup):
         fn()
@@ -33,7 +43,9 @@ def main() -> int:
     parser.add_argument("--mode", choices=["headline", "full"], default="headline")
     parser.add_argument("--warmup", type=int, default=50)
     parser.add_argument("--iters", type=int, default=200)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
     shapes = [(16, 32000), (16, 248320)] if args.mode == "headline" else [
diff --git a/transformer-layout-primitives/benchmarks/benchmark.py b/transformer-layout-primitives/benchmarks/benchmark.py
index 339ca00..04114a0 100644
--- a/transformer-layout-primitives/benchmarks/benchmark.py
+++ b/transformer-layout-primitives/benchmarks/benchmark.py
@@ -20,6 +20,16 @@
 )
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def load_ops(backend: str, artifact: str | None):
     if backend == "source":
         return load_source_ops()
@@ -53,7 +63,9 @@ def main() -> int:
     parser.add_argument("--mode", choices=["headline", "full"], default="headline")
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iters", type=int, default=100)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     ops = load_ops(args.backend, args.artifact)
 
     print("workload,shape,op,flashrt_us,torch_eager_us,speedup")
diff --git a/turboquant-kv/benchmarks/benchmark.py b/turboquant-kv/benchmarks/benchmark.py
index 661b31c..0d27068 100644
--- a/turboquant-kv/benchmarks/benchmark.py
+++ b/turboquant-kv/benchmarks/benchmark.py
@@ -15,6 +15,16 @@
 from test_turboquant_kv import load_installed_ops, load_source_ops, ref_unpack  # noqa: E402
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def bench(fn, warmup: int, iters: int) -> float:
     for _ in range(warmup):
         fn()
@@ -44,7 +54,9 @@ def main() -> int:
     parser.add_argument("--artifact", default=None)
     parser.add_argument("--warmup", type=int, default=100)
     parser.add_argument("--iters", type=int, default=1000)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required")
diff --git a/vl-transformer-primitives/benchmarks/benchmark.py b/vl-transformer-primitives/benchmarks/benchmark.py
index cdd5fe4..2e3eaf8 100644
--- a/vl-transformer-primitives/benchmarks/benchmark.py
+++ b/vl-transformer-primitives/benchmarks/benchmark.py
@@ -22,6 +22,16 @@
 )
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def bench(fn, warmup: int, iters: int) -> float:
     for _ in range(warmup):
         fn()
@@ -42,7 +52,9 @@ def main() -> int:
     parser.add_argument("--artifact", default=None)
     parser.add_argument("--warmup", type=int, default=50)
     parser.add_argument("--iters", type=int, default=500)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
 
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required")
diff --git a/weight-only-ffn/benchmarks/benchmark.py b/weight-only-ffn/benchmarks/benchmark.py
index 31f0309..0f6293b 100644
--- a/weight-only-ffn/benchmarks/benchmark.py
+++ b/weight-only-ffn/benchmarks/benchmark.py
@@ -15,6 +15,16 @@
 import torch.nn.functional as F
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 SHAPES = {
     "llm_m1": (1, 4096, 11008, 4096),
     "llm_m2": (2, 4096, 11008, 4096),
@@ -426,7 +436,9 @@ def main() -> int:
     parser.add_argument("--warmup", type=int, default=20)
     parser.add_argument("--iterations", type=int, default=100)
     parser.add_argument("--json-out")
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     module = load_module(args.backend, args.artifact)
     names = ["llm_m1"] if args.mode == "smoke" else list(SHAPES)
     rows = []
diff --git a/world-model-conv/benchmarks/benchmark.py b/world-model-conv/benchmarks/benchmark.py
index ea3ac01..d5f8bee 100644
--- a/world-model-conv/benchmarks/benchmark.py
+++ b/world-model-conv/benchmarks/benchmark.py
@@ -14,6 +14,16 @@
 from test_world_model_conv import load_installed_ops, load_source_ops, ref_conv  # noqa: E402
 
 
+def _apply_mem_cap(max_mem_gb: float = 30.0) -> None:
+    if not torch.cuda.is_available() or max_mem_gb <= 0:
+        return
+    total = torch.cuda.get_device_properties(0).total_memory
+    cap = int(max_mem_gb * 1024**3)
+    if total <= 0 or cap >= total:
+        return
+    torch.cuda.set_per_process_memory_fraction(cap / total)
+
+
 def bench(fn, warmup: int, iters: int) -> float:
     for _ in range(warmup):
         fn()
@@ -34,7 +44,9 @@ def main() -> int:
     parser.add_argument("--artifact", default=None)
     parser.add_argument("--warmup", type=int, default=50)
     parser.add_argument("--iters", type=int, default=500)
+    parser.add_argument("--max-mem-gb", type=float, default=30.0)
     args = parser.parse_args()
+    _apply_mem_cap(args.max_mem_gb)
     if not torch.cuda.is_available():
         raise RuntimeError("CUDA is required")
     torch.manual_seed(123)