Hugging Face's logo

Kernels:
drbh
/
fp8linear

Kernel card Files
xet
 Community
drbh HF Staff commited on
Commit
e0b89a9
·
verified ·
1 Parent(s): f5125c7

Uploaded using `kernel-builder`.

Browse files
Files changed (32) hide show
  1. build/torch211-cxx11-cu126-x86_64-linux/__init__.py +9 -0
  2. build/torch211-cxx11-cu126-x86_64-linux/{_fp8linear_cuda_91cad61.abi3.so → _fp8linear_cuda_999b53f.abi3.so} +2 -2
  3. build/torch211-cxx11-cu126-x86_64-linux/_ops.py +3 -3
  4. build/torch211-cxx11-cu126-x86_64-linux/layers.py +36 -0
  5. build/torch211-cxx11-cu126-x86_64-linux/metadata.json +5 -5
  6. build/torch211-cxx11-cu128-x86_64-linux/__init__.py +9 -0
  7. build/torch211-cxx11-cu128-x86_64-linux/{_fp8linear_cuda_91cad61.abi3.so → _fp8linear_cuda_999b53f.abi3.so} +2 -2
  8. build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
  9. build/torch211-cxx11-cu128-x86_64-linux/layers.py +36 -0
  10. build/torch211-cxx11-cu128-x86_64-linux/metadata.json +5 -5
  11. build/torch211-cxx11-cu130-x86_64-linux/__init__.py +9 -0
  12. build/torch211-cxx11-cu130-x86_64-linux/{_fp8linear_cuda_91cad61.abi3.so → _fp8linear_cuda_999b53f.abi3.so} +2 -2
  13. build/torch211-cxx11-cu130-x86_64-linux/_ops.py +3 -3
  14. build/torch211-cxx11-cu130-x86_64-linux/layers.py +36 -0
  15. build/torch211-cxx11-cu130-x86_64-linux/metadata.json +5 -5
  16. build/torch212-cxx11-cu126-x86_64-linux/__init__.py +9 -0
  17. build/torch212-cxx11-cu126-x86_64-linux/{_fp8linear_cuda_91cad61.abi3.so → _fp8linear_cuda_999b53f.abi3.so} +2 -2
  18. build/torch212-cxx11-cu126-x86_64-linux/_ops.py +3 -3
  19. build/torch212-cxx11-cu126-x86_64-linux/layers.py +36 -0
  20. build/torch212-cxx11-cu126-x86_64-linux/metadata.json +5 -5
  21. build/torch212-cxx11-cu130-x86_64-linux/__init__.py +9 -0
  22. build/torch212-cxx11-cu130-x86_64-linux/_fp8linear_cuda_91cad61.abi3.so +0 -3
  23. build/torch212-cxx11-cu130-x86_64-linux/_fp8linear_cuda_999b53f.abi3.so +3 -0
  24. build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
  25. build/torch212-cxx11-cu130-x86_64-linux/layers.py +36 -0
  26. build/torch212-cxx11-cu130-x86_64-linux/metadata.json +5 -5
  27. build/torch212-cxx11-cu132-x86_64-linux/__init__.py +9 -0
  28. build/torch212-cxx11-cu132-x86_64-linux/_fp8linear_cuda_91cad61.abi3.so +0 -3
  29. build/torch212-cxx11-cu132-x86_64-linux/_fp8linear_cuda_999b53f.abi3.so +3 -0
  30. build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
  31. build/torch212-cxx11-cu132-x86_64-linux/layers.py +36 -0
  32. build/torch212-cxx11-cu132-x86_64-linux/metadata.json +5 -5
build/torch211-cxx11-cu126-x86_64-linux/__init__.py CHANGED
@@ -4,15 +4,19 @@ from .layers import (
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
 
7
  quantize_,
8
  quantize_weight,
9
  quantize_weight_mxfp8,
 
10
  )
11
 
12
  quantize_fp8 = ops.quantize_fp8
13
  fp8_gemm = ops.fp8_gemm
14
  quantize_mxfp8 = ops.quantize_mxfp8
15
  mxfp8_gemm = ops.mxfp8_gemm
 
 
16
 
17
  __all__ = [
18
  "ops",
@@ -28,4 +32,9 @@ __all__ = [
28
  "quantize_weight_mxfp8",
29
  "quantize_mxfp8",
30
  "mxfp8_gemm",
 
 
 
 
 
31
  ]
 
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
7
+ nvfp4_linear,
8
  quantize_,
9
  quantize_weight,
10
  quantize_weight_mxfp8,
11
+ quantize_weight_nvfp4,
12
  )
13
 
14
  quantize_fp8 = ops.quantize_fp8
15
  fp8_gemm = ops.fp8_gemm
16
  quantize_mxfp8 = ops.quantize_mxfp8
17
  mxfp8_gemm = ops.mxfp8_gemm
18
+ quantize_nvfp4 = ops.quantize_nvfp4
19
+ nvfp4_gemm = ops.nvfp4_gemm
20
 
21
  __all__ = [
22
  "ops",
 
32
  "quantize_weight_mxfp8",
33
  "quantize_mxfp8",
34
  "mxfp8_gemm",
35
+ # Blackwell NVFP4 path
36
+ "nvfp4_linear",
37
+ "quantize_weight_nvfp4",
38
+ "quantize_nvfp4",
39
+ "nvfp4_gemm",
40
  ]
build/torch211-cxx11-cu126-x86_64-linux/{_fp8linear_cuda_91cad61.abi3.so → _fp8linear_cuda_999b53f.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:16237c721fd76af1831b1f6b8676e4758c0a067059da4abfe24195e494360a54
3
- size 246032
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:bfbef942a42a1b843800e210491cf50b4d5e055765f8737ec7b333772753fe4d
3
+ size 296120
build/torch211-cxx11-cu126-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp8linear_cuda_91cad61
3
- ops = torch.ops._fp8linear_cuda_91cad61
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp8linear_cuda_91cad61::{op_name}"
 
1
  import torch
2
+ from . import _fp8linear_cuda_999b53f
3
+ ops = torch.ops._fp8linear_cuda_999b53f
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp8linear_cuda_999b53f::{op_name}"
build/torch211-cxx11-cu126-x86_64-linux/layers.py CHANGED
@@ -47,6 +47,25 @@ def _mxfp8_linear_fake(x, wq, w_scales, bias=None):
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
  MX_BLOCK = 32
51
 
52
 
@@ -67,6 +86,23 @@ def mxfp8_linear(x, wq, w_scales, bias=None):
67
  return ops.mxfp8_linear(x, wq, w_scales, bias)
68
 
69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70
  def quantize_weight(weight: torch.Tensor):
71
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
72
 
 
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
50
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_nvfp4"))
51
+ def _quantize_nvfp4_fake(x):
52
+ xq = x.new_empty((x.shape[0], x.shape[1] // 2), dtype=torch.uint8)
53
+ nrb = (x.shape[0] + 127) // 128
54
+ ncb = (x.shape[1] // 16 + 3) // 4
55
+ scales = x.new_empty((nrb * ncb * 512,), dtype=torch.uint8)
56
+ return xq, scales
57
+
58
+
59
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm"))
60
+ def _nvfp4_gemm_fake(xq, x_scales, wq, w_scales, bias=None):
61
+ return xq.new_empty((xq.shape[0], wq.shape[0]), dtype=torch.float16)
62
+
63
+
64
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_linear"))
65
+ def _nvfp4_linear_fake(x, wq, w_scales, bias=None):
66
+ return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
67
+
68
+
69
  MX_BLOCK = 32
70
 
71
 
 
86
  return ops.mxfp8_linear(x, wq, w_scales, bias)
87
 
88
 
89
+ def quantize_weight_nvfp4(weight: torch.Tensor):
90
+ """NVFP4 block-quantize a [N, K] weight (offline) via the CUDA kernel.
91
+
92
+ Returns (wq[N, K/2] packed e2m1 uint8, scales uint8/e4m3 in the tiled cuBLASLt
93
+ layout) for `nvfp4_gemm`/`nvfp4_linear`.
94
+ """
95
+ w = weight.detach()
96
+ if w.dtype not in (torch.float16, torch.bfloat16):
97
+ w = w.to(torch.bfloat16)
98
+ return ops.quantize_nvfp4(w.contiguous())
99
+
100
+
101
+ def nvfp4_linear(x, wq, w_scales, bias=None):
102
+ """NVFP4 linear (Blackwell sm_100+). `wq`/`w_scales` from quantize_weight_nvfp4."""
103
+ return ops.nvfp4_linear(x, wq, w_scales, bias)
104
+
105
+
106
  def quantize_weight(weight: torch.Tensor):
107
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
108
 
build/torch211-cxx11-cu126-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp8linear",
3
- "id": "_fp8linear_cuda_91cad61",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
@@ -14,11 +14,11 @@
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
- "__init__.py": "Ui5JQop95S6F9SugXfDwNRZya/rLpxJepzhCFp+3Co8=",
18
- "_fp8linear_cuda_91cad61.abi3.so": "FiN8ch/XavGDGx9rhnbkdYwKBnBZ2kq/4kGV5JQ2ClQ=",
19
- "_ops.py": "dPGt3db4u3Jj6m6kXzb4hLHnA2Narw1cOO+50c7lg4k=",
20
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
21
- "layers.py": "lIe+IYmTAunoNiYdUkHLDjguDQvMsWV6p4koNPTjuo0="
22
  }
23
  }
24
  }
 
1
  {
2
  "name": "fp8linear",
3
+ "id": "_fp8linear_cuda_999b53f",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
 
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
+ "__init__.py": "kuYXA/yMzeG8mUeRuFseTWi5nnWR13ZFrBVvAZqJPeY=",
18
+ "_fp8linear_cuda_999b53f.abi3.so": "v775QqQqG4Q4AOIQSRz1C01eBVdl+HN+x7MzdydT/k0=",
19
+ "_ops.py": "DKOF2voEdv6S/MArdkqK14wb9yDH06Fmy9QNdlQ4Kgo=",
20
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
21
+ "layers.py": "6JzuLWloloQ29zkTO9Dv9jJX4aV/yrFU5MbkJqiNr5w="
22
  }
23
  }
24
  }
build/torch211-cxx11-cu128-x86_64-linux/__init__.py CHANGED
@@ -4,15 +4,19 @@ from .layers import (
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
 
7
  quantize_,
8
  quantize_weight,
9
  quantize_weight_mxfp8,
 
10
  )
11
 
12
  quantize_fp8 = ops.quantize_fp8
13
  fp8_gemm = ops.fp8_gemm
14
  quantize_mxfp8 = ops.quantize_mxfp8
15
  mxfp8_gemm = ops.mxfp8_gemm
 
 
16
 
17
  __all__ = [
18
  "ops",
@@ -28,4 +32,9 @@ __all__ = [
28
  "quantize_weight_mxfp8",
29
  "quantize_mxfp8",
30
  "mxfp8_gemm",
 
 
 
 
 
31
  ]
 
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
7
+ nvfp4_linear,
8
  quantize_,
9
  quantize_weight,
10
  quantize_weight_mxfp8,
11
+ quantize_weight_nvfp4,
12
  )
13
 
14
  quantize_fp8 = ops.quantize_fp8
15
  fp8_gemm = ops.fp8_gemm
16
  quantize_mxfp8 = ops.quantize_mxfp8
17
  mxfp8_gemm = ops.mxfp8_gemm
18
+ quantize_nvfp4 = ops.quantize_nvfp4
19
+ nvfp4_gemm = ops.nvfp4_gemm
20
 
21
  __all__ = [
22
  "ops",
 
32
  "quantize_weight_mxfp8",
33
  "quantize_mxfp8",
34
  "mxfp8_gemm",
35
+ # Blackwell NVFP4 path
36
+ "nvfp4_linear",
37
+ "quantize_weight_nvfp4",
38
+ "quantize_nvfp4",
39
+ "nvfp4_gemm",
40
  ]
build/torch211-cxx11-cu128-x86_64-linux/{_fp8linear_cuda_91cad61.abi3.so → _fp8linear_cuda_999b53f.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:7a1f8d4b93aaf7cee8b7eee3a0d52fa9ab81c0b4c360c29e6196142fc880e121
3
- size 409592
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:60f6f8c7db62207a53928f0cdd0dee6053a5df728d1796e960def6731ff46bba
3
+ size 508760
build/torch211-cxx11-cu128-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp8linear_cuda_91cad61
3
- ops = torch.ops._fp8linear_cuda_91cad61
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp8linear_cuda_91cad61::{op_name}"
 
1
  import torch
2
+ from . import _fp8linear_cuda_999b53f
3
+ ops = torch.ops._fp8linear_cuda_999b53f
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp8linear_cuda_999b53f::{op_name}"
build/torch211-cxx11-cu128-x86_64-linux/layers.py CHANGED
@@ -47,6 +47,25 @@ def _mxfp8_linear_fake(x, wq, w_scales, bias=None):
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
  MX_BLOCK = 32
51
 
52
 
@@ -67,6 +86,23 @@ def mxfp8_linear(x, wq, w_scales, bias=None):
67
  return ops.mxfp8_linear(x, wq, w_scales, bias)
68
 
69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70
  def quantize_weight(weight: torch.Tensor):
71
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
72
 
 
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
50
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_nvfp4"))
51
+ def _quantize_nvfp4_fake(x):
52
+ xq = x.new_empty((x.shape[0], x.shape[1] // 2), dtype=torch.uint8)
53
+ nrb = (x.shape[0] + 127) // 128
54
+ ncb = (x.shape[1] // 16 + 3) // 4
55
+ scales = x.new_empty((nrb * ncb * 512,), dtype=torch.uint8)
56
+ return xq, scales
57
+
58
+
59
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm"))
60
+ def _nvfp4_gemm_fake(xq, x_scales, wq, w_scales, bias=None):
61
+ return xq.new_empty((xq.shape[0], wq.shape[0]), dtype=torch.float16)
62
+
63
+
64
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_linear"))
65
+ def _nvfp4_linear_fake(x, wq, w_scales, bias=None):
66
+ return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
67
+
68
+
69
  MX_BLOCK = 32
70
 
71
 
 
86
  return ops.mxfp8_linear(x, wq, w_scales, bias)
87
 
88
 
89
+ def quantize_weight_nvfp4(weight: torch.Tensor):
90
+ """NVFP4 block-quantize a [N, K] weight (offline) via the CUDA kernel.
91
+
92
+ Returns (wq[N, K/2] packed e2m1 uint8, scales uint8/e4m3 in the tiled cuBLASLt
93
+ layout) for `nvfp4_gemm`/`nvfp4_linear`.
94
+ """
95
+ w = weight.detach()
96
+ if w.dtype not in (torch.float16, torch.bfloat16):
97
+ w = w.to(torch.bfloat16)
98
+ return ops.quantize_nvfp4(w.contiguous())
99
+
100
+
101
+ def nvfp4_linear(x, wq, w_scales, bias=None):
102
+ """NVFP4 linear (Blackwell sm_100+). `wq`/`w_scales` from quantize_weight_nvfp4."""
103
+ return ops.nvfp4_linear(x, wq, w_scales, bias)
104
+
105
+
106
  def quantize_weight(weight: torch.Tensor):
107
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
108
 
build/torch211-cxx11-cu128-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp8linear",
3
- "id": "_fp8linear_cuda_91cad61",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
@@ -16,11 +16,11 @@
16
  "digest": {
17
  "algorithm": "sha256",
18
  "files": {
19
- "__init__.py": "Ui5JQop95S6F9SugXfDwNRZya/rLpxJepzhCFp+3Co8=",
20
- "_fp8linear_cuda_91cad61.abi3.so": "eh+NS5Oq987ot+7joNUvqauBwLTDYMKeYZYUL8iA4SE=",
21
- "_ops.py": "dPGt3db4u3Jj6m6kXzb4hLHnA2Narw1cOO+50c7lg4k=",
22
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
23
- "layers.py": "lIe+IYmTAunoNiYdUkHLDjguDQvMsWV6p4koNPTjuo0="
24
  }
25
  }
26
  }
 
1
  {
2
  "name": "fp8linear",
3
+ "id": "_fp8linear_cuda_999b53f",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
 
16
  "digest": {
17
  "algorithm": "sha256",
18
  "files": {
19
+ "__init__.py": "kuYXA/yMzeG8mUeRuFseTWi5nnWR13ZFrBVvAZqJPeY=",
20
+ "_fp8linear_cuda_999b53f.abi3.so": "YPb4x9tiIHpTko8M3Q3uYFOl33KNF5bpYN72cx/0a7o=",
21
+ "_ops.py": "DKOF2voEdv6S/MArdkqK14wb9yDH06Fmy9QNdlQ4Kgo=",
22
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
23
+ "layers.py": "6JzuLWloloQ29zkTO9Dv9jJX4aV/yrFU5MbkJqiNr5w="
24
  }
25
  }
26
  }
build/torch211-cxx11-cu130-x86_64-linux/__init__.py CHANGED
@@ -4,15 +4,19 @@ from .layers import (
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
 
7
  quantize_,
8
  quantize_weight,
9
  quantize_weight_mxfp8,
 
10
  )
11
 
12
  quantize_fp8 = ops.quantize_fp8
13
  fp8_gemm = ops.fp8_gemm
14
  quantize_mxfp8 = ops.quantize_mxfp8
15
  mxfp8_gemm = ops.mxfp8_gemm
 
 
16
 
17
  __all__ = [
18
  "ops",
@@ -28,4 +32,9 @@ __all__ = [
28
  "quantize_weight_mxfp8",
29
  "quantize_mxfp8",
30
  "mxfp8_gemm",
 
 
 
 
 
31
  ]
 
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
7
+ nvfp4_linear,
8
  quantize_,
9
  quantize_weight,
10
  quantize_weight_mxfp8,
11
+ quantize_weight_nvfp4,
12
  )
13
 
14
  quantize_fp8 = ops.quantize_fp8
15
  fp8_gemm = ops.fp8_gemm
16
  quantize_mxfp8 = ops.quantize_mxfp8
17
  mxfp8_gemm = ops.mxfp8_gemm
18
+ quantize_nvfp4 = ops.quantize_nvfp4
19
+ nvfp4_gemm = ops.nvfp4_gemm
20
 
21
  __all__ = [
22
  "ops",
 
32
  "quantize_weight_mxfp8",
33
  "quantize_mxfp8",
34
  "mxfp8_gemm",
35
+ # Blackwell NVFP4 path
36
+ "nvfp4_linear",
37
+ "quantize_weight_nvfp4",
38
+ "quantize_nvfp4",
39
+ "nvfp4_gemm",
40
  ]
build/torch211-cxx11-cu130-x86_64-linux/{_fp8linear_cuda_91cad61.abi3.so → _fp8linear_cuda_999b53f.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:a9939e8e918aef3b054c9f8dcf85b1155241e6799eee9b904a7d78010c069c10
3
- size 528752
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:53262026dcbfbc64f053acfe84f5d778b2627c8f7168ca1049854ac043b36881
3
+ size 653816
build/torch211-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp8linear_cuda_91cad61
3
- ops = torch.ops._fp8linear_cuda_91cad61
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp8linear_cuda_91cad61::{op_name}"
 
1
  import torch
2
+ from . import _fp8linear_cuda_999b53f
3
+ ops = torch.ops._fp8linear_cuda_999b53f
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp8linear_cuda_999b53f::{op_name}"
build/torch211-cxx11-cu130-x86_64-linux/layers.py CHANGED
@@ -47,6 +47,25 @@ def _mxfp8_linear_fake(x, wq, w_scales, bias=None):
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
  MX_BLOCK = 32
51
 
52
 
@@ -67,6 +86,23 @@ def mxfp8_linear(x, wq, w_scales, bias=None):
67
  return ops.mxfp8_linear(x, wq, w_scales, bias)
68
 
69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70
  def quantize_weight(weight: torch.Tensor):
71
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
72
 
 
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
50
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_nvfp4"))
51
+ def _quantize_nvfp4_fake(x):
52
+ xq = x.new_empty((x.shape[0], x.shape[1] // 2), dtype=torch.uint8)
53
+ nrb = (x.shape[0] + 127) // 128
54
+ ncb = (x.shape[1] // 16 + 3) // 4
55
+ scales = x.new_empty((nrb * ncb * 512,), dtype=torch.uint8)
56
+ return xq, scales
57
+
58
+
59
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm"))
60
+ def _nvfp4_gemm_fake(xq, x_scales, wq, w_scales, bias=None):
61
+ return xq.new_empty((xq.shape[0], wq.shape[0]), dtype=torch.float16)
62
+
63
+
64
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_linear"))
65
+ def _nvfp4_linear_fake(x, wq, w_scales, bias=None):
66
+ return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
67
+
68
+
69
  MX_BLOCK = 32
70
 
71
 
 
86
  return ops.mxfp8_linear(x, wq, w_scales, bias)
87
 
88
 
89
+ def quantize_weight_nvfp4(weight: torch.Tensor):
90
+ """NVFP4 block-quantize a [N, K] weight (offline) via the CUDA kernel.
91
+
92
+ Returns (wq[N, K/2] packed e2m1 uint8, scales uint8/e4m3 in the tiled cuBLASLt
93
+ layout) for `nvfp4_gemm`/`nvfp4_linear`.
94
+ """
95
+ w = weight.detach()
96
+ if w.dtype not in (torch.float16, torch.bfloat16):
97
+ w = w.to(torch.bfloat16)
98
+ return ops.quantize_nvfp4(w.contiguous())
99
+
100
+
101
+ def nvfp4_linear(x, wq, w_scales, bias=None):
102
+ """NVFP4 linear (Blackwell sm_100+). `wq`/`w_scales` from quantize_weight_nvfp4."""
103
+ return ops.nvfp4_linear(x, wq, w_scales, bias)
104
+
105
+
106
  def quantize_weight(weight: torch.Tensor):
107
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
108
 
build/torch211-cxx11-cu130-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp8linear",
3
- "id": "_fp8linear_cuda_91cad61",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
@@ -17,11 +17,11 @@
17
  "digest": {
18
  "algorithm": "sha256",
19
  "files": {
20
- "__init__.py": "Ui5JQop95S6F9SugXfDwNRZya/rLpxJepzhCFp+3Co8=",
21
- "_fp8linear_cuda_91cad61.abi3.so": "qZOejpGK7zsFTJ+Nz4WxFVJB5nme7puQSn14AQwGnBA=",
22
- "_ops.py": "dPGt3db4u3Jj6m6kXzb4hLHnA2Narw1cOO+50c7lg4k=",
23
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
24
- "layers.py": "lIe+IYmTAunoNiYdUkHLDjguDQvMsWV6p4koNPTjuo0="
25
  }
26
  }
27
  }
 
1
  {
2
  "name": "fp8linear",
3
+ "id": "_fp8linear_cuda_999b53f",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
 
17
  "digest": {
18
  "algorithm": "sha256",
19
  "files": {
20
+ "__init__.py": "kuYXA/yMzeG8mUeRuFseTWi5nnWR13ZFrBVvAZqJPeY=",
21
+ "_fp8linear_cuda_999b53f.abi3.so": "UyYgJty/vGTwU6z+hPXXeLJifI9xaMoQSYVKwEOzaIE=",
22
+ "_ops.py": "DKOF2voEdv6S/MArdkqK14wb9yDH06Fmy9QNdlQ4Kgo=",
23
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
24
+ "layers.py": "6JzuLWloloQ29zkTO9Dv9jJX4aV/yrFU5MbkJqiNr5w="
25
  }
26
  }
27
  }
build/torch212-cxx11-cu126-x86_64-linux/__init__.py CHANGED
@@ -4,15 +4,19 @@ from .layers import (
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
 
7
  quantize_,
8
  quantize_weight,
9
  quantize_weight_mxfp8,
 
10
  )
11
 
12
  quantize_fp8 = ops.quantize_fp8
13
  fp8_gemm = ops.fp8_gemm
14
  quantize_mxfp8 = ops.quantize_mxfp8
15
  mxfp8_gemm = ops.mxfp8_gemm
 
 
16
 
17
  __all__ = [
18
  "ops",
@@ -28,4 +32,9 @@ __all__ = [
28
  "quantize_weight_mxfp8",
29
  "quantize_mxfp8",
30
  "mxfp8_gemm",
 
 
 
 
 
31
  ]
 
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
7
+ nvfp4_linear,
8
  quantize_,
9
  quantize_weight,
10
  quantize_weight_mxfp8,
11
+ quantize_weight_nvfp4,
12
  )
13
 
14
  quantize_fp8 = ops.quantize_fp8
15
  fp8_gemm = ops.fp8_gemm
16
  quantize_mxfp8 = ops.quantize_mxfp8
17
  mxfp8_gemm = ops.mxfp8_gemm
18
+ quantize_nvfp4 = ops.quantize_nvfp4
19
+ nvfp4_gemm = ops.nvfp4_gemm
20
 
21
  __all__ = [
22
  "ops",
 
32
  "quantize_weight_mxfp8",
33
  "quantize_mxfp8",
34
  "mxfp8_gemm",
35
+ # Blackwell NVFP4 path
36
+ "nvfp4_linear",
37
+ "quantize_weight_nvfp4",
38
+ "quantize_nvfp4",
39
+ "nvfp4_gemm",
40
  ]
build/torch212-cxx11-cu126-x86_64-linux/{_fp8linear_cuda_91cad61.abi3.so → _fp8linear_cuda_999b53f.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:8a6f17eda5748e6e85b7dc72410ea70db6911cf396f79c07a0b72fdae6adb0b9
3
- size 245248
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6fe5c7462cddd5f946175249b544fa71192e1001d34f28b465de3ba032a6432e
3
+ size 291240
build/torch212-cxx11-cu126-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp8linear_cuda_91cad61
3
- ops = torch.ops._fp8linear_cuda_91cad61
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp8linear_cuda_91cad61::{op_name}"
 
1
  import torch
2
+ from . import _fp8linear_cuda_999b53f
3
+ ops = torch.ops._fp8linear_cuda_999b53f
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp8linear_cuda_999b53f::{op_name}"
build/torch212-cxx11-cu126-x86_64-linux/layers.py CHANGED
@@ -47,6 +47,25 @@ def _mxfp8_linear_fake(x, wq, w_scales, bias=None):
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
  MX_BLOCK = 32
51
 
52
 
@@ -67,6 +86,23 @@ def mxfp8_linear(x, wq, w_scales, bias=None):
67
  return ops.mxfp8_linear(x, wq, w_scales, bias)
68
 
69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70
  def quantize_weight(weight: torch.Tensor):
71
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
72
 
 
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
50
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_nvfp4"))
51
+ def _quantize_nvfp4_fake(x):
52
+ xq = x.new_empty((x.shape[0], x.shape[1] // 2), dtype=torch.uint8)
53
+ nrb = (x.shape[0] + 127) // 128
54
+ ncb = (x.shape[1] // 16 + 3) // 4
55
+ scales = x.new_empty((nrb * ncb * 512,), dtype=torch.uint8)
56
+ return xq, scales
57
+
58
+
59
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm"))
60
+ def _nvfp4_gemm_fake(xq, x_scales, wq, w_scales, bias=None):
61
+ return xq.new_empty((xq.shape[0], wq.shape[0]), dtype=torch.float16)
62
+
63
+
64
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_linear"))
65
+ def _nvfp4_linear_fake(x, wq, w_scales, bias=None):
66
+ return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
67
+
68
+
69
  MX_BLOCK = 32
70
 
71
 
 
86
  return ops.mxfp8_linear(x, wq, w_scales, bias)
87
 
88
 
89
+ def quantize_weight_nvfp4(weight: torch.Tensor):
90
+ """NVFP4 block-quantize a [N, K] weight (offline) via the CUDA kernel.
91
+
92
+ Returns (wq[N, K/2] packed e2m1 uint8, scales uint8/e4m3 in the tiled cuBLASLt
93
+ layout) for `nvfp4_gemm`/`nvfp4_linear`.
94
+ """
95
+ w = weight.detach()
96
+ if w.dtype not in (torch.float16, torch.bfloat16):
97
+ w = w.to(torch.bfloat16)
98
+ return ops.quantize_nvfp4(w.contiguous())
99
+
100
+
101
+ def nvfp4_linear(x, wq, w_scales, bias=None):
102
+ """NVFP4 linear (Blackwell sm_100+). `wq`/`w_scales` from quantize_weight_nvfp4."""
103
+ return ops.nvfp4_linear(x, wq, w_scales, bias)
104
+
105
+
106
  def quantize_weight(weight: torch.Tensor):
107
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
108
 
build/torch212-cxx11-cu126-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp8linear",
3
- "id": "_fp8linear_cuda_91cad61",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
@@ -14,11 +14,11 @@
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
- "__init__.py": "Ui5JQop95S6F9SugXfDwNRZya/rLpxJepzhCFp+3Co8=",
18
- "_fp8linear_cuda_91cad61.abi3.so": "im8X7aV0jm6Ft9xyQQ6nDbaRHPOW95wHoLcv2uatsLk=",
19
- "_ops.py": "dPGt3db4u3Jj6m6kXzb4hLHnA2Narw1cOO+50c7lg4k=",
20
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
21
- "layers.py": "lIe+IYmTAunoNiYdUkHLDjguDQvMsWV6p4koNPTjuo0="
22
  }
23
  }
24
  }
 
1
  {
2
  "name": "fp8linear",
3
+ "id": "_fp8linear_cuda_999b53f",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
 
14
  "digest": {
15
  "algorithm": "sha256",
16
  "files": {
17
+ "__init__.py": "kuYXA/yMzeG8mUeRuFseTWi5nnWR13ZFrBVvAZqJPeY=",
18
+ "_fp8linear_cuda_999b53f.abi3.so": "b+XHRizd1flGF1JJtUT6cRkuEAHTTyi0Zd47oDKmQy4=",
19
+ "_ops.py": "DKOF2voEdv6S/MArdkqK14wb9yDH06Fmy9QNdlQ4Kgo=",
20
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
21
+ "layers.py": "6JzuLWloloQ29zkTO9Dv9jJX4aV/yrFU5MbkJqiNr5w="
22
  }
23
  }
24
  }
build/torch212-cxx11-cu130-x86_64-linux/__init__.py CHANGED
@@ -4,15 +4,19 @@ from .layers import (
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
 
7
  quantize_,
8
  quantize_weight,
9
  quantize_weight_mxfp8,
 
10
  )
11
 
12
  quantize_fp8 = ops.quantize_fp8
13
  fp8_gemm = ops.fp8_gemm
14
  quantize_mxfp8 = ops.quantize_mxfp8
15
  mxfp8_gemm = ops.mxfp8_gemm
 
 
16
 
17
  __all__ = [
18
  "ops",
@@ -28,4 +32,9 @@ __all__ = [
28
  "quantize_weight_mxfp8",
29
  "quantize_mxfp8",
30
  "mxfp8_gemm",
 
 
 
 
 
31
  ]
 
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
7
+ nvfp4_linear,
8
  quantize_,
9
  quantize_weight,
10
  quantize_weight_mxfp8,
11
+ quantize_weight_nvfp4,
12
  )
13
 
14
  quantize_fp8 = ops.quantize_fp8
15
  fp8_gemm = ops.fp8_gemm
16
  quantize_mxfp8 = ops.quantize_mxfp8
17
  mxfp8_gemm = ops.mxfp8_gemm
18
+ quantize_nvfp4 = ops.quantize_nvfp4
19
+ nvfp4_gemm = ops.nvfp4_gemm
20
 
21
  __all__ = [
22
  "ops",
 
32
  "quantize_weight_mxfp8",
33
  "quantize_mxfp8",
34
  "mxfp8_gemm",
35
+ # Blackwell NVFP4 path
36
+ "nvfp4_linear",
37
+ "quantize_weight_nvfp4",
38
+ "quantize_nvfp4",
39
+ "nvfp4_gemm",
40
  ]
build/torch212-cxx11-cu130-x86_64-linux/_fp8linear_cuda_91cad61.abi3.so DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:3c3c3adab2d82001ef96d31278c5e5134638894995c25439db9b0f433d79f37a
3
- size 528400
 
 
 
 
build/torch212-cxx11-cu130-x86_64-linux/_fp8linear_cuda_999b53f.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5714d408fc21891bf66b73ae3302cf026a62cef9ab6cfcf2c71cd703ef40db9a
3
+ size 653472
build/torch212-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp8linear_cuda_91cad61
3
- ops = torch.ops._fp8linear_cuda_91cad61
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp8linear_cuda_91cad61::{op_name}"
 
1
  import torch
2
+ from . import _fp8linear_cuda_999b53f
3
+ ops = torch.ops._fp8linear_cuda_999b53f
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp8linear_cuda_999b53f::{op_name}"
build/torch212-cxx11-cu130-x86_64-linux/layers.py CHANGED
@@ -47,6 +47,25 @@ def _mxfp8_linear_fake(x, wq, w_scales, bias=None):
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
  MX_BLOCK = 32
51
 
52
 
@@ -67,6 +86,23 @@ def mxfp8_linear(x, wq, w_scales, bias=None):
67
  return ops.mxfp8_linear(x, wq, w_scales, bias)
68
 
69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70
  def quantize_weight(weight: torch.Tensor):
71
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
72
 
 
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
50
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_nvfp4"))
51
+ def _quantize_nvfp4_fake(x):
52
+ xq = x.new_empty((x.shape[0], x.shape[1] // 2), dtype=torch.uint8)
53
+ nrb = (x.shape[0] + 127) // 128
54
+ ncb = (x.shape[1] // 16 + 3) // 4
55
+ scales = x.new_empty((nrb * ncb * 512,), dtype=torch.uint8)
56
+ return xq, scales
57
+
58
+
59
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm"))
60
+ def _nvfp4_gemm_fake(xq, x_scales, wq, w_scales, bias=None):
61
+ return xq.new_empty((xq.shape[0], wq.shape[0]), dtype=torch.float16)
62
+
63
+
64
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_linear"))
65
+ def _nvfp4_linear_fake(x, wq, w_scales, bias=None):
66
+ return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
67
+
68
+
69
  MX_BLOCK = 32
70
 
71
 
 
86
  return ops.mxfp8_linear(x, wq, w_scales, bias)
87
 
88
 
89
+ def quantize_weight_nvfp4(weight: torch.Tensor):
90
+ """NVFP4 block-quantize a [N, K] weight (offline) via the CUDA kernel.
91
+
92
+ Returns (wq[N, K/2] packed e2m1 uint8, scales uint8/e4m3 in the tiled cuBLASLt
93
+ layout) for `nvfp4_gemm`/`nvfp4_linear`.
94
+ """
95
+ w = weight.detach()
96
+ if w.dtype not in (torch.float16, torch.bfloat16):
97
+ w = w.to(torch.bfloat16)
98
+ return ops.quantize_nvfp4(w.contiguous())
99
+
100
+
101
+ def nvfp4_linear(x, wq, w_scales, bias=None):
102
+ """NVFP4 linear (Blackwell sm_100+). `wq`/`w_scales` from quantize_weight_nvfp4."""
103
+ return ops.nvfp4_linear(x, wq, w_scales, bias)
104
+
105
+
106
  def quantize_weight(weight: torch.Tensor):
107
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
108
 
build/torch212-cxx11-cu130-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp8linear",
3
- "id": "_fp8linear_cuda_91cad61",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
@@ -17,11 +17,11 @@
17
  "digest": {
18
  "algorithm": "sha256",
19
  "files": {
20
- "__init__.py": "Ui5JQop95S6F9SugXfDwNRZya/rLpxJepzhCFp+3Co8=",
21
- "_fp8linear_cuda_91cad61.abi3.so": "PDw62rLYIAHvltMSeMXlE0Y4iUmVwlQ525sPQz1583o=",
22
- "_ops.py": "dPGt3db4u3Jj6m6kXzb4hLHnA2Narw1cOO+50c7lg4k=",
23
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
24
- "layers.py": "lIe+IYmTAunoNiYdUkHLDjguDQvMsWV6p4koNPTjuo0="
25
  }
26
  }
27
  }
 
1
  {
2
  "name": "fp8linear",
3
+ "id": "_fp8linear_cuda_999b53f",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
 
17
  "digest": {
18
  "algorithm": "sha256",
19
  "files": {
20
+ "__init__.py": "kuYXA/yMzeG8mUeRuFseTWi5nnWR13ZFrBVvAZqJPeY=",
21
+ "_fp8linear_cuda_999b53f.abi3.so": "VxTUCPwhiRv2a3OuMwLPAmpizvmrbPzyxxzXA+9A25o=",
22
+ "_ops.py": "DKOF2voEdv6S/MArdkqK14wb9yDH06Fmy9QNdlQ4Kgo=",
23
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
24
+ "layers.py": "6JzuLWloloQ29zkTO9Dv9jJX4aV/yrFU5MbkJqiNr5w="
25
  }
26
  }
27
  }
build/torch212-cxx11-cu132-x86_64-linux/__init__.py CHANGED
@@ -4,15 +4,19 @@ from .layers import (
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
 
7
  quantize_,
8
  quantize_weight,
9
  quantize_weight_mxfp8,
 
10
  )
11
 
12
  quantize_fp8 = ops.quantize_fp8
13
  fp8_gemm = ops.fp8_gemm
14
  quantize_mxfp8 = ops.quantize_mxfp8
15
  mxfp8_gemm = ops.mxfp8_gemm
 
 
16
 
17
  __all__ = [
18
  "ops",
@@ -28,4 +32,9 @@ __all__ = [
28
  "quantize_weight_mxfp8",
29
  "quantize_mxfp8",
30
  "mxfp8_gemm",
 
 
 
 
 
31
  ]
 
4
  Fp8Linear,
5
  fp8_linear,
6
  mxfp8_linear,
7
+ nvfp4_linear,
8
  quantize_,
9
  quantize_weight,
10
  quantize_weight_mxfp8,
11
+ quantize_weight_nvfp4,
12
  )
13
 
14
  quantize_fp8 = ops.quantize_fp8
15
  fp8_gemm = ops.fp8_gemm
16
  quantize_mxfp8 = ops.quantize_mxfp8
17
  mxfp8_gemm = ops.mxfp8_gemm
18
+ quantize_nvfp4 = ops.quantize_nvfp4
19
+ nvfp4_gemm = ops.nvfp4_gemm
20
 
21
  __all__ = [
22
  "ops",
 
32
  "quantize_weight_mxfp8",
33
  "quantize_mxfp8",
34
  "mxfp8_gemm",
35
+ # Blackwell NVFP4 path
36
+ "nvfp4_linear",
37
+ "quantize_weight_nvfp4",
38
+ "quantize_nvfp4",
39
+ "nvfp4_gemm",
40
  ]
build/torch212-cxx11-cu132-x86_64-linux/_fp8linear_cuda_91cad61.abi3.so DELETED
@@ -1,3 +0,0 @@
1
- version https://git-lfs.github.com/spec/v1
2
- oid sha256:dd694ce980ce13b0e67f67346c56be0f91a0ff134483ff60e36eca751e83f814
3
- size 600152
 
 
 
 
build/torch212-cxx11-cu132-x86_64-linux/_fp8linear_cuda_999b53f.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c5212b7322375abb616c64a440e138d35f41a3a6308ef0c4958c94ccd534a1dd
3
+ size 725224
build/torch212-cxx11-cu132-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp8linear_cuda_91cad61
3
- ops = torch.ops._fp8linear_cuda_91cad61
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp8linear_cuda_91cad61::{op_name}"
 
1
  import torch
2
+ from . import _fp8linear_cuda_999b53f
3
+ ops = torch.ops._fp8linear_cuda_999b53f
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp8linear_cuda_999b53f::{op_name}"
build/torch212-cxx11-cu132-x86_64-linux/layers.py CHANGED
@@ -47,6 +47,25 @@ def _mxfp8_linear_fake(x, wq, w_scales, bias=None):
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
50
  MX_BLOCK = 32
51
 
52
 
@@ -67,6 +86,23 @@ def mxfp8_linear(x, wq, w_scales, bias=None):
67
  return ops.mxfp8_linear(x, wq, w_scales, bias)
68
 
69
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70
  def quantize_weight(weight: torch.Tensor):
71
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
72
 
 
47
  return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
48
 
49
 
50
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_nvfp4"))
51
+ def _quantize_nvfp4_fake(x):
52
+ xq = x.new_empty((x.shape[0], x.shape[1] // 2), dtype=torch.uint8)
53
+ nrb = (x.shape[0] + 127) // 128
54
+ ncb = (x.shape[1] // 16 + 3) // 4
55
+ scales = x.new_empty((nrb * ncb * 512,), dtype=torch.uint8)
56
+ return xq, scales
57
+
58
+
59
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm"))
60
+ def _nvfp4_gemm_fake(xq, x_scales, wq, w_scales, bias=None):
61
+ return xq.new_empty((xq.shape[0], wq.shape[0]), dtype=torch.float16)
62
+
63
+
64
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_linear"))
65
+ def _nvfp4_linear_fake(x, wq, w_scales, bias=None):
66
+ return x.new_empty((*x.shape[:-1], wq.shape[0]), dtype=torch.float16)
67
+
68
+
69
  MX_BLOCK = 32
70
 
71
 
 
86
  return ops.mxfp8_linear(x, wq, w_scales, bias)
87
 
88
 
89
+ def quantize_weight_nvfp4(weight: torch.Tensor):
90
+ """NVFP4 block-quantize a [N, K] weight (offline) via the CUDA kernel.
91
+
92
+ Returns (wq[N, K/2] packed e2m1 uint8, scales uint8/e4m3 in the tiled cuBLASLt
93
+ layout) for `nvfp4_gemm`/`nvfp4_linear`.
94
+ """
95
+ w = weight.detach()
96
+ if w.dtype not in (torch.float16, torch.bfloat16):
97
+ w = w.to(torch.bfloat16)
98
+ return ops.quantize_nvfp4(w.contiguous())
99
+
100
+
101
+ def nvfp4_linear(x, wq, w_scales, bias=None):
102
+ """NVFP4 linear (Blackwell sm_100+). `wq`/`w_scales` from quantize_weight_nvfp4."""
103
+ return ops.nvfp4_linear(x, wq, w_scales, bias)
104
+
105
+
106
  def quantize_weight(weight: torch.Tensor):
107
  """Arch-aware weight quantization (the choice lives in C++ `ops.quantize_weight`).
108
 
build/torch212-cxx11-cu132-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp8linear",
3
- "id": "_fp8linear_cuda_91cad61",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
@@ -17,11 +17,11 @@
17
  "digest": {
18
  "algorithm": "sha256",
19
  "files": {
20
- "__init__.py": "Ui5JQop95S6F9SugXfDwNRZya/rLpxJepzhCFp+3Co8=",
21
- "_fp8linear_cuda_91cad61.abi3.so": "3WlM6YDOE7Dmf2c0bFa+D5Gg/xNEg/9g427KdR6D+BQ=",
22
- "_ops.py": "dPGt3db4u3Jj6m6kXzb4hLHnA2Narw1cOO+50c7lg4k=",
23
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
24
- "layers.py": "lIe+IYmTAunoNiYdUkHLDjguDQvMsWV6p4koNPTjuo0="
25
  }
26
  }
27
  }
 
1
  {
2
  "name": "fp8linear",
3
+ "id": "_fp8linear_cuda_999b53f",
4
  "version": 1,
5
  "license": "MIT",
6
  "python-depends": [],
 
17
  "digest": {
18
  "algorithm": "sha256",
19
  "files": {
20
+ "__init__.py": "kuYXA/yMzeG8mUeRuFseTWi5nnWR13ZFrBVvAZqJPeY=",
21
+ "_fp8linear_cuda_999b53f.abi3.so": "xSErcyI3WrthbGSkQOE4019Bo6YwjvDElYyUzNU0od0=",
22
+ "_ops.py": "DKOF2voEdv6S/MArdkqK14wb9yDH06Fmy9QNdlQ4Kgo=",
23
  "fp8linear/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
24
+ "layers.py": "6JzuLWloloQ29zkTO9Dv9jJX4aV/yrFU5MbkJqiNr5w="
25
  }
26
  }
27
  }