Hugging Face's logo

Kernels:
flashrt
/
fp4-gemm

Kernel card Files
xet
 Community
liangsu9988 commited on
Commit
20b01a8
·
verified ·
1 Parent(s): 3a6e5ec

Uploaded using `kernel-builder`.

Browse files
Files changed (16) hide show
  1. build/torch212-cxx11-cu130-x86_64-linux/__init__.py +77 -0
  2. build/torch212-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_230e467.abi3.so → _fp4_gemm_cuda_b6588aa.abi3.so} +2 -2
  3. build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
  4. build/torch212-cxx11-cu130-x86_64-linux/metadata.json +5 -5
  5. build/torch212-cxx11-cu132-x86_64-linux/__init__.py +77 -0
  6. build/torch212-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_230e467.abi3.so → _fp4_gemm_cuda_b6588aa.abi3.so} +2 -2
  7. build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
  8. build/torch212-cxx11-cu132-x86_64-linux/metadata.json +5 -5
  9. build/torch213-cxx11-cu130-x86_64-linux/__init__.py +77 -0
  10. build/torch213-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_230e467.abi3.so → _fp4_gemm_cuda_b6588aa.abi3.so} +2 -2
  11. build/torch213-cxx11-cu130-x86_64-linux/_ops.py +3 -3
  12. build/torch213-cxx11-cu130-x86_64-linux/metadata.json +5 -5
  13. build/torch213-cxx11-cu132-x86_64-linux/__init__.py +77 -0
  14. build/torch213-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_230e467.abi3.so → _fp4_gemm_cuda_b6588aa.abi3.so} +2 -2
  15. build/torch213-cxx11-cu132-x86_64-linux/_ops.py +3 -3
  16. build/torch213-cxx11-cu132-x86_64-linux/metadata.json +5 -5
build/torch212-cxx11-cu130-x86_64-linux/__init__.py CHANGED
@@ -42,6 +42,16 @@ def _linear_fp16_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, var
42
  return None
43
 
44
 
 
 
 
 
 
 
 
 
 
 
45
  @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
46
  def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
47
  return None
@@ -94,6 +104,11 @@ def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb
94
  return None
95
 
96
 
 
 
 
 
 
97
  @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
98
  def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
99
  return None
@@ -158,6 +173,19 @@ def quantize_fp4_sfa_fp16(
158
  return packed, sfa
159
 
160
 
 
 
 
 
 
 
 
 
 
 
 
 
 
161
  def quantize_e0m3_sfa_fp16(
162
  x: torch.Tensor,
163
  packed: torch.Tensor | None = None,
@@ -276,6 +304,52 @@ def nvfp4_gemm_fp16(
276
  return out
277
 
278
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
279
  def nvfp4_gemm_geglu_nvfp4_fp16(
280
  a_packed: torch.Tensor,
281
  b_interleaved_packed: torch.Tensor,
@@ -533,6 +607,8 @@ __all__ = [
533
  "fp4_w4a4_gemv_warpsplit_bf16",
534
  "nvfp4_gemm_bf16",
535
  "nvfp4_gemm_fp16",
 
 
536
  "nvfp4_gemm_geglu_nvfp4_fp16",
537
  "nvfp4_gemm_bias_gelu_nvfp4_fp16",
538
  "nvfp4_gemm_bias_residual_fp16",
@@ -545,6 +621,7 @@ __all__ = [
545
  "nvfp4_gemm_streamk_bf16",
546
  "nvfp4_gemm_streamk_bias_bf16",
547
  "quantize_fp4_sfa_fp16",
 
548
  "quantize_e0m3_sfa_fp16",
549
  "quantize_fp4_sfa_bf16",
550
  "sfa_size_bytes",
 
42
  return None
43
 
44
 
45
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_variant_bf16"))
46
+ def _linear_variant_bf16_fake(a, b, sfa, sfb, out, alpha: float = 1.0, variant: int = 7) -> None:
47
+ return None
48
+
49
+
50
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_nvfp4"))
51
+ def _linear_fp4out_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
52
+ return None
53
+
54
+
55
  @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
56
  def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
57
  return None
 
104
  return None
105
 
106
 
107
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_fp16"))
108
+ def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
109
+ return None
110
+
111
+
112
  @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
113
  def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
114
  return None
 
173
  return packed, sfa
174
 
175
 
176
+ def quantize_fp4_sfa_mse_fp16(
177
+ x: torch.Tensor,
178
+ packed: torch.Tensor | None = None,
179
+ sfa: torch.Tensor | None = None,
180
+ is_sfb: bool = False,
181
+ ):
182
+ """Pack FP16 weights using per-block reconstruction-MSE scale search."""
183
+ if packed is None or sfa is None:
184
+ packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
185
+ ops.quantize_fp4_sfa_mse_fp16(x, packed, sfa, bool(is_sfb))
186
+ return packed, sfa
187
+
188
+
189
  def quantize_e0m3_sfa_fp16(
190
  x: torch.Tensor,
191
  packed: torch.Tensor | None = None,
 
304
  return out
305
 
306
 
307
+ def nvfp4_gemm_variant_bf16(
308
+ a_packed: torch.Tensor,
309
+ b_packed: torch.Tensor,
310
+ sfa: torch.Tensor,
311
+ sfb: torch.Tensor,
312
+ alpha: float = 1.0,
313
+ out: torch.Tensor | None = None,
314
+ variant: int = 7,
315
+ ) -> torch.Tensor:
316
+ """SM110 native NVFP4 GEMM with BF16 output.
317
+
318
+ ``variant=7`` serves large-M encoder/SigLIP shapes; ``variant=10`` is
319
+ the narrow-N PI0.5 decoder tile. The indices match FlashRT native.
320
+ """
321
+ if out is None:
322
+ out = torch.empty(
323
+ (a_packed.shape[0], b_packed.shape[0]),
324
+ device=a_packed.device,
325
+ dtype=torch.bfloat16,
326
+ )
327
+ ops.nvfp4_gemm_variant_bf16(
328
+ a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant)
329
+ )
330
+ return out
331
+
332
+
333
+ def nvfp4_gemm_nvfp4(
334
+ a_packed: torch.Tensor,
335
+ b_packed: torch.Tensor,
336
+ sfa: torch.Tensor,
337
+ sfb: torch.Tensor,
338
+ *,
339
+ out_packed: torch.Tensor | None = None,
340
+ out_sfa: torch.Tensor | None = None,
341
+ ) -> tuple[torch.Tensor, torch.Tensor]:
342
+ """SM110 NVFP4 GEMM with an NVFP4 output and CUTLASS SFA layout."""
343
+ if out_packed is None or out_sfa is None:
344
+ out_packed, out_sfa = _alloc_fp4(
345
+ a_packed.shape[0], b_packed.shape[0], a_packed.device
346
+ )
347
+ ops.nvfp4_gemm_nvfp4(
348
+ a_packed, b_packed, sfa, sfb, out_packed, out_sfa
349
+ )
350
+ return out_packed, out_sfa
351
+
352
+
353
  def nvfp4_gemm_geglu_nvfp4_fp16(
354
  a_packed: torch.Tensor,
355
  b_interleaved_packed: torch.Tensor,
 
607
  "fp4_w4a4_gemv_warpsplit_bf16",
608
  "nvfp4_gemm_bf16",
609
  "nvfp4_gemm_fp16",
610
+ "nvfp4_gemm_variant_bf16",
611
+ "nvfp4_gemm_nvfp4",
612
  "nvfp4_gemm_geglu_nvfp4_fp16",
613
  "nvfp4_gemm_bias_gelu_nvfp4_fp16",
614
  "nvfp4_gemm_bias_residual_fp16",
 
621
  "nvfp4_gemm_streamk_bf16",
622
  "nvfp4_gemm_streamk_bias_bf16",
623
  "quantize_fp4_sfa_fp16",
624
+ "quantize_fp4_sfa_mse_fp16",
625
  "quantize_e0m3_sfa_fp16",
626
  "quantize_fp4_sfa_bf16",
627
  "sfa_size_bytes",
build/torch212-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_230e467.abi3.so → _fp4_gemm_cuda_b6588aa.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:5e319eb3cf472a0ead98df64005aca885558f238a7b682e0767a27189d170e46
3
- size 7181592
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2d9be54b8589e5131efa5a4ff9ddd2699519e41aa5603d54b5b0a8410c7c2a07
3
+ size 7718880
build/torch212-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_230e467
3
- ops = torch.ops._fp4_gemm_cuda_230e467
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_230e467::{op_name}"
 
1
  import torch
2
+ from . import _fp4_gemm_cuda_b6588aa
3
+ ops = torch.ops._fp4_gemm_cuda_b6588aa
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp4_gemm_cuda_b6588aa::{op_name}"
build/torch212-cxx11-cu130-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp4-gemm",
3
- "id": "_fp4_gemm_cuda_230e467",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
@@ -14,9 +14,9 @@
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
- "__init__.py": "Qfdvg/9EedWLGh34xAalDfC0IsAHCOL44L4ShaP/w5Y=",
18
- "_fp4_gemm_cuda_230e467.abi3.so": "XjGes89HKg6tmN9kAFrKiFVY8jintoLgdnonGJ0XDkY=",
19
- "_ops.py": "CxLLKrT37bp1ie5kAVTryKkFcmkRJqJ2fYXrSYuMHcA="
20
  }
21
  },
22
  "provenance": {
@@ -26,7 +26,7 @@
26
  "dirty": false
27
  },
28
  "kernel": {
29
- "sha": "230e467129aa8e5fc687611fbb2aa40a12e09ad1",
30
  "dirty": false
31
  }
32
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_b6588aa",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
+ "__init__.py": "Fb/FDYOpDw+YGlevblUsv4AkILDQCJdB807qdDUDI1c=",
18
+ "_fp4_gemm_cuda_b6588aa.abi3.so": "LZvlS4WJ5RMe+lpP+d3SaZUZ5BqlYD1UtbCoQQx8Kgc=",
19
+ "_ops.py": "KKClchvuzBz2g1A148Gi2R6QIx0Ipj0CKlTT511GrUY="
20
  }
21
  },
22
  "provenance": {
 
26
  "dirty": false
27
  },
28
  "kernel": {
29
+ "sha": "b6588aa59dd4f8c646f90b58312a0a66886e061f",
30
  "dirty": false
31
  }
32
  }
build/torch212-cxx11-cu132-x86_64-linux/__init__.py CHANGED
@@ -42,6 +42,16 @@ def _linear_fp16_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, var
42
  return None
43
 
44
 
 
 
 
 
 
 
 
 
 
 
45
  @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
46
  def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
47
  return None
@@ -94,6 +104,11 @@ def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb
94
  return None
95
 
96
 
 
 
 
 
 
97
  @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
98
  def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
99
  return None
@@ -158,6 +173,19 @@ def quantize_fp4_sfa_fp16(
158
  return packed, sfa
159
 
160
 
 
 
 
 
 
 
 
 
 
 
 
 
 
161
  def quantize_e0m3_sfa_fp16(
162
  x: torch.Tensor,
163
  packed: torch.Tensor | None = None,
@@ -276,6 +304,52 @@ def nvfp4_gemm_fp16(
276
  return out
277
 
278
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
279
  def nvfp4_gemm_geglu_nvfp4_fp16(
280
  a_packed: torch.Tensor,
281
  b_interleaved_packed: torch.Tensor,
@@ -533,6 +607,8 @@ __all__ = [
533
  "fp4_w4a4_gemv_warpsplit_bf16",
534
  "nvfp4_gemm_bf16",
535
  "nvfp4_gemm_fp16",
 
 
536
  "nvfp4_gemm_geglu_nvfp4_fp16",
537
  "nvfp4_gemm_bias_gelu_nvfp4_fp16",
538
  "nvfp4_gemm_bias_residual_fp16",
@@ -545,6 +621,7 @@ __all__ = [
545
  "nvfp4_gemm_streamk_bf16",
546
  "nvfp4_gemm_streamk_bias_bf16",
547
  "quantize_fp4_sfa_fp16",
 
548
  "quantize_e0m3_sfa_fp16",
549
  "quantize_fp4_sfa_bf16",
550
  "sfa_size_bytes",
 
42
  return None
43
 
44
 
45
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_variant_bf16"))
46
+ def _linear_variant_bf16_fake(a, b, sfa, sfb, out, alpha: float = 1.0, variant: int = 7) -> None:
47
+ return None
48
+
49
+
50
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_nvfp4"))
51
+ def _linear_fp4out_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
52
+ return None
53
+
54
+
55
  @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
56
  def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
57
  return None
 
104
  return None
105
 
106
 
107
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_fp16"))
108
+ def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
109
+ return None
110
+
111
+
112
  @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
113
  def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
114
  return None
 
173
  return packed, sfa
174
 
175
 
176
+ def quantize_fp4_sfa_mse_fp16(
177
+ x: torch.Tensor,
178
+ packed: torch.Tensor | None = None,
179
+ sfa: torch.Tensor | None = None,
180
+ is_sfb: bool = False,
181
+ ):
182
+ """Pack FP16 weights using per-block reconstruction-MSE scale search."""
183
+ if packed is None or sfa is None:
184
+ packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
185
+ ops.quantize_fp4_sfa_mse_fp16(x, packed, sfa, bool(is_sfb))
186
+ return packed, sfa
187
+
188
+
189
  def quantize_e0m3_sfa_fp16(
190
  x: torch.Tensor,
191
  packed: torch.Tensor | None = None,
 
304
  return out
305
 
306
 
307
+ def nvfp4_gemm_variant_bf16(
308
+ a_packed: torch.Tensor,
309
+ b_packed: torch.Tensor,
310
+ sfa: torch.Tensor,
311
+ sfb: torch.Tensor,
312
+ alpha: float = 1.0,
313
+ out: torch.Tensor | None = None,
314
+ variant: int = 7,
315
+ ) -> torch.Tensor:
316
+ """SM110 native NVFP4 GEMM with BF16 output.
317
+
318
+ ``variant=7`` serves large-M encoder/SigLIP shapes; ``variant=10`` is
319
+ the narrow-N PI0.5 decoder tile. The indices match FlashRT native.
320
+ """
321
+ if out is None:
322
+ out = torch.empty(
323
+ (a_packed.shape[0], b_packed.shape[0]),
324
+ device=a_packed.device,
325
+ dtype=torch.bfloat16,
326
+ )
327
+ ops.nvfp4_gemm_variant_bf16(
328
+ a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant)
329
+ )
330
+ return out
331
+
332
+
333
+ def nvfp4_gemm_nvfp4(
334
+ a_packed: torch.Tensor,
335
+ b_packed: torch.Tensor,
336
+ sfa: torch.Tensor,
337
+ sfb: torch.Tensor,
338
+ *,
339
+ out_packed: torch.Tensor | None = None,
340
+ out_sfa: torch.Tensor | None = None,
341
+ ) -> tuple[torch.Tensor, torch.Tensor]:
342
+ """SM110 NVFP4 GEMM with an NVFP4 output and CUTLASS SFA layout."""
343
+ if out_packed is None or out_sfa is None:
344
+ out_packed, out_sfa = _alloc_fp4(
345
+ a_packed.shape[0], b_packed.shape[0], a_packed.device
346
+ )
347
+ ops.nvfp4_gemm_nvfp4(
348
+ a_packed, b_packed, sfa, sfb, out_packed, out_sfa
349
+ )
350
+ return out_packed, out_sfa
351
+
352
+
353
  def nvfp4_gemm_geglu_nvfp4_fp16(
354
  a_packed: torch.Tensor,
355
  b_interleaved_packed: torch.Tensor,
 
607
  "fp4_w4a4_gemv_warpsplit_bf16",
608
  "nvfp4_gemm_bf16",
609
  "nvfp4_gemm_fp16",
610
+ "nvfp4_gemm_variant_bf16",
611
+ "nvfp4_gemm_nvfp4",
612
  "nvfp4_gemm_geglu_nvfp4_fp16",
613
  "nvfp4_gemm_bias_gelu_nvfp4_fp16",
614
  "nvfp4_gemm_bias_residual_fp16",
 
621
  "nvfp4_gemm_streamk_bf16",
622
  "nvfp4_gemm_streamk_bias_bf16",
623
  "quantize_fp4_sfa_fp16",
624
+ "quantize_fp4_sfa_mse_fp16",
625
  "quantize_e0m3_sfa_fp16",
626
  "quantize_fp4_sfa_bf16",
627
  "sfa_size_bytes",
build/torch212-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_230e467.abi3.so → _fp4_gemm_cuda_b6588aa.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:42c8b744abb8d1df65520ba3135b88c58a002b31bb4fcc1b66b823c961ee815d
3
- size 7169248
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:308bf04520e36c5f5cbdcb45644cc86b01449729e246a10f87cc835b2935be25
3
+ size 7706544
build/torch212-cxx11-cu132-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_230e467
3
- ops = torch.ops._fp4_gemm_cuda_230e467
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_230e467::{op_name}"
 
1
  import torch
2
+ from . import _fp4_gemm_cuda_b6588aa
3
+ ops = torch.ops._fp4_gemm_cuda_b6588aa
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp4_gemm_cuda_b6588aa::{op_name}"
build/torch212-cxx11-cu132-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp4-gemm",
3
- "id": "_fp4_gemm_cuda_230e467",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
@@ -14,9 +14,9 @@
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
- "__init__.py": "Qfdvg/9EedWLGh34xAalDfC0IsAHCOL44L4ShaP/w5Y=",
18
- "_fp4_gemm_cuda_230e467.abi3.so": "Qsi3RKu40d9lUgujE1uIxYoAKzG7T8wbZrgjyWHugV0=",
19
- "_ops.py": "CxLLKrT37bp1ie5kAVTryKkFcmkRJqJ2fYXrSYuMHcA="
20
  }
21
  },
22
  "provenance": {
@@ -26,7 +26,7 @@
26
  "dirty": false
27
  },
28
  "kernel": {
29
- "sha": "230e467129aa8e5fc687611fbb2aa40a12e09ad1",
30
  "dirty": false
31
  }
32
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_b6588aa",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
+ "__init__.py": "Fb/FDYOpDw+YGlevblUsv4AkILDQCJdB807qdDUDI1c=",
18
+ "_fp4_gemm_cuda_b6588aa.abi3.so": "MIvwRSDjbF9cvctFZEzIawFElyniRqEPh8yDWyk1viU=",
19
+ "_ops.py": "KKClchvuzBz2g1A148Gi2R6QIx0Ipj0CKlTT511GrUY="
20
  }
21
  },
22
  "provenance": {
 
26
  "dirty": false
27
  },
28
  "kernel": {
29
+ "sha": "b6588aa59dd4f8c646f90b58312a0a66886e061f",
30
  "dirty": false
31
  }
32
  }
build/torch213-cxx11-cu130-x86_64-linux/__init__.py CHANGED
@@ -42,6 +42,16 @@ def _linear_fp16_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, var
42
  return None
43
 
44
 
 
 
 
 
 
 
 
 
 
 
45
  @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
46
  def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
47
  return None
@@ -94,6 +104,11 @@ def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb
94
  return None
95
 
96
 
 
 
 
 
 
97
  @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
98
  def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
99
  return None
@@ -158,6 +173,19 @@ def quantize_fp4_sfa_fp16(
158
  return packed, sfa
159
 
160
 
 
 
 
 
 
 
 
 
 
 
 
 
 
161
  def quantize_e0m3_sfa_fp16(
162
  x: torch.Tensor,
163
  packed: torch.Tensor | None = None,
@@ -276,6 +304,52 @@ def nvfp4_gemm_fp16(
276
  return out
277
 
278
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
279
  def nvfp4_gemm_geglu_nvfp4_fp16(
280
  a_packed: torch.Tensor,
281
  b_interleaved_packed: torch.Tensor,
@@ -533,6 +607,8 @@ __all__ = [
533
  "fp4_w4a4_gemv_warpsplit_bf16",
534
  "nvfp4_gemm_bf16",
535
  "nvfp4_gemm_fp16",
 
 
536
  "nvfp4_gemm_geglu_nvfp4_fp16",
537
  "nvfp4_gemm_bias_gelu_nvfp4_fp16",
538
  "nvfp4_gemm_bias_residual_fp16",
@@ -545,6 +621,7 @@ __all__ = [
545
  "nvfp4_gemm_streamk_bf16",
546
  "nvfp4_gemm_streamk_bias_bf16",
547
  "quantize_fp4_sfa_fp16",
 
548
  "quantize_e0m3_sfa_fp16",
549
  "quantize_fp4_sfa_bf16",
550
  "sfa_size_bytes",
 
42
  return None
43
 
44
 
45
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_variant_bf16"))
46
+ def _linear_variant_bf16_fake(a, b, sfa, sfb, out, alpha: float = 1.0, variant: int = 7) -> None:
47
+ return None
48
+
49
+
50
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_nvfp4"))
51
+ def _linear_fp4out_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
52
+ return None
53
+
54
+
55
  @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
56
  def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
57
  return None
 
104
  return None
105
 
106
 
107
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_fp16"))
108
+ def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
109
+ return None
110
+
111
+
112
  @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
113
  def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
114
  return None
 
173
  return packed, sfa
174
 
175
 
176
+ def quantize_fp4_sfa_mse_fp16(
177
+ x: torch.Tensor,
178
+ packed: torch.Tensor | None = None,
179
+ sfa: torch.Tensor | None = None,
180
+ is_sfb: bool = False,
181
+ ):
182
+ """Pack FP16 weights using per-block reconstruction-MSE scale search."""
183
+ if packed is None or sfa is None:
184
+ packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
185
+ ops.quantize_fp4_sfa_mse_fp16(x, packed, sfa, bool(is_sfb))
186
+ return packed, sfa
187
+
188
+
189
  def quantize_e0m3_sfa_fp16(
190
  x: torch.Tensor,
191
  packed: torch.Tensor | None = None,
 
304
  return out
305
 
306
 
307
+ def nvfp4_gemm_variant_bf16(
308
+ a_packed: torch.Tensor,
309
+ b_packed: torch.Tensor,
310
+ sfa: torch.Tensor,
311
+ sfb: torch.Tensor,
312
+ alpha: float = 1.0,
313
+ out: torch.Tensor | None = None,
314
+ variant: int = 7,
315
+ ) -> torch.Tensor:
316
+ """SM110 native NVFP4 GEMM with BF16 output.
317
+
318
+ ``variant=7`` serves large-M encoder/SigLIP shapes; ``variant=10`` is
319
+ the narrow-N PI0.5 decoder tile. The indices match FlashRT native.
320
+ """
321
+ if out is None:
322
+ out = torch.empty(
323
+ (a_packed.shape[0], b_packed.shape[0]),
324
+ device=a_packed.device,
325
+ dtype=torch.bfloat16,
326
+ )
327
+ ops.nvfp4_gemm_variant_bf16(
328
+ a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant)
329
+ )
330
+ return out
331
+
332
+
333
+ def nvfp4_gemm_nvfp4(
334
+ a_packed: torch.Tensor,
335
+ b_packed: torch.Tensor,
336
+ sfa: torch.Tensor,
337
+ sfb: torch.Tensor,
338
+ *,
339
+ out_packed: torch.Tensor | None = None,
340
+ out_sfa: torch.Tensor | None = None,
341
+ ) -> tuple[torch.Tensor, torch.Tensor]:
342
+ """SM110 NVFP4 GEMM with an NVFP4 output and CUTLASS SFA layout."""
343
+ if out_packed is None or out_sfa is None:
344
+ out_packed, out_sfa = _alloc_fp4(
345
+ a_packed.shape[0], b_packed.shape[0], a_packed.device
346
+ )
347
+ ops.nvfp4_gemm_nvfp4(
348
+ a_packed, b_packed, sfa, sfb, out_packed, out_sfa
349
+ )
350
+ return out_packed, out_sfa
351
+
352
+
353
  def nvfp4_gemm_geglu_nvfp4_fp16(
354
  a_packed: torch.Tensor,
355
  b_interleaved_packed: torch.Tensor,
 
607
  "fp4_w4a4_gemv_warpsplit_bf16",
608
  "nvfp4_gemm_bf16",
609
  "nvfp4_gemm_fp16",
610
+ "nvfp4_gemm_variant_bf16",
611
+ "nvfp4_gemm_nvfp4",
612
  "nvfp4_gemm_geglu_nvfp4_fp16",
613
  "nvfp4_gemm_bias_gelu_nvfp4_fp16",
614
  "nvfp4_gemm_bias_residual_fp16",
 
621
  "nvfp4_gemm_streamk_bf16",
622
  "nvfp4_gemm_streamk_bias_bf16",
623
  "quantize_fp4_sfa_fp16",
624
+ "quantize_fp4_sfa_mse_fp16",
625
  "quantize_e0m3_sfa_fp16",
626
  "quantize_fp4_sfa_bf16",
627
  "sfa_size_bytes",
build/torch213-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_230e467.abi3.so → _fp4_gemm_cuda_b6588aa.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:c5c6596baac50bb9b9ebb3ac5acc544c62cb33959c4ea64bdd8e8fd039bdb294
3
- size 7181432
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:67ad5a71cafc0766b46303485eb638fcd03d048a6b48e97f3a0a90e17585ee2a
3
+ size 7714632
build/torch213-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_230e467
3
- ops = torch.ops._fp4_gemm_cuda_230e467
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_230e467::{op_name}"
 
1
  import torch
2
+ from . import _fp4_gemm_cuda_b6588aa
3
+ ops = torch.ops._fp4_gemm_cuda_b6588aa
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp4_gemm_cuda_b6588aa::{op_name}"
build/torch213-cxx11-cu130-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp4-gemm",
3
- "id": "_fp4_gemm_cuda_230e467",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
@@ -14,9 +14,9 @@
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
- "__init__.py": "Qfdvg/9EedWLGh34xAalDfC0IsAHCOL44L4ShaP/w5Y=",
18
- "_fp4_gemm_cuda_230e467.abi3.so": "xcZZa6rFC7m567OsWsxUTGLLM5WcTqZL3Y6P0Dm9spQ=",
19
- "_ops.py": "CxLLKrT37bp1ie5kAVTryKkFcmkRJqJ2fYXrSYuMHcA="
20
  }
21
  },
22
  "provenance": {
@@ -26,7 +26,7 @@
26
  "dirty": false
27
  },
28
  "kernel": {
29
- "sha": "230e467129aa8e5fc687611fbb2aa40a12e09ad1",
30
  "dirty": false
31
  }
32
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_b6588aa",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
+ "__init__.py": "Fb/FDYOpDw+YGlevblUsv4AkILDQCJdB807qdDUDI1c=",
18
+ "_fp4_gemm_cuda_b6588aa.abi3.so": "Z61accr8B2a0YwNIXrY4/NA9BIprSOl/OgqQ4XWF7io=",
19
+ "_ops.py": "KKClchvuzBz2g1A148Gi2R6QIx0Ipj0CKlTT511GrUY="
20
  }
21
  },
22
  "provenance": {
 
26
  "dirty": false
27
  },
28
  "kernel": {
29
+ "sha": "b6588aa59dd4f8c646f90b58312a0a66886e061f",
30
  "dirty": false
31
  }
32
  }
build/torch213-cxx11-cu132-x86_64-linux/__init__.py CHANGED
@@ -42,6 +42,16 @@ def _linear_fp16_fake(a_packed, b_packed, sfa, sfb, out, alpha: float = 1.0, var
42
  return None
43
 
44
 
 
 
 
 
 
 
 
 
 
 
45
  @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
46
  def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
47
  return None
@@ -94,6 +104,11 @@ def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb
94
  return None
95
 
96
 
 
 
 
 
 
97
  @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
98
  def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
99
  return None
@@ -158,6 +173,19 @@ def quantize_fp4_sfa_fp16(
158
  return packed, sfa
159
 
160
 
 
 
 
 
 
 
 
 
 
 
 
 
 
161
  def quantize_e0m3_sfa_fp16(
162
  x: torch.Tensor,
163
  packed: torch.Tensor | None = None,
@@ -276,6 +304,52 @@ def nvfp4_gemm_fp16(
276
  return out
277
 
278
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
279
  def nvfp4_gemm_geglu_nvfp4_fp16(
280
  a_packed: torch.Tensor,
281
  b_interleaved_packed: torch.Tensor,
@@ -533,6 +607,8 @@ __all__ = [
533
  "fp4_w4a4_gemv_warpsplit_bf16",
534
  "nvfp4_gemm_bf16",
535
  "nvfp4_gemm_fp16",
 
 
536
  "nvfp4_gemm_geglu_nvfp4_fp16",
537
  "nvfp4_gemm_bias_gelu_nvfp4_fp16",
538
  "nvfp4_gemm_bias_residual_fp16",
@@ -545,6 +621,7 @@ __all__ = [
545
  "nvfp4_gemm_streamk_bf16",
546
  "nvfp4_gemm_streamk_bias_bf16",
547
  "quantize_fp4_sfa_fp16",
 
548
  "quantize_e0m3_sfa_fp16",
549
  "quantize_fp4_sfa_bf16",
550
  "sfa_size_bytes",
 
42
  return None
43
 
44
 
45
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_variant_bf16"))
46
+ def _linear_variant_bf16_fake(a, b, sfa, sfb, out, alpha: float = 1.0, variant: int = 7) -> None:
47
+ return None
48
+
49
+
50
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_nvfp4"))
51
+ def _linear_fp4out_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
52
+ return None
53
+
54
+
55
  @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_geglu_nvfp4_fp16"))
56
  def _geglu_fp4_fake(a, b, sfa, sfb, scratch, out_packed, out_sfa, skinny: bool = False) -> None:
57
  return None
 
104
  return None
105
 
106
 
107
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_mse_fp16"))
108
+ def _quant_mse_fake(x, packed, sfa, is_sfb: bool = False) -> None:
109
+ return None
110
+
111
+
112
  @torch.library.register_fake(add_op_namespace_prefix("quantize_e0m3_sfa_fp16"))
113
  def _quant_e0m3_fake(x, packed, sfa, is_sfb: bool = False) -> None:
114
  return None
 
173
  return packed, sfa
174
 
175
 
176
+ def quantize_fp4_sfa_mse_fp16(
177
+ x: torch.Tensor,
178
+ packed: torch.Tensor | None = None,
179
+ sfa: torch.Tensor | None = None,
180
+ is_sfb: bool = False,
181
+ ):
182
+ """Pack FP16 weights using per-block reconstruction-MSE scale search."""
183
+ if packed is None or sfa is None:
184
+ packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
185
+ ops.quantize_fp4_sfa_mse_fp16(x, packed, sfa, bool(is_sfb))
186
+ return packed, sfa
187
+
188
+
189
  def quantize_e0m3_sfa_fp16(
190
  x: torch.Tensor,
191
  packed: torch.Tensor | None = None,
 
304
  return out
305
 
306
 
307
+ def nvfp4_gemm_variant_bf16(
308
+ a_packed: torch.Tensor,
309
+ b_packed: torch.Tensor,
310
+ sfa: torch.Tensor,
311
+ sfb: torch.Tensor,
312
+ alpha: float = 1.0,
313
+ out: torch.Tensor | None = None,
314
+ variant: int = 7,
315
+ ) -> torch.Tensor:
316
+ """SM110 native NVFP4 GEMM with BF16 output.
317
+
318
+ ``variant=7`` serves large-M encoder/SigLIP shapes; ``variant=10`` is
319
+ the narrow-N PI0.5 decoder tile. The indices match FlashRT native.
320
+ """
321
+ if out is None:
322
+ out = torch.empty(
323
+ (a_packed.shape[0], b_packed.shape[0]),
324
+ device=a_packed.device,
325
+ dtype=torch.bfloat16,
326
+ )
327
+ ops.nvfp4_gemm_variant_bf16(
328
+ a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant)
329
+ )
330
+ return out
331
+
332
+
333
+ def nvfp4_gemm_nvfp4(
334
+ a_packed: torch.Tensor,
335
+ b_packed: torch.Tensor,
336
+ sfa: torch.Tensor,
337
+ sfb: torch.Tensor,
338
+ *,
339
+ out_packed: torch.Tensor | None = None,
340
+ out_sfa: torch.Tensor | None = None,
341
+ ) -> tuple[torch.Tensor, torch.Tensor]:
342
+ """SM110 NVFP4 GEMM with an NVFP4 output and CUTLASS SFA layout."""
343
+ if out_packed is None or out_sfa is None:
344
+ out_packed, out_sfa = _alloc_fp4(
345
+ a_packed.shape[0], b_packed.shape[0], a_packed.device
346
+ )
347
+ ops.nvfp4_gemm_nvfp4(
348
+ a_packed, b_packed, sfa, sfb, out_packed, out_sfa
349
+ )
350
+ return out_packed, out_sfa
351
+
352
+
353
  def nvfp4_gemm_geglu_nvfp4_fp16(
354
  a_packed: torch.Tensor,
355
  b_interleaved_packed: torch.Tensor,
 
607
  "fp4_w4a4_gemv_warpsplit_bf16",
608
  "nvfp4_gemm_bf16",
609
  "nvfp4_gemm_fp16",
610
+ "nvfp4_gemm_variant_bf16",
611
+ "nvfp4_gemm_nvfp4",
612
  "nvfp4_gemm_geglu_nvfp4_fp16",
613
  "nvfp4_gemm_bias_gelu_nvfp4_fp16",
614
  "nvfp4_gemm_bias_residual_fp16",
 
621
  "nvfp4_gemm_streamk_bf16",
622
  "nvfp4_gemm_streamk_bias_bf16",
623
  "quantize_fp4_sfa_fp16",
624
+ "quantize_fp4_sfa_mse_fp16",
625
  "quantize_e0m3_sfa_fp16",
626
  "quantize_fp4_sfa_bf16",
627
  "sfa_size_bytes",
build/torch213-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_230e467.abi3.so → _fp4_gemm_cuda_b6588aa.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:e941f703da2103bf557f43020d96a34f7f3f2d18d5bfdf5e37731945e6d6e792
3
- size 7169088
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4b27d4b88246c79fd017c8845f6f99cd1c9e7946039fda8f6abec2635a4edd65
3
+ size 7706384
build/torch213-cxx11-cu132-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_230e467
3
- ops = torch.ops._fp4_gemm_cuda_230e467
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_230e467::{op_name}"
 
1
  import torch
2
+ from . import _fp4_gemm_cuda_b6588aa
3
+ ops = torch.ops._fp4_gemm_cuda_b6588aa
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp4_gemm_cuda_b6588aa::{op_name}"
build/torch213-cxx11-cu132-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp4-gemm",
3
- "id": "_fp4_gemm_cuda_230e467",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
@@ -14,9 +14,9 @@
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
- "__init__.py": "Qfdvg/9EedWLGh34xAalDfC0IsAHCOL44L4ShaP/w5Y=",
18
- "_fp4_gemm_cuda_230e467.abi3.so": "6UH3A9ohA79Vf0MCDZajT38/LRjVv99eN3MZRebW55I=",
19
- "_ops.py": "CxLLKrT37bp1ie5kAVTryKkFcmkRJqJ2fYXrSYuMHcA="
20
  }
21
  },
22
  "provenance": {
@@ -26,7 +26,7 @@
26
  "dirty": false
27
  },
28
  "kernel": {
29
- "sha": "230e467129aa8e5fc687611fbb2aa40a12e09ad1",
30
  "dirty": false
31
  }
32
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_b6588aa",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
+ "__init__.py": "Fb/FDYOpDw+YGlevblUsv4AkILDQCJdB807qdDUDI1c=",
18
+ "_fp4_gemm_cuda_b6588aa.abi3.so": "SyfUuIJGx5/QF8iEX2+ZzRyeeUYDn9qPar7CY1pO3WU=",
19
+ "_ops.py": "KKClchvuzBz2g1A148Gi2R6QIx0Ipj0CKlTT511GrUY="
20
  }
21
  },
22
  "provenance": {
 
26
  "dirty": false
27
  },
28
  "kernel": {
29
+ "sha": "b6588aa59dd4f8c646f90b58312a0a66886e061f",
30
  "dirty": false
31
  }
32
  }