Hugging Face's logo

Kernels:
flashrt
/
MiniMaxAI-msa-blackwell

Kernel card Files
xet
 Community
liangsu9988 commited on
Commit
f3eeba1
·
1 Parent(s): 77a04b3

Update Blackwell MSA API card

Browse files
Files changed (2) hide show
  1. CARD.md +56 -9
  2. README.md +56 -9
CARD.md CHANGED
@@ -47,9 +47,9 @@ msa = get_kernel(
47
  | `dequantize_nvfp4_128x4_to_bf16` | Available reference dequantizer. |
48
  | `swizzle_nvfp4_scale_to_128x4` | Available scale-layout helper. |
49
  | `nvfp4_global_scale_from_amax` | Available scale helper. |
50
- | `sparse_atten_func` | Exported, but raises `NotImplementedError`: SM100 CSR prefill kernel is not ported here yet. |
51
- | `sparse_atten_nvfp4_kv_func` | Exported, but raises `NotImplementedError`: SM100 NVFP4 CSR prefill kernel is not ported here yet. |
52
- | `fp4_indexer_block_scores` | Exported, but raises `NotImplementedError`: SM100 FP4 CUTE indexer is not ported here yet. |
53
 
54
  ### FlashRT Blackwell helper names
55
 
@@ -121,6 +121,51 @@ wrapper.plan(
121
  out = wrapper.run(q, k, v)
122
  ```
123
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
124
  ## Direct FlashRT Decode Path
125
 
126
  Use this lower-level path if you already have `topk_idx` in FlashRT's paged KV
@@ -161,12 +206,14 @@ out = msa.flash_decode_with_gqa_share_sparse(
161
  This package contains:
162
 
163
  - native CUDA score-to-top-k helper;
164
- - Blackwell-validated Triton CUDA sparse decode attention;
165
- - MiniMaxAI/msa-compatible Python API layer for decode, CSR, and NVFP4 helpers.
166
-
167
- The upstream SM100 CSR prefill attention and FP4 CUTE indexer are not claimed as
168
- ported in this package yet. Their root names are exported so callers see a clear
169
- `NotImplementedError` instead of silently getting wrong behavior.
 
 
170
 
171
  Source provenance and validation details are documented in `SYNC.md` and
172
  `VALIDATION.md`.
 
47
  | `dequantize_nvfp4_128x4_to_bf16` | Available reference dequantizer. |
48
  | `swizzle_nvfp4_scale_to_128x4` | Available scale-layout helper. |
49
  | `nvfp4_global_scale_from_amax` | Available scale helper. |
50
+ | `sparse_atten_func` | Available. Official CSR sparse prefill API backed by the Blackwell Triton BF16/FP16 prefill kernel. |
51
+ | `sparse_atten_nvfp4_kv_func` | Available. NVFP4 KV compatibility path: dequantizes KV with 128x4 metadata, then calls Blackwell sparse prefill. |
52
+ | `fp4_indexer_block_scores` | Available. Correctness-first FP4 block-score fallback returning the official `[Hq, ceil(max_seqlen_k/128), total_q]` score layout. |
53
 
54
  ### FlashRT Blackwell helper names
55
 
 
121
  out = wrapper.run(q, k, v)
122
  ```
123
 
124
+ ## Prefill Example
125
+
126
+ This example uses the official MiniMax CSR prefill-facing name
127
+ `sparse_atten_func`.
128
+
129
+ ```python
130
+ import torch
131
+ from kernels import get_kernel
132
+
133
+ msa = get_kernel("flashrt/MiniMaxAI-msa-blackwell", version=1, trust_remote_code=True)
134
+
135
+ T, Hq, Hkv, D = 512, 64, 4, 128
136
+ page_size = 128
137
+ topk = 16
138
+ device = "cuda"
139
+ dtype = torch.bfloat16
140
+
141
+ q = torch.randn(T, Hq, D, device=device, dtype=dtype)
142
+ k = torch.randn(T, Hkv, D, device=device, dtype=dtype)
143
+ v = torch.randn_like(k)
144
+ cu = torch.tensor([0, T], device=device, dtype=torch.int32)
145
+
146
+ q2k = torch.full((Hkv, T, topk), -1, device=device, dtype=torch.int32)
147
+ for qi in range(T):
148
+ blocks = torch.arange(qi // page_size + 1, device=device, dtype=torch.int32)
149
+ q2k[:, qi, : min(topk, blocks.numel())] = blocks[:topk]
150
+
151
+ k2q_row_ptr, k2q_q_indices = msa.build_k2q_csr(
152
+ q2k, cu, cu, page_size, total_k=T
153
+ )
154
+ out = msa.sparse_atten_func(
155
+ q,
156
+ k,
157
+ v,
158
+ k2q_row_ptr,
159
+ k2q_q_indices,
160
+ topk,
161
+ cu_seqlens_q=cu,
162
+ cu_seqlens_k=cu,
163
+ max_seqlen_q=T,
164
+ max_seqlen_k=T,
165
+ blk_kv=page_size,
166
+ )
167
+ ```
168
+
169
  ## Direct FlashRT Decode Path
170
 
171
  Use this lower-level path if you already have `topk_idx` in FlashRT's paged KV
 
206
  This package contains:
207
 
208
  - native CUDA score-to-top-k helper;
209
+ - Blackwell-validated Triton CUDA sparse decode and prefill attention;
210
+ - MiniMaxAI/msa-compatible Python API layer for decode, prefill, CSR, NVFP4,
211
+ and FP4 block-score helpers.
212
+
213
+ The optimized SM100 CUTE prefill/indexer bodies are not claimed as ported here.
214
+ For Blackwell, this package provides a validated Triton sparse prefill path and
215
+ correctness-first compatibility fallbacks where the original API requires SM100
216
+ FP4/NVFP4-specific machinery.
217
 
218
  Source provenance and validation details are documented in `SYNC.md` and
219
  `VALIDATION.md`.
README.md CHANGED
@@ -47,9 +47,9 @@ msa = get_kernel(
47
  | `dequantize_nvfp4_128x4_to_bf16` | Available reference dequantizer. |
48
  | `swizzle_nvfp4_scale_to_128x4` | Available scale-layout helper. |
49
  | `nvfp4_global_scale_from_amax` | Available scale helper. |
50
- | `sparse_atten_func` | Exported, but raises `NotImplementedError`: SM100 CSR prefill kernel is not ported here yet. |
51
- | `sparse_atten_nvfp4_kv_func` | Exported, but raises `NotImplementedError`: SM100 NVFP4 CSR prefill kernel is not ported here yet. |
52
- | `fp4_indexer_block_scores` | Exported, but raises `NotImplementedError`: SM100 FP4 CUTE indexer is not ported here yet. |
53
 
54
  ### FlashRT Blackwell helper names
55
 
@@ -121,6 +121,51 @@ wrapper.plan(
121
  out = wrapper.run(q, k, v)
122
  ```
123
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
124
  ## Direct FlashRT Decode Path
125
 
126
  Use this lower-level path if you already have `topk_idx` in FlashRT's paged KV
@@ -161,12 +206,14 @@ out = msa.flash_decode_with_gqa_share_sparse(
161
  This package contains:
162
 
163
  - native CUDA score-to-top-k helper;
164
- - Blackwell-validated Triton CUDA sparse decode attention;
165
- - MiniMaxAI/msa-compatible Python API layer for decode, CSR, and NVFP4 helpers.
166
-
167
- The upstream SM100 CSR prefill attention and FP4 CUTE indexer are not claimed as
168
- ported in this package yet. Their root names are exported so callers see a clear
169
- `NotImplementedError` instead of silently getting wrong behavior.
 
 
170
 
171
  Source provenance and validation details are documented in `SYNC.md` and
172
  `VALIDATION.md`.
 
47
  | `dequantize_nvfp4_128x4_to_bf16` | Available reference dequantizer. |
48
  | `swizzle_nvfp4_scale_to_128x4` | Available scale-layout helper. |
49
  | `nvfp4_global_scale_from_amax` | Available scale helper. |
50
+ | `sparse_atten_func` | Available. Official CSR sparse prefill API backed by the Blackwell Triton BF16/FP16 prefill kernel. |
51
+ | `sparse_atten_nvfp4_kv_func` | Available. NVFP4 KV compatibility path: dequantizes KV with 128x4 metadata, then calls Blackwell sparse prefill. |
52
+ | `fp4_indexer_block_scores` | Available. Correctness-first FP4 block-score fallback returning the official `[Hq, ceil(max_seqlen_k/128), total_q]` score layout. |
53
 
54
  ### FlashRT Blackwell helper names
55
 
 
121
  out = wrapper.run(q, k, v)
122
  ```
123
 
124
+ ## Prefill Example
125
+
126
+ This example uses the official MiniMax CSR prefill-facing name
127
+ `sparse_atten_func`.
128
+
129
+ ```python
130
+ import torch
131
+ from kernels import get_kernel
132
+
133
+ msa = get_kernel("flashrt/MiniMaxAI-msa-blackwell", version=1, trust_remote_code=True)
134
+
135
+ T, Hq, Hkv, D = 512, 64, 4, 128
136
+ page_size = 128
137
+ topk = 16
138
+ device = "cuda"
139
+ dtype = torch.bfloat16
140
+
141
+ q = torch.randn(T, Hq, D, device=device, dtype=dtype)
142
+ k = torch.randn(T, Hkv, D, device=device, dtype=dtype)
143
+ v = torch.randn_like(k)
144
+ cu = torch.tensor([0, T], device=device, dtype=torch.int32)
145
+
146
+ q2k = torch.full((Hkv, T, topk), -1, device=device, dtype=torch.int32)
147
+ for qi in range(T):
148
+ blocks = torch.arange(qi // page_size + 1, device=device, dtype=torch.int32)
149
+ q2k[:, qi, : min(topk, blocks.numel())] = blocks[:topk]
150
+
151
+ k2q_row_ptr, k2q_q_indices = msa.build_k2q_csr(
152
+ q2k, cu, cu, page_size, total_k=T
153
+ )
154
+ out = msa.sparse_atten_func(
155
+ q,
156
+ k,
157
+ v,
158
+ k2q_row_ptr,
159
+ k2q_q_indices,
160
+ topk,
161
+ cu_seqlens_q=cu,
162
+ cu_seqlens_k=cu,
163
+ max_seqlen_q=T,
164
+ max_seqlen_k=T,
165
+ blk_kv=page_size,
166
+ )
167
+ ```
168
+
169
  ## Direct FlashRT Decode Path
170
 
171
  Use this lower-level path if you already have `topk_idx` in FlashRT's paged KV
 
206
  This package contains:
207
 
208
  - native CUDA score-to-top-k helper;
209
+ - Blackwell-validated Triton CUDA sparse decode and prefill attention;
210
+ - MiniMaxAI/msa-compatible Python API layer for decode, prefill, CSR, NVFP4,
211
+ and FP4 block-score helpers.
212
+
213
+ The optimized SM100 CUTE prefill/indexer bodies are not claimed as ported here.
214
+ For Blackwell, this package provides a validated Triton sparse prefill path and
215
+ correctness-first compatibility fallbacks where the original API requires SM100
216
+ FP4/NVFP4-specific machinery.
217
 
218
  Source provenance and validation details are documented in `SYNC.md` and
219
  `VALIDATION.md`.