liangsu9988 commited on Jun 13

Commit

513d68d

verified ·

1 Parent(s): f3eeba1

Uploaded using `kernel-builder`.

Browse files

Files changed (32) hide show

build/torch211-cxx11-cu128-x86_64-linux/__init__.py +11 -3
build/torch211-cxx11-cu128-x86_64-linux/{_minimaxai_msa_blackwell_cuda_58018d1.abi3.so → _minimaxai_msa_blackwell_cuda_25b0067.abi3.so} +2 -2
build/torch211-cxx11-cu128-x86_64-linux/_native.py +195 -1
build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu128-x86_64-linux/api_status.py +4 -3
build/torch211-cxx11-cu128-x86_64-linux/decode/topk_sparse.py +20 -0
build/torch211-cxx11-cu128-x86_64-linux/interface.py +38 -3
build/torch211-cxx11-cu128-x86_64-linux/metadata.json +8 -8
build/torch211-cxx11-cu130-x86_64-linux/__init__.py +11 -3
build/torch211-cxx11-cu130-x86_64-linux/{_minimaxai_msa_blackwell_cuda_58018d1.abi3.so → _minimaxai_msa_blackwell_cuda_25b0067.abi3.so} +2 -2
build/torch211-cxx11-cu130-x86_64-linux/_native.py +195 -1
build/torch211-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch211-cxx11-cu130-x86_64-linux/api_status.py +4 -3
build/torch211-cxx11-cu130-x86_64-linux/decode/topk_sparse.py +20 -0
build/torch211-cxx11-cu130-x86_64-linux/interface.py +38 -3
build/torch211-cxx11-cu130-x86_64-linux/metadata.json +8 -8
build/torch212-cxx11-cu130-x86_64-linux/__init__.py +11 -3
build/torch212-cxx11-cu130-x86_64-linux/{_minimaxai_msa_blackwell_cuda_58018d1.abi3.so → _minimaxai_msa_blackwell_cuda_25b0067.abi3.so} +2 -2
build/torch212-cxx11-cu130-x86_64-linux/_native.py +195 -1
build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu130-x86_64-linux/api_status.py +4 -3
build/torch212-cxx11-cu130-x86_64-linux/decode/topk_sparse.py +20 -0
build/torch212-cxx11-cu130-x86_64-linux/interface.py +38 -3
build/torch212-cxx11-cu130-x86_64-linux/metadata.json +8 -8
build/torch212-cxx11-cu132-x86_64-linux/__init__.py +11 -3
build/torch212-cxx11-cu132-x86_64-linux/{_minimaxai_msa_blackwell_cuda_58018d1.abi3.so → _minimaxai_msa_blackwell_cuda_25b0067.abi3.so} +2 -2
build/torch212-cxx11-cu132-x86_64-linux/_native.py +195 -1
build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
build/torch212-cxx11-cu132-x86_64-linux/api_status.py +4 -3
build/torch212-cxx11-cu132-x86_64-linux/decode/topk_sparse.py +20 -0
build/torch212-cxx11-cu132-x86_64-linux/interface.py +38 -3
build/torch212-cxx11-cu132-x86_64-linux/metadata.json +8 -8

build/torch211-cxx11-cu128-x86_64-linux/__init__.py CHANGED Viewed

@@ -7,9 +7,12 @@ MiniMax-M3 decode-sparse attention path used by FlashRT's MiniMax Spark runtime.
 Implementation status:
   * native CUDA helper: score -> top-k sparse block ids;
-  * Triton CUDA attention kernels: Blackwell-validated decode and prefill paths;
   * MiniMaxAI/msa compatibility wrappers for CSR prefill, decode, NVFP4 helpers,
-    and FP4 block-score fallback.
 The upstream MiniMaxAI/msa package is SM100-only; this package is the Blackwell
 extension path. The public API is Tensor-oriented and independent from
@@ -48,7 +51,11 @@ from .interface import (
     sparse_atten_nvfp4_kv_func,
     sparse_decode_atten_func,
 )
-from ._native import has_native_ops, native_topk_from_scores
 from .naive.flash_with_topk_idx import naive_flash_decode_with_topk_idx
 from .naive.topk_sparse import naive_flash_decode_with_gqa_share_sparse
 from .quantize import (
@@ -88,6 +95,7 @@ __all__ = [
     "flash_decode_with_gqa_share_sparse",
     "has_native_ops",
     "native_topk_from_scores",
     "naive_flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "get_cu_seqblocks",

 Implementation status:
   * native CUDA helper: score -> top-k sparse block ids;
+  * native CUDA tensor-core sparse decode for the MiniMax-M3 Blackwell shape;
+  * native CUDA FP4 block-score indexer;
+  * native CUDA helper: swizzled NVFP4 -> BF16 dequant for the W4A16 quality path;
+  * Blackwell-validated sparse prefill path;
   * MiniMaxAI/msa compatibility wrappers for CSR prefill, decode, NVFP4 helpers,
+    and FP4 block-score helpers.
 The upstream MiniMaxAI/msa package is SM100-only; this package is the Blackwell
 extension path. The public API is Tensor-oriented and independent from
     sparse_atten_nvfp4_kv_func,
     sparse_decode_atten_func,
 )
+from ._native import (
+    has_native_ops,
+    native_nvfp4_dequant_swizzled_to_bf16,
+    native_topk_from_scores,
+)
 from .naive.flash_with_topk_idx import naive_flash_decode_with_topk_idx
 from .naive.topk_sparse import naive_flash_decode_with_gqa_share_sparse
 from .quantize import (
     "flash_decode_with_gqa_share_sparse",
     "has_native_ops",
     "native_topk_from_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
     "naive_flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "get_cu_seqblocks",

build/torch211-cxx11-cu128-x86_64-linux/{_minimaxai_msa_blackwell_cuda_58018d1.abi3.so → _minimaxai_msa_blackwell_cuda_25b0067.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:bffdfdc7f5f7e61a7a828d7887f42ea2412a75e6a259d11699fc24cf0f69b2b5
-size 119536

 version https://git-lfs.github.com/spec/v1
+oid sha256:528883873873d9a482cbe93eb0ea0ec5be7f0eac009dd8b103a0609841c032d5
+size 446928

build/torch211-cxx11-cu128-x86_64-linux/_native.py CHANGED Viewed

@@ -62,4 +62,198 @@ def native_topk_from_scores(
     return out
-__all__ = ["has_native_ops", "native_topk_from_scores"]

     return out
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_decode_sparse_attn_mma_paged")
+    )
+    def _msa_decode_sparse_attn_mma_paged_fake(
+        q: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        seq_lens: torch.Tensor,
+        slot_ids: torch.Tensor,
+        topk_idx: torch.Tensor,
+        block_size: int,
+        sm_scale: float,
+        out: torch.Tensor,
+    ) -> None:
+        return None
+def native_decode_mma_supported(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    *,
+    block_size: int,
+) -> bool:
+    """Whether the tensor-core decode kernel applies to these inputs.
+    The mma kernel is specialized for the M3 MSA shape: head_dim 128, GQA
+    group (Hq/Hkv) 16, bf16, and a block size that is a multiple of 64.
+    """
+    if not _HAS_NATIVE_OPS:
+        return False
+    if q.dtype is not torch.bfloat16 or k_cache.dtype is not torch.bfloat16:
+        return False
+    if q.dim() != 3 or k_cache.dim() != 3:
+        return False
+    head_dim = int(q.shape[-1])
+    num_q_heads = int(q.shape[1])
+    num_kv_heads = int(k_cache.shape[1])
+    if head_dim != 128 or num_kv_heads <= 0:
+        return False
+    if num_q_heads % num_kv_heads != 0 or num_q_heads // num_kv_heads != 16:
+        return False
+    if int(block_size) % 64 != 0:
+        return False
+    return True
+def native_decode_sparse_attn_mma_paged(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    req_to_token: torch.Tensor,
+    seq_lens: torch.Tensor,
+    slot_ids: torch.Tensor,
+    block_size: int,
+    topk_idx: torch.Tensor,
+    sm_scale: float,
+) -> torch.Tensor:
+    """Paged tensor-core block-sparse GQA decode. Returns [B, Hq, D] bf16."""
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    out = torch.empty(
+        (int(q.shape[0]), int(q.shape[1]), int(q.shape[2])),
+        device=q.device,
+        dtype=q.dtype,
+    )
+    ops.msa_decode_sparse_attn_mma_paged(
+        q.contiguous(),
+        k_cache.contiguous(),
+        v_cache.contiguous(),
+        req_to_token.contiguous(),
+        seq_lens.contiguous(),
+        slot_ids.contiguous(),
+        topk_idx.contiguous(),
+        int(block_size),
+        float(sm_scale),
+        out,
+    )
+    return out
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_indexer_block_scores")
+    )
+    def _msa_indexer_block_scores_fake(
+        q: torch.Tensor,
+        k_pages: torch.Tensor,
+        batch_of_q: torch.Tensor,
+        cu_q: torch.Tensor,
+        cu_k: torch.Tensor,
+        cu_pages: torch.Tensor,
+        kv_indices: torch.Tensor,
+        causal: int,
+        scores: torch.Tensor,
+    ) -> None:
+        return None
+def native_indexer_block_scores(
+    q: torch.Tensor,
+    k_pages: torch.Tensor,
+    batch_of_q: torch.Tensor,
+    cu_q: torch.Tensor,
+    cu_k: torch.Tensor,
+    cu_pages: torch.Tensor,
+    kv_indices: torch.Tensor,
+    *,
+    max_blocks: int,
+    causal: bool,
+) -> torch.Tensor:
+    """Block-max QK indexer scores. q/k_pages are bf16 (already dequantized).
+    Returns [Hq, max_blocks, total_q] f32 (-inf where no key is visible).
+    """
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    total_q = int(q.shape[0])
+    hq = int(q.shape[1])
+    scores = torch.full(
+        (hq, int(max_blocks), total_q),
+        float("-inf"),
+        dtype=torch.float32,
+        device=q.device,
+    )
+    ops.msa_indexer_block_scores(
+        q.contiguous(),
+        k_pages.contiguous(),
+        batch_of_q.contiguous(),
+        cu_q.contiguous(),
+        cu_k.contiguous(),
+        cu_pages.contiguous(),
+        kv_indices.contiguous(),
+        1 if causal else 0,
+        scores,
+    )
+    return scores
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_nvfp4_dequant_swizzled_to_bf16")
+    )
+    def _msa_nvfp4_dequant_swizzled_to_bf16_fake(
+        packed: torch.Tensor,
+        scale_128x4: torch.Tensor,
+        global_scale: float,
+        out: torch.Tensor,
+    ) -> None:
+        return None
+def native_nvfp4_dequant_swizzled_to_bf16(
+    packed: torch.Tensor,
+    scale_128x4: torch.Tensor,
+    global_scale: torch.Tensor | float,
+) -> torch.Tensor:
+    """Expand swizzled NVFP4 weights/KV to dense BF16 with the native CUDA path."""
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    if packed.dtype is not torch.uint8:
+        raise TypeError("packed must be torch.uint8")
+    if scale_128x4.dtype is not torch.uint8:
+        raise TypeError("scale_128x4 must be torch.uint8")
+    original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
+    out = torch.empty(original_shape, device=packed.device, dtype=torch.bfloat16)
+    if isinstance(global_scale, torch.Tensor):
+        global_scale_value = float(global_scale.reshape(-1)[0].item())
+    else:
+        global_scale_value = float(global_scale)
+    ops.msa_nvfp4_dequant_swizzled_to_bf16(
+        packed.contiguous(),
+        scale_128x4.contiguous(),
+        global_scale_value,
+        out,
+    )
+    return out
+__all__ = [
+    "has_native_ops",
+    "native_topk_from_scores",
+    "native_decode_mma_supported",
+    "native_decode_sparse_attn_mma_paged",
+    "native_indexer_block_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
+]

build/torch211-cxx11-cu128-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _minimaxai_msa_blackwell_cuda_58018d1
-ops = torch.ops._minimaxai_msa_blackwell_cuda_58018d1
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_minimaxai_msa_blackwell_cuda_58018d1::{op_name}"

 import torch
+from . import _minimaxai_msa_blackwell_cuda_25b0067
+ops = torch.ops._minimaxai_msa_blackwell_cuda_25b0067
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_minimaxai_msa_blackwell_cuda_25b0067::{op_name}"

build/torch211-cxx11-cu128-x86_64-linux/api_status.py CHANGED Viewed

@@ -22,6 +22,7 @@ V1_AVAILABLE_FUNCTIONS = (
     "nvfp4_global_scale_from_amax",
     "has_native_ops",
     "native_topk_from_scores",
     "flash_decode_with_topk_idx",
     "flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
@@ -55,7 +56,7 @@ _OFFICIAL_API_STATUS = {
     "sparse_atten_nvfp4_kv_func": {
         "status": "available",
         "target": "blackwell-prefill-nvfp4",
-        "reason": "Official NVFP4 KV prefill API is available via NVFP4 dequantization followed by the Blackwell prefill wrapper.",
     },
     "sparse_decode_atten_func": {
         "status": "available",
@@ -70,12 +71,12 @@ _OFFICIAL_API_STATUS = {
     "fp4_indexer_block_scores": {
         "status": "available",
         "target": "blackwell-fp4-indexer",
-        "reason": "Official FP4 block-score API is available as a correctness-first PyTorch fallback.",
     },
     "build_k2q_csr": {
         "status": "available",
         "target": "v1",
-        "reason": "Torch CSR construction fallback is available for compatibility.",
     },
     "SparseK2qCsrBuilderSm100": {
         "status": "available",

     "nvfp4_global_scale_from_amax",
     "has_native_ops",
     "native_topk_from_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
     "flash_decode_with_topk_idx",
     "flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "sparse_atten_nvfp4_kv_func": {
         "status": "available",
         "target": "blackwell-prefill-nvfp4",
+        "reason": "Official NVFP4 KV prefill API is available via native swizzled NVFP4 dequantization followed by the Blackwell prefill wrapper.",
     },
     "sparse_decode_atten_func": {
         "status": "available",
     "fp4_indexer_block_scores": {
         "status": "available",
         "target": "blackwell-fp4-indexer",
+        "reason": "Official FP4 block-score API uses the native Blackwell block-score kernel in built artifacts, with a reference fallback in source-tree mode.",
     },
     "build_k2q_csr": {
         "status": "available",
         "target": "v1",
+        "reason": "CSR construction helper is available for the official prefill API.",
     },
     "SparseK2qCsrBuilderSm100": {
         "status": "available",

build/torch211-cxx11-cu128-x86_64-linux/decode/topk_sparse.py CHANGED Viewed

@@ -8,6 +8,10 @@ import triton.language as tl
 from .._compat import is_hip  # vendored: was `from sglang.srt.utils import is_hip`
 from ..common.utils import robust_allocator
 _is_hip = is_hip()
@@ -341,6 +345,22 @@ def flash_decode_with_gqa_share_sparse(
     # sm scale
     if sm_scale is None:
         sm_scale = head_dim**-0.5
     # Pick NUM_TOPK_CHUNKS so total grid ≈ TARGET_GRID. Same constraints as
     # flash_decode_with_topk_idx: must be power of 2 (Triton arange) and must
     # only depend on shape constants (so grid is fixed within a cuda graph).

 from .._compat import is_hip  # vendored: was `from sglang.srt.utils import is_hip`
+from .._native import (
+    native_decode_mma_supported,
+    native_decode_sparse_attn_mma_paged,
+)
 from ..common.utils import robust_allocator
 _is_hip = is_hip()
     # sm scale
     if sm_scale is None:
         sm_scale = head_dim**-0.5
+    # Fast path: native tensor-core decode for the M3 MSA shape. Falls through
+    # to the Triton kernel for other shapes, sink, or fp8 caches.
+    if sink is None and not is_fp8 and native_decode_mma_supported(
+        q, k_cache, block_size=block_size
+    ):
+        return native_decode_sparse_attn_mma_paged(
+            q,
+            k_cache,
+            v_cache,
+            req_to_token,
+            seq_lens,
+            slot_ids,
+            block_size,
+            topk_idx,
+            sm_scale,
+        )
     # Pick NUM_TOPK_CHUNKS so total grid ≈ TARGET_GRID. Same constraints as
     # flash_decode_with_topk_idx: must be power of 2 (Triton arange) and must
     # only depend on shape constants (so grid is fixed within a cuda graph).

build/torch211-cxx11-cu128-x86_64-linux/interface.py CHANGED Viewed

@@ -7,6 +7,11 @@ from typing import Optional
 import torch
 from .decode.topk_sparse import flash_decode_with_gqa_share_sparse
 from .prefill.topk_sparse import flash_prefill_with_gqa_share_sparse
 from .quantize import Nvfp4QuantizedTensor, dequantize_nvfp4_128x4_to_bf16
@@ -488,6 +493,12 @@ def _nvfp4_dequant_kv(
 ) -> torch.Tensor:
     if packed.dtype is not torch.uint8:
         return packed
     original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
     rows = 1
     for dim in original_shape[:-1]:
@@ -572,11 +583,12 @@ def fp4_indexer_block_scores(
     qo_offset: Optional[torch.Tensor] = None,
     scale_layout: str = "preordered_mma",
 ) -> torch.Tensor:
-    """Correctness-first FP4 block-score API for Blackwell.
     Returns the official shape `[Hq, ceil(max_seqlen_k / 128), total_q]`.
-    This fallback is intentionally simple and is used to keep the public API
-    callable on Blackwell while optimized CUTE block scoring remains SM100-only.
     """
     del qo_offset
@@ -599,6 +611,29 @@ def fp4_indexer_block_scores(
     q = _apply_public_fp4_scale(_unpack_e2m1_fp4(q_fp4), q_scale)
     k = _apply_public_fp4_scale(_unpack_e2m1_fp4(k_fp4), k_scale)
     gqa = hq // hkv
     batch = int(cu_seqlens_q.numel()) - 1
     cu_q = cu_seqlens_q.to("cpu", non_blocking=False).tolist()

 import torch
+from ._native import (
+    has_native_ops,
+    native_indexer_block_scores,
+    native_nvfp4_dequant_swizzled_to_bf16,
+)
 from .decode.topk_sparse import flash_decode_with_gqa_share_sparse
 from .prefill.topk_sparse import flash_prefill_with_gqa_share_sparse
 from .quantize import Nvfp4QuantizedTensor, dequantize_nvfp4_128x4_to_bf16
 ) -> torch.Tensor:
     if packed.dtype is not torch.uint8:
         return packed
+    if has_native_ops() and packed.is_cuda:
+        return native_nvfp4_dequant_swizzled_to_bf16(
+            packed,
+            scale_128x4,
+            global_scale,
+        )
     original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
     rows = 1
     for dim in original_shape[:-1]:
     qo_offset: Optional[torch.Tensor] = None,
     scale_layout: str = "preordered_mma",
 ) -> torch.Tensor:
+    """FP4 block-score API for Blackwell.
     Returns the official shape `[Hq, ceil(max_seqlen_k / 128), total_q]`.
+    Built artifacts route CUDA tensors through the native Blackwell block-score
+    kernel. Source-tree mode uses the reference loop below so the API remains
+    testable before the extension is built.
     """
     del qo_offset
     q = _apply_public_fp4_scale(_unpack_e2m1_fp4(q_fp4), q_scale)
     k = _apply_public_fp4_scale(_unpack_e2m1_fp4(k_fp4), k_scale)
+    # Native block-max QK kernel for the cubic scoring loop when available.
+    # Same dequant numerics as the reference (bf16 compute); falls back to the
+    # Python reference below in source-tree mode.
+    if has_native_ops() and q_fp4.is_cuda:
+        batch = int(cu_seqlens_q.numel()) - 1
+        q_lens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(torch.int64)
+        batch_of_q = torch.repeat_interleave(
+            torch.arange(batch, device=q_fp4.device, dtype=torch.int32),
+            q_lens,
+        )
+        return native_indexer_block_scores(
+            q.to(torch.bfloat16),
+            k.to(torch.bfloat16),
+            batch_of_q,
+            cu_seqlens_q.to(torch.int32),
+            cu_seqlens_k.to(torch.int32),
+            cu_page_offsets.to(torch.int32),
+            kv_indices.to(torch.int32),
+            max_blocks=max_blocks,
+            causal=causal,
+        )
     gqa = hq // hkv
     batch = int(cu_seqlens_q.numel()) - 1
     cu_q = cu_seqlens_q.to("cpu", non_blocking=False).tolist()

build/torch211-cxx11-cu128-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "minimaxai-msa-blackwell",
-  "id": "_minimaxai_msa_blackwell_cuda_58018d1",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -13,18 +13,18 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "bH4lC+HdMx2DgfiEWfvxhd7GvQBOt02ftSPAzKo1rE0=",
       "_compat.py": "5Wy4a96HVMD8omRNsGg8rQJ+0qCCeXiDANk0CuyyP0I=",
-      "_minimaxai_msa_blackwell_cuda_58018d1.abi3.so": "v/39x/X35hp6go14h/QuokEqdeaiWdEWmfwkzw9psrU=",
-      "_native.py": "mIDvfkcQVd3H9plDsodJelrIFD2XHiEQgk4uqiiRLLE=",
-      "_ops.py": "kuQPE8cc2FNEl3lIMfjaj074xh5YntsWFa8cYC5CLIQ=",
-      "api_status.py": "VhS7aJhSznFwcgIQAnzXQTTuLJ2WgtysWrhX7RTSXaM=",
       "common/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "common/utils.py": "qtwh+wZoySlsWWALvWGVsmcIHKkdZUwePxUZ42tfWxc=",
       "decode/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "decode/flash_with_topk_idx.py": "Pk70+HJ60IKuy4KJC6hPdMtD2BjwPfSMabFiIBODWiQ=",
-      "decode/topk_sparse.py": "PLwQnOyiZT44pCXcZXzXdlZSixfBeanTZsfbpMOjlpQ=",
-      "interface.py": "TKlKdIqTtzfmieaY5ci3GwhD1QZXqUAZcqfcfPmyeuM=",
       "minimaxai_msa_blackwell/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
       "naive/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "naive/flash_with_topk_idx.py": "K7WQsRdO9+H2sz7pZ6axnJwE1OmnF0egdKtftHuChig=",

 {
   "name": "minimaxai-msa-blackwell",
+  "id": "_minimaxai_msa_blackwell_cuda_25b0067",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "oPB7xQhjHd9+aLVJkEmJkZOXOwMEarmCdj3SBbhH2bM=",
       "_compat.py": "5Wy4a96HVMD8omRNsGg8rQJ+0qCCeXiDANk0CuyyP0I=",
+      "_minimaxai_msa_blackwell_cuda_25b0067.abi3.so": "UoiDhzhz2aSCy+k+sOoOxb5/DqwAndixA6BgmEHAMtU=",
+      "_native.py": "wTh66DiIKxsJ5m8fkAXKSCBYD3CVFVx4PFxceQiwvPY=",
+      "_ops.py": "ikU3Lg+kgswKktYjNSq5f4DmgyTR8dnxV8lH3rmOWvk=",
+      "api_status.py": "7oZzIkb4fuwkAz3r4Q1slHbOfaP3jwVB647Ok223ATk=",
       "common/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "common/utils.py": "qtwh+wZoySlsWWALvWGVsmcIHKkdZUwePxUZ42tfWxc=",
       "decode/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "decode/flash_with_topk_idx.py": "Pk70+HJ60IKuy4KJC6hPdMtD2BjwPfSMabFiIBODWiQ=",
+      "decode/topk_sparse.py": "uINcV9SP5sAf9jYPsPvyQ9lon8sajZ8Qn/yJ+2LA8NU=",
+      "interface.py": "jfhBHGyunFF+OdFycJYyptNNIAngKrXY67buTXAN66M=",
       "minimaxai_msa_blackwell/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
       "naive/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "naive/flash_with_topk_idx.py": "K7WQsRdO9+H2sz7pZ6axnJwE1OmnF0egdKtftHuChig=",

build/torch211-cxx11-cu130-x86_64-linux/__init__.py CHANGED Viewed

@@ -7,9 +7,12 @@ MiniMax-M3 decode-sparse attention path used by FlashRT's MiniMax Spark runtime.
 Implementation status:
   * native CUDA helper: score -> top-k sparse block ids;
-  * Triton CUDA attention kernels: Blackwell-validated decode and prefill paths;
   * MiniMaxAI/msa compatibility wrappers for CSR prefill, decode, NVFP4 helpers,
-    and FP4 block-score fallback.
 The upstream MiniMaxAI/msa package is SM100-only; this package is the Blackwell
 extension path. The public API is Tensor-oriented and independent from
@@ -48,7 +51,11 @@ from .interface import (
     sparse_atten_nvfp4_kv_func,
     sparse_decode_atten_func,
 )
-from ._native import has_native_ops, native_topk_from_scores
 from .naive.flash_with_topk_idx import naive_flash_decode_with_topk_idx
 from .naive.topk_sparse import naive_flash_decode_with_gqa_share_sparse
 from .quantize import (
@@ -88,6 +95,7 @@ __all__ = [
     "flash_decode_with_gqa_share_sparse",
     "has_native_ops",
     "native_topk_from_scores",
     "naive_flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "get_cu_seqblocks",

 Implementation status:
   * native CUDA helper: score -> top-k sparse block ids;
+  * native CUDA tensor-core sparse decode for the MiniMax-M3 Blackwell shape;
+  * native CUDA FP4 block-score indexer;
+  * native CUDA helper: swizzled NVFP4 -> BF16 dequant for the W4A16 quality path;
+  * Blackwell-validated sparse prefill path;
   * MiniMaxAI/msa compatibility wrappers for CSR prefill, decode, NVFP4 helpers,
+    and FP4 block-score helpers.
 The upstream MiniMaxAI/msa package is SM100-only; this package is the Blackwell
 extension path. The public API is Tensor-oriented and independent from
     sparse_atten_nvfp4_kv_func,
     sparse_decode_atten_func,
 )
+from ._native import (
+    has_native_ops,
+    native_nvfp4_dequant_swizzled_to_bf16,
+    native_topk_from_scores,
+)
 from .naive.flash_with_topk_idx import naive_flash_decode_with_topk_idx
 from .naive.topk_sparse import naive_flash_decode_with_gqa_share_sparse
 from .quantize import (
     "flash_decode_with_gqa_share_sparse",
     "has_native_ops",
     "native_topk_from_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
     "naive_flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "get_cu_seqblocks",

build/torch211-cxx11-cu130-x86_64-linux/{_minimaxai_msa_blackwell_cuda_58018d1.abi3.so → _minimaxai_msa_blackwell_cuda_25b0067.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:495a67916122bc24e1c4e4b994c40e9e0b32118061be58af0f76b8f225eee617
-size 145216

 version https://git-lfs.github.com/spec/v1
+oid sha256:a053c1fd69c4df3f94e566a7c4d88ba943f8d5cece714ade52789f6819a961ac
+size 729824

build/torch211-cxx11-cu130-x86_64-linux/_native.py CHANGED Viewed

@@ -62,4 +62,198 @@ def native_topk_from_scores(
     return out
-__all__ = ["has_native_ops", "native_topk_from_scores"]

     return out
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_decode_sparse_attn_mma_paged")
+    )
+    def _msa_decode_sparse_attn_mma_paged_fake(
+        q: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        seq_lens: torch.Tensor,
+        slot_ids: torch.Tensor,
+        topk_idx: torch.Tensor,
+        block_size: int,
+        sm_scale: float,
+        out: torch.Tensor,
+    ) -> None:
+        return None
+def native_decode_mma_supported(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    *,
+    block_size: int,
+) -> bool:
+    """Whether the tensor-core decode kernel applies to these inputs.
+    The mma kernel is specialized for the M3 MSA shape: head_dim 128, GQA
+    group (Hq/Hkv) 16, bf16, and a block size that is a multiple of 64.
+    """
+    if not _HAS_NATIVE_OPS:
+        return False
+    if q.dtype is not torch.bfloat16 or k_cache.dtype is not torch.bfloat16:
+        return False
+    if q.dim() != 3 or k_cache.dim() != 3:
+        return False
+    head_dim = int(q.shape[-1])
+    num_q_heads = int(q.shape[1])
+    num_kv_heads = int(k_cache.shape[1])
+    if head_dim != 128 or num_kv_heads <= 0:
+        return False
+    if num_q_heads % num_kv_heads != 0 or num_q_heads // num_kv_heads != 16:
+        return False
+    if int(block_size) % 64 != 0:
+        return False
+    return True
+def native_decode_sparse_attn_mma_paged(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    req_to_token: torch.Tensor,
+    seq_lens: torch.Tensor,
+    slot_ids: torch.Tensor,
+    block_size: int,
+    topk_idx: torch.Tensor,
+    sm_scale: float,
+) -> torch.Tensor:
+    """Paged tensor-core block-sparse GQA decode. Returns [B, Hq, D] bf16."""
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    out = torch.empty(
+        (int(q.shape[0]), int(q.shape[1]), int(q.shape[2])),
+        device=q.device,
+        dtype=q.dtype,
+    )
+    ops.msa_decode_sparse_attn_mma_paged(
+        q.contiguous(),
+        k_cache.contiguous(),
+        v_cache.contiguous(),
+        req_to_token.contiguous(),
+        seq_lens.contiguous(),
+        slot_ids.contiguous(),
+        topk_idx.contiguous(),
+        int(block_size),
+        float(sm_scale),
+        out,
+    )
+    return out
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_indexer_block_scores")
+    )
+    def _msa_indexer_block_scores_fake(
+        q: torch.Tensor,
+        k_pages: torch.Tensor,
+        batch_of_q: torch.Tensor,
+        cu_q: torch.Tensor,
+        cu_k: torch.Tensor,
+        cu_pages: torch.Tensor,
+        kv_indices: torch.Tensor,
+        causal: int,
+        scores: torch.Tensor,
+    ) -> None:
+        return None
+def native_indexer_block_scores(
+    q: torch.Tensor,
+    k_pages: torch.Tensor,
+    batch_of_q: torch.Tensor,
+    cu_q: torch.Tensor,
+    cu_k: torch.Tensor,
+    cu_pages: torch.Tensor,
+    kv_indices: torch.Tensor,
+    *,
+    max_blocks: int,
+    causal: bool,
+) -> torch.Tensor:
+    """Block-max QK indexer scores. q/k_pages are bf16 (already dequantized).
+    Returns [Hq, max_blocks, total_q] f32 (-inf where no key is visible).
+    """
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    total_q = int(q.shape[0])
+    hq = int(q.shape[1])
+    scores = torch.full(
+        (hq, int(max_blocks), total_q),
+        float("-inf"),
+        dtype=torch.float32,
+        device=q.device,
+    )
+    ops.msa_indexer_block_scores(
+        q.contiguous(),
+        k_pages.contiguous(),
+        batch_of_q.contiguous(),
+        cu_q.contiguous(),
+        cu_k.contiguous(),
+        cu_pages.contiguous(),
+        kv_indices.contiguous(),
+        1 if causal else 0,
+        scores,
+    )
+    return scores
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_nvfp4_dequant_swizzled_to_bf16")
+    )
+    def _msa_nvfp4_dequant_swizzled_to_bf16_fake(
+        packed: torch.Tensor,
+        scale_128x4: torch.Tensor,
+        global_scale: float,
+        out: torch.Tensor,
+    ) -> None:
+        return None
+def native_nvfp4_dequant_swizzled_to_bf16(
+    packed: torch.Tensor,
+    scale_128x4: torch.Tensor,
+    global_scale: torch.Tensor | float,
+) -> torch.Tensor:
+    """Expand swizzled NVFP4 weights/KV to dense BF16 with the native CUDA path."""
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    if packed.dtype is not torch.uint8:
+        raise TypeError("packed must be torch.uint8")
+    if scale_128x4.dtype is not torch.uint8:
+        raise TypeError("scale_128x4 must be torch.uint8")
+    original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
+    out = torch.empty(original_shape, device=packed.device, dtype=torch.bfloat16)
+    if isinstance(global_scale, torch.Tensor):
+        global_scale_value = float(global_scale.reshape(-1)[0].item())
+    else:
+        global_scale_value = float(global_scale)
+    ops.msa_nvfp4_dequant_swizzled_to_bf16(
+        packed.contiguous(),
+        scale_128x4.contiguous(),
+        global_scale_value,
+        out,
+    )
+    return out
+__all__ = [
+    "has_native_ops",
+    "native_topk_from_scores",
+    "native_decode_mma_supported",
+    "native_decode_sparse_attn_mma_paged",
+    "native_indexer_block_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
+]

build/torch211-cxx11-cu130-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _minimaxai_msa_blackwell_cuda_58018d1
-ops = torch.ops._minimaxai_msa_blackwell_cuda_58018d1
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_minimaxai_msa_blackwell_cuda_58018d1::{op_name}"

 import torch
+from . import _minimaxai_msa_blackwell_cuda_25b0067
+ops = torch.ops._minimaxai_msa_blackwell_cuda_25b0067
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_minimaxai_msa_blackwell_cuda_25b0067::{op_name}"

build/torch211-cxx11-cu130-x86_64-linux/api_status.py CHANGED Viewed

@@ -22,6 +22,7 @@ V1_AVAILABLE_FUNCTIONS = (
     "nvfp4_global_scale_from_amax",
     "has_native_ops",
     "native_topk_from_scores",
     "flash_decode_with_topk_idx",
     "flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
@@ -55,7 +56,7 @@ _OFFICIAL_API_STATUS = {
     "sparse_atten_nvfp4_kv_func": {
         "status": "available",
         "target": "blackwell-prefill-nvfp4",
-        "reason": "Official NVFP4 KV prefill API is available via NVFP4 dequantization followed by the Blackwell prefill wrapper.",
     },
     "sparse_decode_atten_func": {
         "status": "available",
@@ -70,12 +71,12 @@ _OFFICIAL_API_STATUS = {
     "fp4_indexer_block_scores": {
         "status": "available",
         "target": "blackwell-fp4-indexer",
-        "reason": "Official FP4 block-score API is available as a correctness-first PyTorch fallback.",
     },
     "build_k2q_csr": {
         "status": "available",
         "target": "v1",
-        "reason": "Torch CSR construction fallback is available for compatibility.",
     },
     "SparseK2qCsrBuilderSm100": {
         "status": "available",

     "nvfp4_global_scale_from_amax",
     "has_native_ops",
     "native_topk_from_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
     "flash_decode_with_topk_idx",
     "flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "sparse_atten_nvfp4_kv_func": {
         "status": "available",
         "target": "blackwell-prefill-nvfp4",
+        "reason": "Official NVFP4 KV prefill API is available via native swizzled NVFP4 dequantization followed by the Blackwell prefill wrapper.",
     },
     "sparse_decode_atten_func": {
         "status": "available",
     "fp4_indexer_block_scores": {
         "status": "available",
         "target": "blackwell-fp4-indexer",
+        "reason": "Official FP4 block-score API uses the native Blackwell block-score kernel in built artifacts, with a reference fallback in source-tree mode.",
     },
     "build_k2q_csr": {
         "status": "available",
         "target": "v1",
+        "reason": "CSR construction helper is available for the official prefill API.",
     },
     "SparseK2qCsrBuilderSm100": {
         "status": "available",

build/torch211-cxx11-cu130-x86_64-linux/decode/topk_sparse.py CHANGED Viewed

@@ -8,6 +8,10 @@ import triton.language as tl
 from .._compat import is_hip  # vendored: was `from sglang.srt.utils import is_hip`
 from ..common.utils import robust_allocator
 _is_hip = is_hip()
@@ -341,6 +345,22 @@ def flash_decode_with_gqa_share_sparse(
     # sm scale
     if sm_scale is None:
         sm_scale = head_dim**-0.5
     # Pick NUM_TOPK_CHUNKS so total grid ≈ TARGET_GRID. Same constraints as
     # flash_decode_with_topk_idx: must be power of 2 (Triton arange) and must
     # only depend on shape constants (so grid is fixed within a cuda graph).

 from .._compat import is_hip  # vendored: was `from sglang.srt.utils import is_hip`
+from .._native import (
+    native_decode_mma_supported,
+    native_decode_sparse_attn_mma_paged,
+)
 from ..common.utils import robust_allocator
 _is_hip = is_hip()
     # sm scale
     if sm_scale is None:
         sm_scale = head_dim**-0.5
+    # Fast path: native tensor-core decode for the M3 MSA shape. Falls through
+    # to the Triton kernel for other shapes, sink, or fp8 caches.
+    if sink is None and not is_fp8 and native_decode_mma_supported(
+        q, k_cache, block_size=block_size
+    ):
+        return native_decode_sparse_attn_mma_paged(
+            q,
+            k_cache,
+            v_cache,
+            req_to_token,
+            seq_lens,
+            slot_ids,
+            block_size,
+            topk_idx,
+            sm_scale,
+        )
     # Pick NUM_TOPK_CHUNKS so total grid ≈ TARGET_GRID. Same constraints as
     # flash_decode_with_topk_idx: must be power of 2 (Triton arange) and must
     # only depend on shape constants (so grid is fixed within a cuda graph).

build/torch211-cxx11-cu130-x86_64-linux/interface.py CHANGED Viewed

@@ -7,6 +7,11 @@ from typing import Optional
 import torch
 from .decode.topk_sparse import flash_decode_with_gqa_share_sparse
 from .prefill.topk_sparse import flash_prefill_with_gqa_share_sparse
 from .quantize import Nvfp4QuantizedTensor, dequantize_nvfp4_128x4_to_bf16
@@ -488,6 +493,12 @@ def _nvfp4_dequant_kv(
 ) -> torch.Tensor:
     if packed.dtype is not torch.uint8:
         return packed
     original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
     rows = 1
     for dim in original_shape[:-1]:
@@ -572,11 +583,12 @@ def fp4_indexer_block_scores(
     qo_offset: Optional[torch.Tensor] = None,
     scale_layout: str = "preordered_mma",
 ) -> torch.Tensor:
-    """Correctness-first FP4 block-score API for Blackwell.
     Returns the official shape `[Hq, ceil(max_seqlen_k / 128), total_q]`.
-    This fallback is intentionally simple and is used to keep the public API
-    callable on Blackwell while optimized CUTE block scoring remains SM100-only.
     """
     del qo_offset
@@ -599,6 +611,29 @@ def fp4_indexer_block_scores(
     q = _apply_public_fp4_scale(_unpack_e2m1_fp4(q_fp4), q_scale)
     k = _apply_public_fp4_scale(_unpack_e2m1_fp4(k_fp4), k_scale)
     gqa = hq // hkv
     batch = int(cu_seqlens_q.numel()) - 1
     cu_q = cu_seqlens_q.to("cpu", non_blocking=False).tolist()

 import torch
+from ._native import (
+    has_native_ops,
+    native_indexer_block_scores,
+    native_nvfp4_dequant_swizzled_to_bf16,
+)
 from .decode.topk_sparse import flash_decode_with_gqa_share_sparse
 from .prefill.topk_sparse import flash_prefill_with_gqa_share_sparse
 from .quantize import Nvfp4QuantizedTensor, dequantize_nvfp4_128x4_to_bf16
 ) -> torch.Tensor:
     if packed.dtype is not torch.uint8:
         return packed
+    if has_native_ops() and packed.is_cuda:
+        return native_nvfp4_dequant_swizzled_to_bf16(
+            packed,
+            scale_128x4,
+            global_scale,
+        )
     original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
     rows = 1
     for dim in original_shape[:-1]:
     qo_offset: Optional[torch.Tensor] = None,
     scale_layout: str = "preordered_mma",
 ) -> torch.Tensor:
+    """FP4 block-score API for Blackwell.
     Returns the official shape `[Hq, ceil(max_seqlen_k / 128), total_q]`.
+    Built artifacts route CUDA tensors through the native Blackwell block-score
+    kernel. Source-tree mode uses the reference loop below so the API remains
+    testable before the extension is built.
     """
     del qo_offset
     q = _apply_public_fp4_scale(_unpack_e2m1_fp4(q_fp4), q_scale)
     k = _apply_public_fp4_scale(_unpack_e2m1_fp4(k_fp4), k_scale)
+    # Native block-max QK kernel for the cubic scoring loop when available.
+    # Same dequant numerics as the reference (bf16 compute); falls back to the
+    # Python reference below in source-tree mode.
+    if has_native_ops() and q_fp4.is_cuda:
+        batch = int(cu_seqlens_q.numel()) - 1
+        q_lens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(torch.int64)
+        batch_of_q = torch.repeat_interleave(
+            torch.arange(batch, device=q_fp4.device, dtype=torch.int32),
+            q_lens,
+        )
+        return native_indexer_block_scores(
+            q.to(torch.bfloat16),
+            k.to(torch.bfloat16),
+            batch_of_q,
+            cu_seqlens_q.to(torch.int32),
+            cu_seqlens_k.to(torch.int32),
+            cu_page_offsets.to(torch.int32),
+            kv_indices.to(torch.int32),
+            max_blocks=max_blocks,
+            causal=causal,
+        )
     gqa = hq // hkv
     batch = int(cu_seqlens_q.numel()) - 1
     cu_q = cu_seqlens_q.to("cpu", non_blocking=False).tolist()

build/torch211-cxx11-cu130-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "minimaxai-msa-blackwell",
-  "id": "_minimaxai_msa_blackwell_cuda_58018d1",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,18 +14,18 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "bH4lC+HdMx2DgfiEWfvxhd7GvQBOt02ftSPAzKo1rE0=",
       "_compat.py": "5Wy4a96HVMD8omRNsGg8rQJ+0qCCeXiDANk0CuyyP0I=",
-      "_minimaxai_msa_blackwell_cuda_58018d1.abi3.so": "SVpnkWEivCThxOS5lMQOngsyEYBhvlivD3a48iXu5hc=",
-      "_native.py": "mIDvfkcQVd3H9plDsodJelrIFD2XHiEQgk4uqiiRLLE=",
-      "_ops.py": "kuQPE8cc2FNEl3lIMfjaj074xh5YntsWFa8cYC5CLIQ=",
-      "api_status.py": "VhS7aJhSznFwcgIQAnzXQTTuLJ2WgtysWrhX7RTSXaM=",
       "common/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "common/utils.py": "qtwh+wZoySlsWWALvWGVsmcIHKkdZUwePxUZ42tfWxc=",
       "decode/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "decode/flash_with_topk_idx.py": "Pk70+HJ60IKuy4KJC6hPdMtD2BjwPfSMabFiIBODWiQ=",
-      "decode/topk_sparse.py": "PLwQnOyiZT44pCXcZXzXdlZSixfBeanTZsfbpMOjlpQ=",
-      "interface.py": "TKlKdIqTtzfmieaY5ci3GwhD1QZXqUAZcqfcfPmyeuM=",
       "minimaxai_msa_blackwell/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
       "naive/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "naive/flash_with_topk_idx.py": "K7WQsRdO9+H2sz7pZ6axnJwE1OmnF0egdKtftHuChig=",

 {
   "name": "minimaxai-msa-blackwell",
+  "id": "_minimaxai_msa_blackwell_cuda_25b0067",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "oPB7xQhjHd9+aLVJkEmJkZOXOwMEarmCdj3SBbhH2bM=",
       "_compat.py": "5Wy4a96HVMD8omRNsGg8rQJ+0qCCeXiDANk0CuyyP0I=",
+      "_minimaxai_msa_blackwell_cuda_25b0067.abi3.so": "oFPB/WnE3z+U5WanxNiLqUP41c7OcUreUnifaBmpYaw=",
+      "_native.py": "wTh66DiIKxsJ5m8fkAXKSCBYD3CVFVx4PFxceQiwvPY=",
+      "_ops.py": "ikU3Lg+kgswKktYjNSq5f4DmgyTR8dnxV8lH3rmOWvk=",
+      "api_status.py": "7oZzIkb4fuwkAz3r4Q1slHbOfaP3jwVB647Ok223ATk=",
       "common/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "common/utils.py": "qtwh+wZoySlsWWALvWGVsmcIHKkdZUwePxUZ42tfWxc=",
       "decode/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "decode/flash_with_topk_idx.py": "Pk70+HJ60IKuy4KJC6hPdMtD2BjwPfSMabFiIBODWiQ=",
+      "decode/topk_sparse.py": "uINcV9SP5sAf9jYPsPvyQ9lon8sajZ8Qn/yJ+2LA8NU=",
+      "interface.py": "jfhBHGyunFF+OdFycJYyptNNIAngKrXY67buTXAN66M=",
       "minimaxai_msa_blackwell/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
       "naive/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "naive/flash_with_topk_idx.py": "K7WQsRdO9+H2sz7pZ6axnJwE1OmnF0egdKtftHuChig=",

build/torch212-cxx11-cu130-x86_64-linux/__init__.py CHANGED Viewed

@@ -7,9 +7,12 @@ MiniMax-M3 decode-sparse attention path used by FlashRT's MiniMax Spark runtime.
 Implementation status:
   * native CUDA helper: score -> top-k sparse block ids;
-  * Triton CUDA attention kernels: Blackwell-validated decode and prefill paths;
   * MiniMaxAI/msa compatibility wrappers for CSR prefill, decode, NVFP4 helpers,
-    and FP4 block-score fallback.
 The upstream MiniMaxAI/msa package is SM100-only; this package is the Blackwell
 extension path. The public API is Tensor-oriented and independent from
@@ -48,7 +51,11 @@ from .interface import (
     sparse_atten_nvfp4_kv_func,
     sparse_decode_atten_func,
 )
-from ._native import has_native_ops, native_topk_from_scores
 from .naive.flash_with_topk_idx import naive_flash_decode_with_topk_idx
 from .naive.topk_sparse import naive_flash_decode_with_gqa_share_sparse
 from .quantize import (
@@ -88,6 +95,7 @@ __all__ = [
     "flash_decode_with_gqa_share_sparse",
     "has_native_ops",
     "native_topk_from_scores",
     "naive_flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "get_cu_seqblocks",

 Implementation status:
   * native CUDA helper: score -> top-k sparse block ids;
+  * native CUDA tensor-core sparse decode for the MiniMax-M3 Blackwell shape;
+  * native CUDA FP4 block-score indexer;
+  * native CUDA helper: swizzled NVFP4 -> BF16 dequant for the W4A16 quality path;
+  * Blackwell-validated sparse prefill path;
   * MiniMaxAI/msa compatibility wrappers for CSR prefill, decode, NVFP4 helpers,
+    and FP4 block-score helpers.
 The upstream MiniMaxAI/msa package is SM100-only; this package is the Blackwell
 extension path. The public API is Tensor-oriented and independent from
     sparse_atten_nvfp4_kv_func,
     sparse_decode_atten_func,
 )
+from ._native import (
+    has_native_ops,
+    native_nvfp4_dequant_swizzled_to_bf16,
+    native_topk_from_scores,
+)
 from .naive.flash_with_topk_idx import naive_flash_decode_with_topk_idx
 from .naive.topk_sparse import naive_flash_decode_with_gqa_share_sparse
 from .quantize import (
     "flash_decode_with_gqa_share_sparse",
     "has_native_ops",
     "native_topk_from_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
     "naive_flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "get_cu_seqblocks",

build/torch212-cxx11-cu130-x86_64-linux/{_minimaxai_msa_blackwell_cuda_58018d1.abi3.so → _minimaxai_msa_blackwell_cuda_25b0067.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:473074c02a1f88aa83de41429947330435156e125b9b4ac1a8b72fabfa49fd66
-size 151920

 version https://git-lfs.github.com/spec/v1
+oid sha256:2ea862050a19dcc4b0c7cbfd3b519112f1b4e054cfc182a14f92876b4ae5b9f7
+size 735888

build/torch212-cxx11-cu130-x86_64-linux/_native.py CHANGED Viewed

@@ -62,4 +62,198 @@ def native_topk_from_scores(
     return out
-__all__ = ["has_native_ops", "native_topk_from_scores"]

     return out
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_decode_sparse_attn_mma_paged")
+    )
+    def _msa_decode_sparse_attn_mma_paged_fake(
+        q: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        seq_lens: torch.Tensor,
+        slot_ids: torch.Tensor,
+        topk_idx: torch.Tensor,
+        block_size: int,
+        sm_scale: float,
+        out: torch.Tensor,
+    ) -> None:
+        return None
+def native_decode_mma_supported(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    *,
+    block_size: int,
+) -> bool:
+    """Whether the tensor-core decode kernel applies to these inputs.
+    The mma kernel is specialized for the M3 MSA shape: head_dim 128, GQA
+    group (Hq/Hkv) 16, bf16, and a block size that is a multiple of 64.
+    """
+    if not _HAS_NATIVE_OPS:
+        return False
+    if q.dtype is not torch.bfloat16 or k_cache.dtype is not torch.bfloat16:
+        return False
+    if q.dim() != 3 or k_cache.dim() != 3:
+        return False
+    head_dim = int(q.shape[-1])
+    num_q_heads = int(q.shape[1])
+    num_kv_heads = int(k_cache.shape[1])
+    if head_dim != 128 or num_kv_heads <= 0:
+        return False
+    if num_q_heads % num_kv_heads != 0 or num_q_heads // num_kv_heads != 16:
+        return False
+    if int(block_size) % 64 != 0:
+        return False
+    return True
+def native_decode_sparse_attn_mma_paged(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    req_to_token: torch.Tensor,
+    seq_lens: torch.Tensor,
+    slot_ids: torch.Tensor,
+    block_size: int,
+    topk_idx: torch.Tensor,
+    sm_scale: float,
+) -> torch.Tensor:
+    """Paged tensor-core block-sparse GQA decode. Returns [B, Hq, D] bf16."""
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    out = torch.empty(
+        (int(q.shape[0]), int(q.shape[1]), int(q.shape[2])),
+        device=q.device,
+        dtype=q.dtype,
+    )
+    ops.msa_decode_sparse_attn_mma_paged(
+        q.contiguous(),
+        k_cache.contiguous(),
+        v_cache.contiguous(),
+        req_to_token.contiguous(),
+        seq_lens.contiguous(),
+        slot_ids.contiguous(),
+        topk_idx.contiguous(),
+        int(block_size),
+        float(sm_scale),
+        out,
+    )
+    return out
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_indexer_block_scores")
+    )
+    def _msa_indexer_block_scores_fake(
+        q: torch.Tensor,
+        k_pages: torch.Tensor,
+        batch_of_q: torch.Tensor,
+        cu_q: torch.Tensor,
+        cu_k: torch.Tensor,
+        cu_pages: torch.Tensor,
+        kv_indices: torch.Tensor,
+        causal: int,
+        scores: torch.Tensor,
+    ) -> None:
+        return None
+def native_indexer_block_scores(
+    q: torch.Tensor,
+    k_pages: torch.Tensor,
+    batch_of_q: torch.Tensor,
+    cu_q: torch.Tensor,
+    cu_k: torch.Tensor,
+    cu_pages: torch.Tensor,
+    kv_indices: torch.Tensor,
+    *,
+    max_blocks: int,
+    causal: bool,
+) -> torch.Tensor:
+    """Block-max QK indexer scores. q/k_pages are bf16 (already dequantized).
+    Returns [Hq, max_blocks, total_q] f32 (-inf where no key is visible).
+    """
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    total_q = int(q.shape[0])
+    hq = int(q.shape[1])
+    scores = torch.full(
+        (hq, int(max_blocks), total_q),
+        float("-inf"),
+        dtype=torch.float32,
+        device=q.device,
+    )
+    ops.msa_indexer_block_scores(
+        q.contiguous(),
+        k_pages.contiguous(),
+        batch_of_q.contiguous(),
+        cu_q.contiguous(),
+        cu_k.contiguous(),
+        cu_pages.contiguous(),
+        kv_indices.contiguous(),
+        1 if causal else 0,
+        scores,
+    )
+    return scores
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_nvfp4_dequant_swizzled_to_bf16")
+    )
+    def _msa_nvfp4_dequant_swizzled_to_bf16_fake(
+        packed: torch.Tensor,
+        scale_128x4: torch.Tensor,
+        global_scale: float,
+        out: torch.Tensor,
+    ) -> None:
+        return None
+def native_nvfp4_dequant_swizzled_to_bf16(
+    packed: torch.Tensor,
+    scale_128x4: torch.Tensor,
+    global_scale: torch.Tensor | float,
+) -> torch.Tensor:
+    """Expand swizzled NVFP4 weights/KV to dense BF16 with the native CUDA path."""
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    if packed.dtype is not torch.uint8:
+        raise TypeError("packed must be torch.uint8")
+    if scale_128x4.dtype is not torch.uint8:
+        raise TypeError("scale_128x4 must be torch.uint8")
+    original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
+    out = torch.empty(original_shape, device=packed.device, dtype=torch.bfloat16)
+    if isinstance(global_scale, torch.Tensor):
+        global_scale_value = float(global_scale.reshape(-1)[0].item())
+    else:
+        global_scale_value = float(global_scale)
+    ops.msa_nvfp4_dequant_swizzled_to_bf16(
+        packed.contiguous(),
+        scale_128x4.contiguous(),
+        global_scale_value,
+        out,
+    )
+    return out
+__all__ = [
+    "has_native_ops",
+    "native_topk_from_scores",
+    "native_decode_mma_supported",
+    "native_decode_sparse_attn_mma_paged",
+    "native_indexer_block_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
+]

build/torch212-cxx11-cu130-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _minimaxai_msa_blackwell_cuda_58018d1
-ops = torch.ops._minimaxai_msa_blackwell_cuda_58018d1
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_minimaxai_msa_blackwell_cuda_58018d1::{op_name}"

 import torch
+from . import _minimaxai_msa_blackwell_cuda_25b0067
+ops = torch.ops._minimaxai_msa_blackwell_cuda_25b0067
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_minimaxai_msa_blackwell_cuda_25b0067::{op_name}"

build/torch212-cxx11-cu130-x86_64-linux/api_status.py CHANGED Viewed

@@ -22,6 +22,7 @@ V1_AVAILABLE_FUNCTIONS = (
     "nvfp4_global_scale_from_amax",
     "has_native_ops",
     "native_topk_from_scores",
     "flash_decode_with_topk_idx",
     "flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
@@ -55,7 +56,7 @@ _OFFICIAL_API_STATUS = {
     "sparse_atten_nvfp4_kv_func": {
         "status": "available",
         "target": "blackwell-prefill-nvfp4",
-        "reason": "Official NVFP4 KV prefill API is available via NVFP4 dequantization followed by the Blackwell prefill wrapper.",
     },
     "sparse_decode_atten_func": {
         "status": "available",
@@ -70,12 +71,12 @@ _OFFICIAL_API_STATUS = {
     "fp4_indexer_block_scores": {
         "status": "available",
         "target": "blackwell-fp4-indexer",
-        "reason": "Official FP4 block-score API is available as a correctness-first PyTorch fallback.",
     },
     "build_k2q_csr": {
         "status": "available",
         "target": "v1",
-        "reason": "Torch CSR construction fallback is available for compatibility.",
     },
     "SparseK2qCsrBuilderSm100": {
         "status": "available",

     "nvfp4_global_scale_from_amax",
     "has_native_ops",
     "native_topk_from_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
     "flash_decode_with_topk_idx",
     "flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "sparse_atten_nvfp4_kv_func": {
         "status": "available",
         "target": "blackwell-prefill-nvfp4",
+        "reason": "Official NVFP4 KV prefill API is available via native swizzled NVFP4 dequantization followed by the Blackwell prefill wrapper.",
     },
     "sparse_decode_atten_func": {
         "status": "available",
     "fp4_indexer_block_scores": {
         "status": "available",
         "target": "blackwell-fp4-indexer",
+        "reason": "Official FP4 block-score API uses the native Blackwell block-score kernel in built artifacts, with a reference fallback in source-tree mode.",
     },
     "build_k2q_csr": {
         "status": "available",
         "target": "v1",
+        "reason": "CSR construction helper is available for the official prefill API.",
     },
     "SparseK2qCsrBuilderSm100": {
         "status": "available",

build/torch212-cxx11-cu130-x86_64-linux/decode/topk_sparse.py CHANGED Viewed

@@ -8,6 +8,10 @@ import triton.language as tl
 from .._compat import is_hip  # vendored: was `from sglang.srt.utils import is_hip`
 from ..common.utils import robust_allocator
 _is_hip = is_hip()
@@ -341,6 +345,22 @@ def flash_decode_with_gqa_share_sparse(
     # sm scale
     if sm_scale is None:
         sm_scale = head_dim**-0.5
     # Pick NUM_TOPK_CHUNKS so total grid ≈ TARGET_GRID. Same constraints as
     # flash_decode_with_topk_idx: must be power of 2 (Triton arange) and must
     # only depend on shape constants (so grid is fixed within a cuda graph).

 from .._compat import is_hip  # vendored: was `from sglang.srt.utils import is_hip`
+from .._native import (
+    native_decode_mma_supported,
+    native_decode_sparse_attn_mma_paged,
+)
 from ..common.utils import robust_allocator
 _is_hip = is_hip()
     # sm scale
     if sm_scale is None:
         sm_scale = head_dim**-0.5
+    # Fast path: native tensor-core decode for the M3 MSA shape. Falls through
+    # to the Triton kernel for other shapes, sink, or fp8 caches.
+    if sink is None and not is_fp8 and native_decode_mma_supported(
+        q, k_cache, block_size=block_size
+    ):
+        return native_decode_sparse_attn_mma_paged(
+            q,
+            k_cache,
+            v_cache,
+            req_to_token,
+            seq_lens,
+            slot_ids,
+            block_size,
+            topk_idx,
+            sm_scale,
+        )
     # Pick NUM_TOPK_CHUNKS so total grid ≈ TARGET_GRID. Same constraints as
     # flash_decode_with_topk_idx: must be power of 2 (Triton arange) and must
     # only depend on shape constants (so grid is fixed within a cuda graph).

build/torch212-cxx11-cu130-x86_64-linux/interface.py CHANGED Viewed

@@ -7,6 +7,11 @@ from typing import Optional
 import torch
 from .decode.topk_sparse import flash_decode_with_gqa_share_sparse
 from .prefill.topk_sparse import flash_prefill_with_gqa_share_sparse
 from .quantize import Nvfp4QuantizedTensor, dequantize_nvfp4_128x4_to_bf16
@@ -488,6 +493,12 @@ def _nvfp4_dequant_kv(
 ) -> torch.Tensor:
     if packed.dtype is not torch.uint8:
         return packed
     original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
     rows = 1
     for dim in original_shape[:-1]:
@@ -572,11 +583,12 @@ def fp4_indexer_block_scores(
     qo_offset: Optional[torch.Tensor] = None,
     scale_layout: str = "preordered_mma",
 ) -> torch.Tensor:
-    """Correctness-first FP4 block-score API for Blackwell.
     Returns the official shape `[Hq, ceil(max_seqlen_k / 128), total_q]`.
-    This fallback is intentionally simple and is used to keep the public API
-    callable on Blackwell while optimized CUTE block scoring remains SM100-only.
     """
     del qo_offset
@@ -599,6 +611,29 @@ def fp4_indexer_block_scores(
     q = _apply_public_fp4_scale(_unpack_e2m1_fp4(q_fp4), q_scale)
     k = _apply_public_fp4_scale(_unpack_e2m1_fp4(k_fp4), k_scale)
     gqa = hq // hkv
     batch = int(cu_seqlens_q.numel()) - 1
     cu_q = cu_seqlens_q.to("cpu", non_blocking=False).tolist()

 import torch
+from ._native import (
+    has_native_ops,
+    native_indexer_block_scores,
+    native_nvfp4_dequant_swizzled_to_bf16,
+)
 from .decode.topk_sparse import flash_decode_with_gqa_share_sparse
 from .prefill.topk_sparse import flash_prefill_with_gqa_share_sparse
 from .quantize import Nvfp4QuantizedTensor, dequantize_nvfp4_128x4_to_bf16
 ) -> torch.Tensor:
     if packed.dtype is not torch.uint8:
         return packed
+    if has_native_ops() and packed.is_cuda:
+        return native_nvfp4_dequant_swizzled_to_bf16(
+            packed,
+            scale_128x4,
+            global_scale,
+        )
     original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
     rows = 1
     for dim in original_shape[:-1]:
     qo_offset: Optional[torch.Tensor] = None,
     scale_layout: str = "preordered_mma",
 ) -> torch.Tensor:
+    """FP4 block-score API for Blackwell.
     Returns the official shape `[Hq, ceil(max_seqlen_k / 128), total_q]`.
+    Built artifacts route CUDA tensors through the native Blackwell block-score
+    kernel. Source-tree mode uses the reference loop below so the API remains
+    testable before the extension is built.
     """
     del qo_offset
     q = _apply_public_fp4_scale(_unpack_e2m1_fp4(q_fp4), q_scale)
     k = _apply_public_fp4_scale(_unpack_e2m1_fp4(k_fp4), k_scale)
+    # Native block-max QK kernel for the cubic scoring loop when available.
+    # Same dequant numerics as the reference (bf16 compute); falls back to the
+    # Python reference below in source-tree mode.
+    if has_native_ops() and q_fp4.is_cuda:
+        batch = int(cu_seqlens_q.numel()) - 1
+        q_lens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(torch.int64)
+        batch_of_q = torch.repeat_interleave(
+            torch.arange(batch, device=q_fp4.device, dtype=torch.int32),
+            q_lens,
+        )
+        return native_indexer_block_scores(
+            q.to(torch.bfloat16),
+            k.to(torch.bfloat16),
+            batch_of_q,
+            cu_seqlens_q.to(torch.int32),
+            cu_seqlens_k.to(torch.int32),
+            cu_page_offsets.to(torch.int32),
+            kv_indices.to(torch.int32),
+            max_blocks=max_blocks,
+            causal=causal,
+        )
     gqa = hq // hkv
     batch = int(cu_seqlens_q.numel()) - 1
     cu_q = cu_seqlens_q.to("cpu", non_blocking=False).tolist()

build/torch212-cxx11-cu130-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "minimaxai-msa-blackwell",
-  "id": "_minimaxai_msa_blackwell_cuda_58018d1",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,18 +14,18 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "bH4lC+HdMx2DgfiEWfvxhd7GvQBOt02ftSPAzKo1rE0=",
       "_compat.py": "5Wy4a96HVMD8omRNsGg8rQJ+0qCCeXiDANk0CuyyP0I=",
-      "_minimaxai_msa_blackwell_cuda_58018d1.abi3.so": "RzB0wCofiKqD3kFCmUczBDUVbhJbm0rBqLcvq/pJ/WY=",
-      "_native.py": "mIDvfkcQVd3H9plDsodJelrIFD2XHiEQgk4uqiiRLLE=",
-      "_ops.py": "kuQPE8cc2FNEl3lIMfjaj074xh5YntsWFa8cYC5CLIQ=",
-      "api_status.py": "VhS7aJhSznFwcgIQAnzXQTTuLJ2WgtysWrhX7RTSXaM=",
       "common/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "common/utils.py": "qtwh+wZoySlsWWALvWGVsmcIHKkdZUwePxUZ42tfWxc=",
       "decode/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "decode/flash_with_topk_idx.py": "Pk70+HJ60IKuy4KJC6hPdMtD2BjwPfSMabFiIBODWiQ=",
-      "decode/topk_sparse.py": "PLwQnOyiZT44pCXcZXzXdlZSixfBeanTZsfbpMOjlpQ=",
-      "interface.py": "TKlKdIqTtzfmieaY5ci3GwhD1QZXqUAZcqfcfPmyeuM=",
       "minimaxai_msa_blackwell/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
       "naive/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "naive/flash_with_topk_idx.py": "K7WQsRdO9+H2sz7pZ6axnJwE1OmnF0egdKtftHuChig=",

 {
   "name": "minimaxai-msa-blackwell",
+  "id": "_minimaxai_msa_blackwell_cuda_25b0067",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "oPB7xQhjHd9+aLVJkEmJkZOXOwMEarmCdj3SBbhH2bM=",
       "_compat.py": "5Wy4a96HVMD8omRNsGg8rQJ+0qCCeXiDANk0CuyyP0I=",
+      "_minimaxai_msa_blackwell_cuda_25b0067.abi3.so": "LqhiBQoZ3MSwx8v9O1GREvG04FTPwYKhT5KHa0rlufc=",
+      "_native.py": "wTh66DiIKxsJ5m8fkAXKSCBYD3CVFVx4PFxceQiwvPY=",
+      "_ops.py": "ikU3Lg+kgswKktYjNSq5f4DmgyTR8dnxV8lH3rmOWvk=",
+      "api_status.py": "7oZzIkb4fuwkAz3r4Q1slHbOfaP3jwVB647Ok223ATk=",
       "common/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "common/utils.py": "qtwh+wZoySlsWWALvWGVsmcIHKkdZUwePxUZ42tfWxc=",
       "decode/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "decode/flash_with_topk_idx.py": "Pk70+HJ60IKuy4KJC6hPdMtD2BjwPfSMabFiIBODWiQ=",
+      "decode/topk_sparse.py": "uINcV9SP5sAf9jYPsPvyQ9lon8sajZ8Qn/yJ+2LA8NU=",
+      "interface.py": "jfhBHGyunFF+OdFycJYyptNNIAngKrXY67buTXAN66M=",
       "minimaxai_msa_blackwell/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
       "naive/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "naive/flash_with_topk_idx.py": "K7WQsRdO9+H2sz7pZ6axnJwE1OmnF0egdKtftHuChig=",

build/torch212-cxx11-cu132-x86_64-linux/__init__.py CHANGED Viewed

@@ -7,9 +7,12 @@ MiniMax-M3 decode-sparse attention path used by FlashRT's MiniMax Spark runtime.
 Implementation status:
   * native CUDA helper: score -> top-k sparse block ids;
-  * Triton CUDA attention kernels: Blackwell-validated decode and prefill paths;
   * MiniMaxAI/msa compatibility wrappers for CSR prefill, decode, NVFP4 helpers,
-    and FP4 block-score fallback.
 The upstream MiniMaxAI/msa package is SM100-only; this package is the Blackwell
 extension path. The public API is Tensor-oriented and independent from
@@ -48,7 +51,11 @@ from .interface import (
     sparse_atten_nvfp4_kv_func,
     sparse_decode_atten_func,
 )
-from ._native import has_native_ops, native_topk_from_scores
 from .naive.flash_with_topk_idx import naive_flash_decode_with_topk_idx
 from .naive.topk_sparse import naive_flash_decode_with_gqa_share_sparse
 from .quantize import (
@@ -88,6 +95,7 @@ __all__ = [
     "flash_decode_with_gqa_share_sparse",
     "has_native_ops",
     "native_topk_from_scores",
     "naive_flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "get_cu_seqblocks",

 Implementation status:
   * native CUDA helper: score -> top-k sparse block ids;
+  * native CUDA tensor-core sparse decode for the MiniMax-M3 Blackwell shape;
+  * native CUDA FP4 block-score indexer;
+  * native CUDA helper: swizzled NVFP4 -> BF16 dequant for the W4A16 quality path;
+  * Blackwell-validated sparse prefill path;
   * MiniMaxAI/msa compatibility wrappers for CSR prefill, decode, NVFP4 helpers,
+    and FP4 block-score helpers.
 The upstream MiniMaxAI/msa package is SM100-only; this package is the Blackwell
 extension path. The public API is Tensor-oriented and independent from
     sparse_atten_nvfp4_kv_func,
     sparse_decode_atten_func,
 )
+from ._native import (
+    has_native_ops,
+    native_nvfp4_dequant_swizzled_to_bf16,
+    native_topk_from_scores,
+)
 from .naive.flash_with_topk_idx import naive_flash_decode_with_topk_idx
 from .naive.topk_sparse import naive_flash_decode_with_gqa_share_sparse
 from .quantize import (
     "flash_decode_with_gqa_share_sparse",
     "has_native_ops",
     "native_topk_from_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
     "naive_flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "get_cu_seqblocks",

build/torch212-cxx11-cu132-x86_64-linux/{_minimaxai_msa_blackwell_cuda_58018d1.abi3.so → _minimaxai_msa_blackwell_cuda_25b0067.abi3.so} RENAMED Viewed

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3c2f06a4139531e1dffcede1973b0c0605afa3b62f9818a87504ed794adc874d
-size 164208

 version https://git-lfs.github.com/spec/v1
+oid sha256:4eb01c1a96c3e640474b3a52f2d38da4be2af0290c4b945c50f90e867d9e907f
+size 744080

build/torch212-cxx11-cu132-x86_64-linux/_native.py CHANGED Viewed

@@ -62,4 +62,198 @@ def native_topk_from_scores(
     return out
-__all__ = ["has_native_ops", "native_topk_from_scores"]

     return out
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_decode_sparse_attn_mma_paged")
+    )
+    def _msa_decode_sparse_attn_mma_paged_fake(
+        q: torch.Tensor,
+        k_cache: torch.Tensor,
+        v_cache: torch.Tensor,
+        req_to_token: torch.Tensor,
+        seq_lens: torch.Tensor,
+        slot_ids: torch.Tensor,
+        topk_idx: torch.Tensor,
+        block_size: int,
+        sm_scale: float,
+        out: torch.Tensor,
+    ) -> None:
+        return None
+def native_decode_mma_supported(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    *,
+    block_size: int,
+) -> bool:
+    """Whether the tensor-core decode kernel applies to these inputs.
+    The mma kernel is specialized for the M3 MSA shape: head_dim 128, GQA
+    group (Hq/Hkv) 16, bf16, and a block size that is a multiple of 64.
+    """
+    if not _HAS_NATIVE_OPS:
+        return False
+    if q.dtype is not torch.bfloat16 or k_cache.dtype is not torch.bfloat16:
+        return False
+    if q.dim() != 3 or k_cache.dim() != 3:
+        return False
+    head_dim = int(q.shape[-1])
+    num_q_heads = int(q.shape[1])
+    num_kv_heads = int(k_cache.shape[1])
+    if head_dim != 128 or num_kv_heads <= 0:
+        return False
+    if num_q_heads % num_kv_heads != 0 or num_q_heads // num_kv_heads != 16:
+        return False
+    if int(block_size) % 64 != 0:
+        return False
+    return True
+def native_decode_sparse_attn_mma_paged(
+    q: torch.Tensor,
+    k_cache: torch.Tensor,
+    v_cache: torch.Tensor,
+    req_to_token: torch.Tensor,
+    seq_lens: torch.Tensor,
+    slot_ids: torch.Tensor,
+    block_size: int,
+    topk_idx: torch.Tensor,
+    sm_scale: float,
+) -> torch.Tensor:
+    """Paged tensor-core block-sparse GQA decode. Returns [B, Hq, D] bf16."""
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    out = torch.empty(
+        (int(q.shape[0]), int(q.shape[1]), int(q.shape[2])),
+        device=q.device,
+        dtype=q.dtype,
+    )
+    ops.msa_decode_sparse_attn_mma_paged(
+        q.contiguous(),
+        k_cache.contiguous(),
+        v_cache.contiguous(),
+        req_to_token.contiguous(),
+        seq_lens.contiguous(),
+        slot_ids.contiguous(),
+        topk_idx.contiguous(),
+        int(block_size),
+        float(sm_scale),
+        out,
+    )
+    return out
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_indexer_block_scores")
+    )
+    def _msa_indexer_block_scores_fake(
+        q: torch.Tensor,
+        k_pages: torch.Tensor,
+        batch_of_q: torch.Tensor,
+        cu_q: torch.Tensor,
+        cu_k: torch.Tensor,
+        cu_pages: torch.Tensor,
+        kv_indices: torch.Tensor,
+        causal: int,
+        scores: torch.Tensor,
+    ) -> None:
+        return None
+def native_indexer_block_scores(
+    q: torch.Tensor,
+    k_pages: torch.Tensor,
+    batch_of_q: torch.Tensor,
+    cu_q: torch.Tensor,
+    cu_k: torch.Tensor,
+    cu_pages: torch.Tensor,
+    kv_indices: torch.Tensor,
+    *,
+    max_blocks: int,
+    causal: bool,
+) -> torch.Tensor:
+    """Block-max QK indexer scores. q/k_pages are bf16 (already dequantized).
+    Returns [Hq, max_blocks, total_q] f32 (-inf where no key is visible).
+    """
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    total_q = int(q.shape[0])
+    hq = int(q.shape[1])
+    scores = torch.full(
+        (hq, int(max_blocks), total_q),
+        float("-inf"),
+        dtype=torch.float32,
+        device=q.device,
+    )
+    ops.msa_indexer_block_scores(
+        q.contiguous(),
+        k_pages.contiguous(),
+        batch_of_q.contiguous(),
+        cu_q.contiguous(),
+        cu_k.contiguous(),
+        cu_pages.contiguous(),
+        kv_indices.contiguous(),
+        1 if causal else 0,
+        scores,
+    )
+    return scores
+if _HAS_NATIVE_OPS:
+    @torch.library.register_fake(
+        add_op_namespace_prefix("msa_nvfp4_dequant_swizzled_to_bf16")
+    )
+    def _msa_nvfp4_dequant_swizzled_to_bf16_fake(
+        packed: torch.Tensor,
+        scale_128x4: torch.Tensor,
+        global_scale: float,
+        out: torch.Tensor,
+    ) -> None:
+        return None
+def native_nvfp4_dequant_swizzled_to_bf16(
+    packed: torch.Tensor,
+    scale_128x4: torch.Tensor,
+    global_scale: torch.Tensor | float,
+) -> torch.Tensor:
+    """Expand swizzled NVFP4 weights/KV to dense BF16 with the native CUDA path."""
+    if not _HAS_NATIVE_OPS:
+        raise RuntimeError("native MiniMax MSA ops are not available in source-tree mode")
+    if packed.dtype is not torch.uint8:
+        raise TypeError("packed must be torch.uint8")
+    if scale_128x4.dtype is not torch.uint8:
+        raise TypeError("scale_128x4 must be torch.uint8")
+    original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
+    out = torch.empty(original_shape, device=packed.device, dtype=torch.bfloat16)
+    if isinstance(global_scale, torch.Tensor):
+        global_scale_value = float(global_scale.reshape(-1)[0].item())
+    else:
+        global_scale_value = float(global_scale)
+    ops.msa_nvfp4_dequant_swizzled_to_bf16(
+        packed.contiguous(),
+        scale_128x4.contiguous(),
+        global_scale_value,
+        out,
+    )
+    return out
+__all__ = [
+    "has_native_ops",
+    "native_topk_from_scores",
+    "native_decode_mma_supported",
+    "native_decode_sparse_attn_mma_paged",
+    "native_indexer_block_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
+]

build/torch212-cxx11-cu132-x86_64-linux/_ops.py CHANGED Viewed

@@ -1,9 +1,9 @@
 import torch
-from . import _minimaxai_msa_blackwell_cuda_58018d1
-ops = torch.ops._minimaxai_msa_blackwell_cuda_58018d1
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
-    return f"_minimaxai_msa_blackwell_cuda_58018d1::{op_name}"

 import torch
+from . import _minimaxai_msa_blackwell_cuda_25b0067
+ops = torch.ops._minimaxai_msa_blackwell_cuda_25b0067
 def add_op_namespace_prefix(op_name: str):
     """
     Prefix op by namespace.
     """
+    return f"_minimaxai_msa_blackwell_cuda_25b0067::{op_name}"

build/torch212-cxx11-cu132-x86_64-linux/api_status.py CHANGED Viewed

@@ -22,6 +22,7 @@ V1_AVAILABLE_FUNCTIONS = (
     "nvfp4_global_scale_from_amax",
     "has_native_ops",
     "native_topk_from_scores",
     "flash_decode_with_topk_idx",
     "flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
@@ -55,7 +56,7 @@ _OFFICIAL_API_STATUS = {
     "sparse_atten_nvfp4_kv_func": {
         "status": "available",
         "target": "blackwell-prefill-nvfp4",
-        "reason": "Official NVFP4 KV prefill API is available via NVFP4 dequantization followed by the Blackwell prefill wrapper.",
     },
     "sparse_decode_atten_func": {
         "status": "available",
@@ -70,12 +71,12 @@ _OFFICIAL_API_STATUS = {
     "fp4_indexer_block_scores": {
         "status": "available",
         "target": "blackwell-fp4-indexer",
-        "reason": "Official FP4 block-score API is available as a correctness-first PyTorch fallback.",
     },
     "build_k2q_csr": {
         "status": "available",
         "target": "v1",
-        "reason": "Torch CSR construction fallback is available for compatibility.",
     },
     "SparseK2qCsrBuilderSm100": {
         "status": "available",

     "nvfp4_global_scale_from_amax",
     "has_native_ops",
     "native_topk_from_scores",
+    "native_nvfp4_dequant_swizzled_to_bf16",
     "flash_decode_with_topk_idx",
     "flash_decode_with_gqa_share_sparse",
     "naive_flash_decode_with_topk_idx",
     "sparse_atten_nvfp4_kv_func": {
         "status": "available",
         "target": "blackwell-prefill-nvfp4",
+        "reason": "Official NVFP4 KV prefill API is available via native swizzled NVFP4 dequantization followed by the Blackwell prefill wrapper.",
     },
     "sparse_decode_atten_func": {
         "status": "available",
     "fp4_indexer_block_scores": {
         "status": "available",
         "target": "blackwell-fp4-indexer",
+        "reason": "Official FP4 block-score API uses the native Blackwell block-score kernel in built artifacts, with a reference fallback in source-tree mode.",
     },
     "build_k2q_csr": {
         "status": "available",
         "target": "v1",
+        "reason": "CSR construction helper is available for the official prefill API.",
     },
     "SparseK2qCsrBuilderSm100": {
         "status": "available",

build/torch212-cxx11-cu132-x86_64-linux/decode/topk_sparse.py CHANGED Viewed

@@ -8,6 +8,10 @@ import triton.language as tl
 from .._compat import is_hip  # vendored: was `from sglang.srt.utils import is_hip`
 from ..common.utils import robust_allocator
 _is_hip = is_hip()
@@ -341,6 +345,22 @@ def flash_decode_with_gqa_share_sparse(
     # sm scale
     if sm_scale is None:
         sm_scale = head_dim**-0.5
     # Pick NUM_TOPK_CHUNKS so total grid ≈ TARGET_GRID. Same constraints as
     # flash_decode_with_topk_idx: must be power of 2 (Triton arange) and must
     # only depend on shape constants (so grid is fixed within a cuda graph).

 from .._compat import is_hip  # vendored: was `from sglang.srt.utils import is_hip`
+from .._native import (
+    native_decode_mma_supported,
+    native_decode_sparse_attn_mma_paged,
+)
 from ..common.utils import robust_allocator
 _is_hip = is_hip()
     # sm scale
     if sm_scale is None:
         sm_scale = head_dim**-0.5
+    # Fast path: native tensor-core decode for the M3 MSA shape. Falls through
+    # to the Triton kernel for other shapes, sink, or fp8 caches.
+    if sink is None and not is_fp8 and native_decode_mma_supported(
+        q, k_cache, block_size=block_size
+    ):
+        return native_decode_sparse_attn_mma_paged(
+            q,
+            k_cache,
+            v_cache,
+            req_to_token,
+            seq_lens,
+            slot_ids,
+            block_size,
+            topk_idx,
+            sm_scale,
+        )
     # Pick NUM_TOPK_CHUNKS so total grid ≈ TARGET_GRID. Same constraints as
     # flash_decode_with_topk_idx: must be power of 2 (Triton arange) and must
     # only depend on shape constants (so grid is fixed within a cuda graph).

build/torch212-cxx11-cu132-x86_64-linux/interface.py CHANGED Viewed

@@ -7,6 +7,11 @@ from typing import Optional
 import torch
 from .decode.topk_sparse import flash_decode_with_gqa_share_sparse
 from .prefill.topk_sparse import flash_prefill_with_gqa_share_sparse
 from .quantize import Nvfp4QuantizedTensor, dequantize_nvfp4_128x4_to_bf16
@@ -488,6 +493,12 @@ def _nvfp4_dequant_kv(
 ) -> torch.Tensor:
     if packed.dtype is not torch.uint8:
         return packed
     original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
     rows = 1
     for dim in original_shape[:-1]:
@@ -572,11 +583,12 @@ def fp4_indexer_block_scores(
     qo_offset: Optional[torch.Tensor] = None,
     scale_layout: str = "preordered_mma",
 ) -> torch.Tensor:
-    """Correctness-first FP4 block-score API for Blackwell.
     Returns the official shape `[Hq, ceil(max_seqlen_k / 128), total_q]`.
-    This fallback is intentionally simple and is used to keep the public API
-    callable on Blackwell while optimized CUTE block scoring remains SM100-only.
     """
     del qo_offset
@@ -599,6 +611,29 @@ def fp4_indexer_block_scores(
     q = _apply_public_fp4_scale(_unpack_e2m1_fp4(q_fp4), q_scale)
     k = _apply_public_fp4_scale(_unpack_e2m1_fp4(k_fp4), k_scale)
     gqa = hq // hkv
     batch = int(cu_seqlens_q.numel()) - 1
     cu_q = cu_seqlens_q.to("cpu", non_blocking=False).tolist()

 import torch
+from ._native import (
+    has_native_ops,
+    native_indexer_block_scores,
+    native_nvfp4_dequant_swizzled_to_bf16,
+)
 from .decode.topk_sparse import flash_decode_with_gqa_share_sparse
 from .prefill.topk_sparse import flash_prefill_with_gqa_share_sparse
 from .quantize import Nvfp4QuantizedTensor, dequantize_nvfp4_128x4_to_bf16
 ) -> torch.Tensor:
     if packed.dtype is not torch.uint8:
         return packed
+    if has_native_ops() and packed.is_cuda:
+        return native_nvfp4_dequant_swizzled_to_bf16(
+            packed,
+            scale_128x4,
+            global_scale,
+        )
     original_shape = tuple(int(v) for v in packed.shape[:-1]) + (int(packed.shape[-1]) * 2,)
     rows = 1
     for dim in original_shape[:-1]:
     qo_offset: Optional[torch.Tensor] = None,
     scale_layout: str = "preordered_mma",
 ) -> torch.Tensor:
+    """FP4 block-score API for Blackwell.
     Returns the official shape `[Hq, ceil(max_seqlen_k / 128), total_q]`.
+    Built artifacts route CUDA tensors through the native Blackwell block-score
+    kernel. Source-tree mode uses the reference loop below so the API remains
+    testable before the extension is built.
     """
     del qo_offset
     q = _apply_public_fp4_scale(_unpack_e2m1_fp4(q_fp4), q_scale)
     k = _apply_public_fp4_scale(_unpack_e2m1_fp4(k_fp4), k_scale)
+    # Native block-max QK kernel for the cubic scoring loop when available.
+    # Same dequant numerics as the reference (bf16 compute); falls back to the
+    # Python reference below in source-tree mode.
+    if has_native_ops() and q_fp4.is_cuda:
+        batch = int(cu_seqlens_q.numel()) - 1
+        q_lens = (cu_seqlens_q[1:] - cu_seqlens_q[:-1]).to(torch.int64)
+        batch_of_q = torch.repeat_interleave(
+            torch.arange(batch, device=q_fp4.device, dtype=torch.int32),
+            q_lens,
+        )
+        return native_indexer_block_scores(
+            q.to(torch.bfloat16),
+            k.to(torch.bfloat16),
+            batch_of_q,
+            cu_seqlens_q.to(torch.int32),
+            cu_seqlens_k.to(torch.int32),
+            cu_page_offsets.to(torch.int32),
+            kv_indices.to(torch.int32),
+            max_blocks=max_blocks,
+            causal=causal,
+        )
     gqa = hq // hkv
     batch = int(cu_seqlens_q.numel()) - 1
     cu_q = cu_seqlens_q.to("cpu", non_blocking=False).tolist()

build/torch212-cxx11-cu132-x86_64-linux/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "minimaxai-msa-blackwell",
-  "id": "_minimaxai_msa_blackwell_cuda_58018d1",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -14,18 +14,18 @@
   "digest": {
     "algorithm": "sha256",
     "files": {
-      "__init__.py": "bH4lC+HdMx2DgfiEWfvxhd7GvQBOt02ftSPAzKo1rE0=",
       "_compat.py": "5Wy4a96HVMD8omRNsGg8rQJ+0qCCeXiDANk0CuyyP0I=",
-      "_minimaxai_msa_blackwell_cuda_58018d1.abi3.so": "PC8GpBOVMeHf/O3hlzsMBgWvo7YvmBiodQTteUrch00=",
-      "_native.py": "mIDvfkcQVd3H9plDsodJelrIFD2XHiEQgk4uqiiRLLE=",
-      "_ops.py": "kuQPE8cc2FNEl3lIMfjaj074xh5YntsWFa8cYC5CLIQ=",
-      "api_status.py": "VhS7aJhSznFwcgIQAnzXQTTuLJ2WgtysWrhX7RTSXaM=",
       "common/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "common/utils.py": "qtwh+wZoySlsWWALvWGVsmcIHKkdZUwePxUZ42tfWxc=",
       "decode/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "decode/flash_with_topk_idx.py": "Pk70+HJ60IKuy4KJC6hPdMtD2BjwPfSMabFiIBODWiQ=",
-      "decode/topk_sparse.py": "PLwQnOyiZT44pCXcZXzXdlZSixfBeanTZsfbpMOjlpQ=",
-      "interface.py": "TKlKdIqTtzfmieaY5ci3GwhD1QZXqUAZcqfcfPmyeuM=",
       "minimaxai_msa_blackwell/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
       "naive/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "naive/flash_with_topk_idx.py": "K7WQsRdO9+H2sz7pZ6axnJwE1OmnF0egdKtftHuChig=",

 {
   "name": "minimaxai-msa-blackwell",
+  "id": "_minimaxai_msa_blackwell_cuda_25b0067",
   "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "digest": {
     "algorithm": "sha256",
     "files": {
+      "__init__.py": "oPB7xQhjHd9+aLVJkEmJkZOXOwMEarmCdj3SBbhH2bM=",
       "_compat.py": "5Wy4a96HVMD8omRNsGg8rQJ+0qCCeXiDANk0CuyyP0I=",
+      "_minimaxai_msa_blackwell_cuda_25b0067.abi3.so": "TrAcGpbD5kBHSzpS8tONpL4q8CkMS5RcUPkOhn2ekH8=",
+      "_native.py": "wTh66DiIKxsJ5m8fkAXKSCBYD3CVFVx4PFxceQiwvPY=",
+      "_ops.py": "ikU3Lg+kgswKktYjNSq5f4DmgyTR8dnxV8lH3rmOWvk=",
+      "api_status.py": "7oZzIkb4fuwkAz3r4Q1slHbOfaP3jwVB647Ok223ATk=",
       "common/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "common/utils.py": "qtwh+wZoySlsWWALvWGVsmcIHKkdZUwePxUZ42tfWxc=",
       "decode/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "decode/flash_with_topk_idx.py": "Pk70+HJ60IKuy4KJC6hPdMtD2BjwPfSMabFiIBODWiQ=",
+      "decode/topk_sparse.py": "uINcV9SP5sAf9jYPsPvyQ9lon8sajZ8Qn/yJ+2LA8NU=",
+      "interface.py": "jfhBHGyunFF+OdFycJYyptNNIAngKrXY67buTXAN66M=",
       "minimaxai_msa_blackwell/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
       "naive/__init__.py": "bG7Ffen4LvQtgnYPFEpFccsWs81t4zqqeqn9ZeirH6E=",
       "naive/flash_with_topk_idx.py": "K7WQsRdO9+H2sz7pZ6axnJwE1OmnF0egdKtftHuChig=",