feat: Thor (SM110) portable SIMT fallbacks, benchmark memory guard, and validation docs

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())

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}")

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`

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")

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))

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`.