Hugging Face's logo

Kernels:
flashrt
/
fp4-gemm

Kernel card Files
xet
 Community
liangsu9988 commited on
Commit
8bc04d8
·
verified ·
1 Parent(s): ee38e70

Uploaded using `kernel-builder`.

Browse files
Files changed (16) hide show
  1. build/torch212-cxx11-cu130-x86_64-linux/__init__.py +74 -0
  2. build/torch212-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_55c4885.abi3.so → _fp4_gemm_cuda_230e467.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 +74 -0
  6. build/torch212-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_55c4885.abi3.so → _fp4_gemm_cuda_230e467.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 +74 -0
  10. build/torch213-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_55c4885.abi3.so → _fp4_gemm_cuda_230e467.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 +74 -0
  14. build/torch213-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_55c4885.abi3.so → _fp4_gemm_cuda_230e467.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
@@ -94,6 +94,21 @@ 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_fp4_sfa_bf16"))
98
  def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
99
  return None
@@ -143,6 +158,62 @@ def quantize_fp4_sfa_fp16(
143
  return packed, sfa
144
 
145
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
146
  def quantize_fp4_sfa_bf16(
147
  x: torch.Tensor,
148
  packed: torch.Tensor | None = None,
@@ -457,6 +528,7 @@ def nvfp4_gemm_streamk_bias_bf16(
457
 
458
  __all__ = [
459
  "dequantize_fp4_sfa_fp16",
 
460
  "fp4_w4a16_linear_bf16",
461
  "fp4_w4a4_gemv_warpsplit_bf16",
462
  "nvfp4_gemm_bf16",
@@ -469,9 +541,11 @@ __all__ = [
469
  "nvfp4_gemm_bias_gelu_nvfp4",
470
  "nvfp4_gemm_bias_residual_bf16",
471
  "nvfp4_gemm_residual_bf16",
 
472
  "nvfp4_gemm_streamk_bf16",
473
  "nvfp4_gemm_streamk_bias_bf16",
474
  "quantize_fp4_sfa_fp16",
 
475
  "quantize_fp4_sfa_bf16",
476
  "sfa_size_bytes",
477
  ]
 
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
100
+
101
+
102
+ @torch.library.register_fake(add_op_namespace_prefix("e0m3_weight_gemm_fp16"))
103
+ def _e0m3_gemm_fake(a, b, sfa, sfb, out, alpha: float = 1.0, a_format: int = 1) -> None:
104
+ return None
105
+
106
+
107
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_relu2_nvfp4"))
108
+ def _relu2_gemm_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
109
+ return None
110
+
111
+
112
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16"))
113
  def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
114
  return None
 
158
  return packed, sfa
159
 
160
 
161
+ def quantize_e0m3_sfa_fp16(
162
+ x: torch.Tensor,
163
+ packed: torch.Tensor | None = None,
164
+ sfa: torch.Tensor | None = None,
165
+ is_sfb: bool = False,
166
+ ):
167
+ """Quantize FP16 to packed E0M3 with CUTLASS SFA/SFB layout."""
168
+ if packed is None or sfa is None:
169
+ packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
170
+ ops.quantize_e0m3_sfa_fp16(x, packed, sfa, bool(is_sfb))
171
+ return packed, sfa
172
+
173
+
174
+ def e0m3_weight_gemm_fp16(
175
+ a_packed: torch.Tensor,
176
+ b_packed: torch.Tensor,
177
+ sfa: torch.Tensor,
178
+ sfb: torch.Tensor,
179
+ *,
180
+ alpha: float = 1.0,
181
+ a_format: int = 1,
182
+ out: torch.Tensor | None = None,
183
+ ):
184
+ """SM110 GEMM with E0M3 weights and E2M1 or E0M3 activations."""
185
+ if out is None:
186
+ out = torch.empty(
187
+ (a_packed.shape[0], b_packed.shape[0]),
188
+ device=a_packed.device,
189
+ dtype=torch.float16,
190
+ )
191
+ ops.e0m3_weight_gemm_fp16(
192
+ a_packed, b_packed, sfa, sfb, out, float(alpha), int(a_format)
193
+ )
194
+ return out
195
+
196
+
197
+ def nvfp4_gemm_relu2_nvfp4(
198
+ a_packed: torch.Tensor,
199
+ b_packed: torch.Tensor,
200
+ sfa: torch.Tensor,
201
+ sfb: torch.Tensor,
202
+ *,
203
+ out_packed: torch.Tensor | None = None,
204
+ out_sfa: torch.Tensor | None = None,
205
+ ):
206
+ """SM110 NVFP4 GEMM with fused ReLU-squared and NVFP4 output."""
207
+ if out_packed is None or out_sfa is None:
208
+ out_packed, out_sfa = _alloc_fp4(
209
+ a_packed.shape[0], b_packed.shape[0], a_packed.device
210
+ )
211
+ ops.nvfp4_gemm_relu2_nvfp4(
212
+ a_packed, b_packed, sfa, sfb, out_packed, out_sfa
213
+ )
214
+ return out_packed, out_sfa
215
+
216
+
217
  def quantize_fp4_sfa_bf16(
218
  x: torch.Tensor,
219
  packed: torch.Tensor | None = None,
 
528
 
529
  __all__ = [
530
  "dequantize_fp4_sfa_fp16",
531
+ "e0m3_weight_gemm_fp16",
532
  "fp4_w4a16_linear_bf16",
533
  "fp4_w4a4_gemv_warpsplit_bf16",
534
  "nvfp4_gemm_bf16",
 
541
  "nvfp4_gemm_bias_gelu_nvfp4",
542
  "nvfp4_gemm_bias_residual_bf16",
543
  "nvfp4_gemm_residual_bf16",
544
+ "nvfp4_gemm_relu2_nvfp4",
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",
551
  ]
build/torch212-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_55c4885.abi3.so → _fp4_gemm_cuda_230e467.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:c55a4f6233f699e6d2e69205e0f4c5e7753d0dc216129bf75e4477d675daee8e
3
- size 6350504
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5e319eb3cf472a0ead98df64005aca885558f238a7b682e0767a27189d170e46
3
+ size 7181592
build/torch212-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_55c4885
3
- ops = torch.ops._fp4_gemm_cuda_55c4885
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_55c4885::{op_name}"
 
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}"
build/torch212-cxx11-cu130-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp4-gemm",
3
- "id": "_fp4_gemm_cuda_55c4885",
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": "4o968SKYWj8r1Ly7FqSPFK/AYdu4G4bXuVv1UcvYjk8=",
18
- "_fp4_gemm_cuda_55c4885.abi3.so": "xVpPYjP2mebS5pIF4PTF53U9DcIWEpv3XkR31nXa7o4=",
19
- "_ops.py": "IVtq+3pHXwuoCbgYGw3JP7/nKr3GEZqn8D97JJjf7ZE="
20
  }
21
  },
22
  "provenance": {
@@ -26,7 +26,7 @@
26
  "dirty": false
27
  },
28
  "kernel": {
29
- "sha": "55c4885251068f195418bcf9ae541f4b757a6ea0",
30
  "dirty": false
31
  }
32
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_230e467",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
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
  "dirty": false
27
  },
28
  "kernel": {
29
+ "sha": "230e467129aa8e5fc687611fbb2aa40a12e09ad1",
30
  "dirty": false
31
  }
32
  }
build/torch212-cxx11-cu132-x86_64-linux/__init__.py CHANGED
@@ -94,6 +94,21 @@ 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_fp4_sfa_bf16"))
98
  def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
99
  return None
@@ -143,6 +158,62 @@ def quantize_fp4_sfa_fp16(
143
  return packed, sfa
144
 
145
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
146
  def quantize_fp4_sfa_bf16(
147
  x: torch.Tensor,
148
  packed: torch.Tensor | None = None,
@@ -457,6 +528,7 @@ def nvfp4_gemm_streamk_bias_bf16(
457
 
458
  __all__ = [
459
  "dequantize_fp4_sfa_fp16",
 
460
  "fp4_w4a16_linear_bf16",
461
  "fp4_w4a4_gemv_warpsplit_bf16",
462
  "nvfp4_gemm_bf16",
@@ -469,9 +541,11 @@ __all__ = [
469
  "nvfp4_gemm_bias_gelu_nvfp4",
470
  "nvfp4_gemm_bias_residual_bf16",
471
  "nvfp4_gemm_residual_bf16",
 
472
  "nvfp4_gemm_streamk_bf16",
473
  "nvfp4_gemm_streamk_bias_bf16",
474
  "quantize_fp4_sfa_fp16",
 
475
  "quantize_fp4_sfa_bf16",
476
  "sfa_size_bytes",
477
  ]
 
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
100
+
101
+
102
+ @torch.library.register_fake(add_op_namespace_prefix("e0m3_weight_gemm_fp16"))
103
+ def _e0m3_gemm_fake(a, b, sfa, sfb, out, alpha: float = 1.0, a_format: int = 1) -> None:
104
+ return None
105
+
106
+
107
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_relu2_nvfp4"))
108
+ def _relu2_gemm_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
109
+ return None
110
+
111
+
112
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16"))
113
  def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
114
  return None
 
158
  return packed, sfa
159
 
160
 
161
+ def quantize_e0m3_sfa_fp16(
162
+ x: torch.Tensor,
163
+ packed: torch.Tensor | None = None,
164
+ sfa: torch.Tensor | None = None,
165
+ is_sfb: bool = False,
166
+ ):
167
+ """Quantize FP16 to packed E0M3 with CUTLASS SFA/SFB layout."""
168
+ if packed is None or sfa is None:
169
+ packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
170
+ ops.quantize_e0m3_sfa_fp16(x, packed, sfa, bool(is_sfb))
171
+ return packed, sfa
172
+
173
+
174
+ def e0m3_weight_gemm_fp16(
175
+ a_packed: torch.Tensor,
176
+ b_packed: torch.Tensor,
177
+ sfa: torch.Tensor,
178
+ sfb: torch.Tensor,
179
+ *,
180
+ alpha: float = 1.0,
181
+ a_format: int = 1,
182
+ out: torch.Tensor | None = None,
183
+ ):
184
+ """SM110 GEMM with E0M3 weights and E2M1 or E0M3 activations."""
185
+ if out is None:
186
+ out = torch.empty(
187
+ (a_packed.shape[0], b_packed.shape[0]),
188
+ device=a_packed.device,
189
+ dtype=torch.float16,
190
+ )
191
+ ops.e0m3_weight_gemm_fp16(
192
+ a_packed, b_packed, sfa, sfb, out, float(alpha), int(a_format)
193
+ )
194
+ return out
195
+
196
+
197
+ def nvfp4_gemm_relu2_nvfp4(
198
+ a_packed: torch.Tensor,
199
+ b_packed: torch.Tensor,
200
+ sfa: torch.Tensor,
201
+ sfb: torch.Tensor,
202
+ *,
203
+ out_packed: torch.Tensor | None = None,
204
+ out_sfa: torch.Tensor | None = None,
205
+ ):
206
+ """SM110 NVFP4 GEMM with fused ReLU-squared and NVFP4 output."""
207
+ if out_packed is None or out_sfa is None:
208
+ out_packed, out_sfa = _alloc_fp4(
209
+ a_packed.shape[0], b_packed.shape[0], a_packed.device
210
+ )
211
+ ops.nvfp4_gemm_relu2_nvfp4(
212
+ a_packed, b_packed, sfa, sfb, out_packed, out_sfa
213
+ )
214
+ return out_packed, out_sfa
215
+
216
+
217
  def quantize_fp4_sfa_bf16(
218
  x: torch.Tensor,
219
  packed: torch.Tensor | None = None,
 
528
 
529
  __all__ = [
530
  "dequantize_fp4_sfa_fp16",
531
+ "e0m3_weight_gemm_fp16",
532
  "fp4_w4a16_linear_bf16",
533
  "fp4_w4a4_gemv_warpsplit_bf16",
534
  "nvfp4_gemm_bf16",
 
541
  "nvfp4_gemm_bias_gelu_nvfp4",
542
  "nvfp4_gemm_bias_residual_bf16",
543
  "nvfp4_gemm_residual_bf16",
544
+ "nvfp4_gemm_relu2_nvfp4",
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",
551
  ]
build/torch212-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_55c4885.abi3.so → _fp4_gemm_cuda_230e467.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:9745e68973460d1d2fbb943c7ba63ab9fb23d8bb8a59c707bcf386e3891c69cb
3
- size 6342256
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:42c8b744abb8d1df65520ba3135b88c58a002b31bb4fcc1b66b823c961ee815d
3
+ size 7169248
build/torch212-cxx11-cu132-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_55c4885
3
- ops = torch.ops._fp4_gemm_cuda_55c4885
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_55c4885::{op_name}"
 
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}"
build/torch212-cxx11-cu132-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp4-gemm",
3
- "id": "_fp4_gemm_cuda_55c4885",
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": "4o968SKYWj8r1Ly7FqSPFK/AYdu4G4bXuVv1UcvYjk8=",
18
- "_fp4_gemm_cuda_55c4885.abi3.so": "l0XmiXNGDR0vu5Q8e6Y6ufsj2LuKWccHvPOG44kcacs=",
19
- "_ops.py": "IVtq+3pHXwuoCbgYGw3JP7/nKr3GEZqn8D97JJjf7ZE="
20
  }
21
  },
22
  "provenance": {
@@ -26,7 +26,7 @@
26
  "dirty": false
27
  },
28
  "kernel": {
29
- "sha": "55c4885251068f195418bcf9ae541f4b757a6ea0",
30
  "dirty": false
31
  }
32
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_230e467",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
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
  "dirty": false
27
  },
28
  "kernel": {
29
+ "sha": "230e467129aa8e5fc687611fbb2aa40a12e09ad1",
30
  "dirty": false
31
  }
32
  }
build/torch213-cxx11-cu130-x86_64-linux/__init__.py CHANGED
@@ -94,6 +94,21 @@ 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_fp4_sfa_bf16"))
98
  def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
99
  return None
@@ -143,6 +158,62 @@ def quantize_fp4_sfa_fp16(
143
  return packed, sfa
144
 
145
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
146
  def quantize_fp4_sfa_bf16(
147
  x: torch.Tensor,
148
  packed: torch.Tensor | None = None,
@@ -457,6 +528,7 @@ def nvfp4_gemm_streamk_bias_bf16(
457
 
458
  __all__ = [
459
  "dequantize_fp4_sfa_fp16",
 
460
  "fp4_w4a16_linear_bf16",
461
  "fp4_w4a4_gemv_warpsplit_bf16",
462
  "nvfp4_gemm_bf16",
@@ -469,9 +541,11 @@ __all__ = [
469
  "nvfp4_gemm_bias_gelu_nvfp4",
470
  "nvfp4_gemm_bias_residual_bf16",
471
  "nvfp4_gemm_residual_bf16",
 
472
  "nvfp4_gemm_streamk_bf16",
473
  "nvfp4_gemm_streamk_bias_bf16",
474
  "quantize_fp4_sfa_fp16",
 
475
  "quantize_fp4_sfa_bf16",
476
  "sfa_size_bytes",
477
  ]
 
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
100
+
101
+
102
+ @torch.library.register_fake(add_op_namespace_prefix("e0m3_weight_gemm_fp16"))
103
+ def _e0m3_gemm_fake(a, b, sfa, sfb, out, alpha: float = 1.0, a_format: int = 1) -> None:
104
+ return None
105
+
106
+
107
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_relu2_nvfp4"))
108
+ def _relu2_gemm_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
109
+ return None
110
+
111
+
112
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16"))
113
  def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
114
  return None
 
158
  return packed, sfa
159
 
160
 
161
+ def quantize_e0m3_sfa_fp16(
162
+ x: torch.Tensor,
163
+ packed: torch.Tensor | None = None,
164
+ sfa: torch.Tensor | None = None,
165
+ is_sfb: bool = False,
166
+ ):
167
+ """Quantize FP16 to packed E0M3 with CUTLASS SFA/SFB layout."""
168
+ if packed is None or sfa is None:
169
+ packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
170
+ ops.quantize_e0m3_sfa_fp16(x, packed, sfa, bool(is_sfb))
171
+ return packed, sfa
172
+
173
+
174
+ def e0m3_weight_gemm_fp16(
175
+ a_packed: torch.Tensor,
176
+ b_packed: torch.Tensor,
177
+ sfa: torch.Tensor,
178
+ sfb: torch.Tensor,
179
+ *,
180
+ alpha: float = 1.0,
181
+ a_format: int = 1,
182
+ out: torch.Tensor | None = None,
183
+ ):
184
+ """SM110 GEMM with E0M3 weights and E2M1 or E0M3 activations."""
185
+ if out is None:
186
+ out = torch.empty(
187
+ (a_packed.shape[0], b_packed.shape[0]),
188
+ device=a_packed.device,
189
+ dtype=torch.float16,
190
+ )
191
+ ops.e0m3_weight_gemm_fp16(
192
+ a_packed, b_packed, sfa, sfb, out, float(alpha), int(a_format)
193
+ )
194
+ return out
195
+
196
+
197
+ def nvfp4_gemm_relu2_nvfp4(
198
+ a_packed: torch.Tensor,
199
+ b_packed: torch.Tensor,
200
+ sfa: torch.Tensor,
201
+ sfb: torch.Tensor,
202
+ *,
203
+ out_packed: torch.Tensor | None = None,
204
+ out_sfa: torch.Tensor | None = None,
205
+ ):
206
+ """SM110 NVFP4 GEMM with fused ReLU-squared and NVFP4 output."""
207
+ if out_packed is None or out_sfa is None:
208
+ out_packed, out_sfa = _alloc_fp4(
209
+ a_packed.shape[0], b_packed.shape[0], a_packed.device
210
+ )
211
+ ops.nvfp4_gemm_relu2_nvfp4(
212
+ a_packed, b_packed, sfa, sfb, out_packed, out_sfa
213
+ )
214
+ return out_packed, out_sfa
215
+
216
+
217
  def quantize_fp4_sfa_bf16(
218
  x: torch.Tensor,
219
  packed: torch.Tensor | None = None,
 
528
 
529
  __all__ = [
530
  "dequantize_fp4_sfa_fp16",
531
+ "e0m3_weight_gemm_fp16",
532
  "fp4_w4a16_linear_bf16",
533
  "fp4_w4a4_gemv_warpsplit_bf16",
534
  "nvfp4_gemm_bf16",
 
541
  "nvfp4_gemm_bias_gelu_nvfp4",
542
  "nvfp4_gemm_bias_residual_bf16",
543
  "nvfp4_gemm_residual_bf16",
544
+ "nvfp4_gemm_relu2_nvfp4",
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",
551
  ]
build/torch213-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_55c4885.abi3.so → _fp4_gemm_cuda_230e467.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:0a35111ce16d53ac8a6300acc0ea58b3c560abc04ca16e3be4297b4e3bb09335
3
- size 6350344
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c5c6596baac50bb9b9ebb3ac5acc544c62cb33959c4ea64bdd8e8fd039bdb294
3
+ size 7181432
build/torch213-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_55c4885
3
- ops = torch.ops._fp4_gemm_cuda_55c4885
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_55c4885::{op_name}"
 
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}"
build/torch213-cxx11-cu130-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp4-gemm",
3
- "id": "_fp4_gemm_cuda_55c4885",
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": "4o968SKYWj8r1Ly7FqSPFK/AYdu4G4bXuVv1UcvYjk8=",
18
- "_fp4_gemm_cuda_55c4885.abi3.so": "CjURHOFtU6yKYwCswOpYs8Vgq8BMoW475Cl7TjuwkzU=",
19
- "_ops.py": "IVtq+3pHXwuoCbgYGw3JP7/nKr3GEZqn8D97JJjf7ZE="
20
  }
21
  },
22
  "provenance": {
@@ -26,7 +26,7 @@
26
  "dirty": false
27
  },
28
  "kernel": {
29
- "sha": "55c4885251068f195418bcf9ae541f4b757a6ea0",
30
  "dirty": false
31
  }
32
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_230e467",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
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
  "dirty": false
27
  },
28
  "kernel": {
29
+ "sha": "230e467129aa8e5fc687611fbb2aa40a12e09ad1",
30
  "dirty": false
31
  }
32
  }
build/torch213-cxx11-cu132-x86_64-linux/__init__.py CHANGED
@@ -94,6 +94,21 @@ 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_fp4_sfa_bf16"))
98
  def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
99
  return None
@@ -143,6 +158,62 @@ def quantize_fp4_sfa_fp16(
143
  return packed, sfa
144
 
145
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
146
  def quantize_fp4_sfa_bf16(
147
  x: torch.Tensor,
148
  packed: torch.Tensor | None = None,
@@ -457,6 +528,7 @@ def nvfp4_gemm_streamk_bias_bf16(
457
 
458
  __all__ = [
459
  "dequantize_fp4_sfa_fp16",
 
460
  "fp4_w4a16_linear_bf16",
461
  "fp4_w4a4_gemv_warpsplit_bf16",
462
  "nvfp4_gemm_bf16",
@@ -469,9 +541,11 @@ __all__ = [
469
  "nvfp4_gemm_bias_gelu_nvfp4",
470
  "nvfp4_gemm_bias_residual_bf16",
471
  "nvfp4_gemm_residual_bf16",
 
472
  "nvfp4_gemm_streamk_bf16",
473
  "nvfp4_gemm_streamk_bias_bf16",
474
  "quantize_fp4_sfa_fp16",
 
475
  "quantize_fp4_sfa_bf16",
476
  "sfa_size_bytes",
477
  ]
 
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
100
+
101
+
102
+ @torch.library.register_fake(add_op_namespace_prefix("e0m3_weight_gemm_fp16"))
103
+ def _e0m3_gemm_fake(a, b, sfa, sfb, out, alpha: float = 1.0, a_format: int = 1) -> None:
104
+ return None
105
+
106
+
107
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_relu2_nvfp4"))
108
+ def _relu2_gemm_fake(a, b, sfa, sfb, out_packed, out_sfa) -> None:
109
+ return None
110
+
111
+
112
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_bf16"))
113
  def _quant_bf16_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
114
  return None
 
158
  return packed, sfa
159
 
160
 
161
+ def quantize_e0m3_sfa_fp16(
162
+ x: torch.Tensor,
163
+ packed: torch.Tensor | None = None,
164
+ sfa: torch.Tensor | None = None,
165
+ is_sfb: bool = False,
166
+ ):
167
+ """Quantize FP16 to packed E0M3 with CUTLASS SFA/SFB layout."""
168
+ if packed is None or sfa is None:
169
+ packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
170
+ ops.quantize_e0m3_sfa_fp16(x, packed, sfa, bool(is_sfb))
171
+ return packed, sfa
172
+
173
+
174
+ def e0m3_weight_gemm_fp16(
175
+ a_packed: torch.Tensor,
176
+ b_packed: torch.Tensor,
177
+ sfa: torch.Tensor,
178
+ sfb: torch.Tensor,
179
+ *,
180
+ alpha: float = 1.0,
181
+ a_format: int = 1,
182
+ out: torch.Tensor | None = None,
183
+ ):
184
+ """SM110 GEMM with E0M3 weights and E2M1 or E0M3 activations."""
185
+ if out is None:
186
+ out = torch.empty(
187
+ (a_packed.shape[0], b_packed.shape[0]),
188
+ device=a_packed.device,
189
+ dtype=torch.float16,
190
+ )
191
+ ops.e0m3_weight_gemm_fp16(
192
+ a_packed, b_packed, sfa, sfb, out, float(alpha), int(a_format)
193
+ )
194
+ return out
195
+
196
+
197
+ def nvfp4_gemm_relu2_nvfp4(
198
+ a_packed: torch.Tensor,
199
+ b_packed: torch.Tensor,
200
+ sfa: torch.Tensor,
201
+ sfb: torch.Tensor,
202
+ *,
203
+ out_packed: torch.Tensor | None = None,
204
+ out_sfa: torch.Tensor | None = None,
205
+ ):
206
+ """SM110 NVFP4 GEMM with fused ReLU-squared and NVFP4 output."""
207
+ if out_packed is None or out_sfa is None:
208
+ out_packed, out_sfa = _alloc_fp4(
209
+ a_packed.shape[0], b_packed.shape[0], a_packed.device
210
+ )
211
+ ops.nvfp4_gemm_relu2_nvfp4(
212
+ a_packed, b_packed, sfa, sfb, out_packed, out_sfa
213
+ )
214
+ return out_packed, out_sfa
215
+
216
+
217
  def quantize_fp4_sfa_bf16(
218
  x: torch.Tensor,
219
  packed: torch.Tensor | None = None,
 
528
 
529
  __all__ = [
530
  "dequantize_fp4_sfa_fp16",
531
+ "e0m3_weight_gemm_fp16",
532
  "fp4_w4a16_linear_bf16",
533
  "fp4_w4a4_gemv_warpsplit_bf16",
534
  "nvfp4_gemm_bf16",
 
541
  "nvfp4_gemm_bias_gelu_nvfp4",
542
  "nvfp4_gemm_bias_residual_bf16",
543
  "nvfp4_gemm_residual_bf16",
544
+ "nvfp4_gemm_relu2_nvfp4",
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",
551
  ]
build/torch213-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_55c4885.abi3.so → _fp4_gemm_cuda_230e467.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:ce8cdbdd5f77404bb14b639d16ca5cdada3ddd43a40d6edeb5a4cdee2d6a260e
3
- size 6342104
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e941f703da2103bf557f43020d96a34f7f3f2d18d5bfdf5e37731945e6d6e792
3
+ size 7169088
build/torch213-cxx11-cu132-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_55c4885
3
- ops = torch.ops._fp4_gemm_cuda_55c4885
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_55c4885::{op_name}"
 
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}"
build/torch213-cxx11-cu132-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp4-gemm",
3
- "id": "_fp4_gemm_cuda_55c4885",
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": "4o968SKYWj8r1Ly7FqSPFK/AYdu4G4bXuVv1UcvYjk8=",
18
- "_fp4_gemm_cuda_55c4885.abi3.so": "zozb3V93QEuxS2OdFspc2to93UOkDW7etaTN7i1qJg4=",
19
- "_ops.py": "IVtq+3pHXwuoCbgYGw3JP7/nKr3GEZqn8D97JJjf7ZE="
20
  }
21
  },
22
  "provenance": {
@@ -26,7 +26,7 @@
26
  "dirty": false
27
  },
28
  "kernel": {
29
- "sha": "55c4885251068f195418bcf9ae541f4b757a6ea0",
30
  "dirty": false
31
  }
32
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_230e467",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
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
  "dirty": false
27
  },
28
  "kernel": {
29
+ "sha": "230e467129aa8e5fc687611fbb2aa40a12e09ad1",
30
  "dirty": false
31
  }
32
  }