Hugging Face's logo

Kernels:
flashrt
/
fp4-gemm

Kernel card Files
xet
 Community
liangsu9988 commited on
Commit
bf1cd65
·
verified ·
1 Parent(s): 79ce0b9

Uploaded using `kernel-builder`.

Browse files
Files changed (27) hide show
  1. benchmarks/benchmark.py +21 -8
  2. build/torch211-cxx11-cu128-x86_64-linux/__init__.py +39 -3
  3. build/torch211-cxx11-cu128-x86_64-linux/{_fp4_gemm_cuda_7e2e5b9.abi3.so → _fp4_gemm_cuda_752e924.abi3.so} +2 -2
  4. build/torch211-cxx11-cu128-x86_64-linux/_ops.py +3 -3
  5. build/torch211-cxx11-cu128-x86_64-linux/metadata.json +15 -4
  6. build/torch211-cxx11-cu130-x86_64-linux/__init__.py +39 -3
  7. build/torch211-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_7e2e5b9.abi3.so → _fp4_gemm_cuda_752e924.abi3.so} +2 -2
  8. build/torch211-cxx11-cu130-x86_64-linux/_ops.py +3 -3
  9. build/torch211-cxx11-cu130-x86_64-linux/metadata.json +15 -4
  10. build/torch212-cxx11-cu130-x86_64-linux/__init__.py +39 -3
  11. build/torch212-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_7e2e5b9.abi3.so → _fp4_gemm_cuda_752e924.abi3.so} +2 -2
  12. build/torch212-cxx11-cu130-x86_64-linux/_ops.py +3 -3
  13. build/torch212-cxx11-cu130-x86_64-linux/metadata.json +15 -4
  14. build/torch212-cxx11-cu132-x86_64-linux/__init__.py +39 -3
  15. build/torch212-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_7e2e5b9.abi3.so → _fp4_gemm_cuda_752e924.abi3.so} +2 -2
  16. build/torch212-cxx11-cu132-x86_64-linux/_ops.py +3 -3
  17. build/torch212-cxx11-cu132-x86_64-linux/metadata.json +15 -4
  18. build/torch213-cxx11-cu130-x86_64-linux/__init__.py +122 -0
  19. build/torch213-cxx11-cu130-x86_64-linux/_fp4_gemm_cuda_752e924.abi3.so +3 -0
  20. build/torch213-cxx11-cu130-x86_64-linux/_ops.py +9 -0
  21. build/torch213-cxx11-cu130-x86_64-linux/fp4_gemm/__init__.py +26 -0
  22. build/torch213-cxx11-cu130-x86_64-linux/metadata.json +33 -0
  23. build/torch213-cxx11-cu132-x86_64-linux/__init__.py +122 -0
  24. build/torch213-cxx11-cu132-x86_64-linux/_fp4_gemm_cuda_752e924.abi3.so +3 -0
  25. build/torch213-cxx11-cu132-x86_64-linux/_ops.py +9 -0
  26. build/torch213-cxx11-cu132-x86_64-linux/fp4_gemm/__init__.py +26 -0
  27. build/torch213-cxx11-cu132-x86_64-linux/metadata.json +33 -0
benchmarks/benchmark.py CHANGED
@@ -25,8 +25,10 @@ class BenchResult:
25
  K: int
26
  variant: int
27
  flashrt_us: float
28
- torch_reference_us: float
29
- speedup_vs_reference: float
 
 
30
  max_abs: float
31
  mean_abs: float
32
  p99_abs: float
@@ -70,15 +72,17 @@ def bench_case(helpers, ops, name: str, shape: tuple[int, int, int], warmup: int
70
  def torch_ref():
71
  return (a_deq.float() @ b_deq.float().T).to(torch.bfloat16)
72
 
73
- torch_us = measure(torch_ref, warmup, iters)
 
 
74
  results: list[BenchResult] = []
75
  for variant in (0, 1, 2):
76
  out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
77
- ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant)
78
  torch.cuda.synchronize()
79
  max_abs, mean_abs, p99_abs, cosine = helpers.metrics(out, expected)
80
  flashrt_us = measure(
81
- lambda: ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant),
82
  warmup,
83
  iters,
84
  )
@@ -90,8 +94,10 @@ def bench_case(helpers, ops, name: str, shape: tuple[int, int, int], warmup: int
90
  K=k,
91
  variant=variant,
92
  flashrt_us=flashrt_us,
93
- torch_reference_us=torch_us,
94
- speedup_vs_reference=torch_us / flashrt_us,
 
 
95
  max_abs=max_abs,
96
  mean_abs=mean_abs,
97
  p99_abs=p99_abs,
@@ -104,6 +110,8 @@ def bench_case(helpers, ops, name: str, shape: tuple[int, int, int], warmup: int
104
 
105
  def main() -> int:
106
  parser = argparse.ArgumentParser()
 
 
107
  parser.add_argument("--mode", choices=["smoke", "headline"], default="headline")
108
  parser.add_argument("--warmup", type=int, default=20)
109
  parser.add_argument("--iterations", type=int, default=100)
@@ -111,7 +119,11 @@ def main() -> int:
111
  args = parser.parse_args()
112
 
113
  helpers = load_helpers()
114
- ops = helpers.load_source_ops()
 
 
 
 
115
  shapes = {
116
  "small_m16_n128_k128": (16, 128, 128),
117
  "small_m32_n256_k256": (32, 256, 256),
@@ -124,6 +136,7 @@ def main() -> int:
124
  results.extend(bench_case(helpers, ops, name, shape, args.warmup, args.iterations))
125
  payload = {
126
  "mode": args.mode,
 
127
  "device": torch.cuda.get_device_name(),
128
  "torch": torch.__version__,
129
  "results": [asdict(item) for item in results],
 
25
  K: int
26
  variant: int
27
  flashrt_us: float
28
+ torch_eager_us: float
29
+ torch_compile_us: float
30
+ speedup_vs_eager: float
31
+ speedup_vs_compile: float
32
  max_abs: float
33
  mean_abs: float
34
  p99_abs: float
 
72
  def torch_ref():
73
  return (a_deq.float() @ b_deq.float().T).to(torch.bfloat16)
74
 
75
+ torch_eager_us = measure(torch_ref, warmup, iters)
76
+ compiled_ref = torch.compile(torch_ref, mode="max-autotune-no-cudagraphs")
77
+ torch_compile_us = measure(compiled_ref, warmup, iters)
78
  results: list[BenchResult] = []
79
  for variant in (0, 1, 2):
80
  out = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
81
+ ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant)
82
  torch.cuda.synchronize()
83
  max_abs, mean_abs, p99_abs, cosine = helpers.metrics(out, expected)
84
  flashrt_us = measure(
85
+ lambda: ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, 1.0, variant),
86
  warmup,
87
  iters,
88
  )
 
94
  K=k,
95
  variant=variant,
96
  flashrt_us=flashrt_us,
97
+ torch_eager_us=torch_eager_us,
98
+ torch_compile_us=torch_compile_us,
99
+ speedup_vs_eager=torch_eager_us / flashrt_us,
100
+ speedup_vs_compile=torch_compile_us / flashrt_us,
101
  max_abs=max_abs,
102
  mean_abs=mean_abs,
103
  p99_abs=p99_abs,
 
110
 
111
  def main() -> int:
112
  parser = argparse.ArgumentParser()
113
+ parser.add_argument("--backend", choices=["source", "installed"], default="source")
114
+ parser.add_argument("--artifact", default=None)
115
  parser.add_argument("--mode", choices=["smoke", "headline"], default="headline")
116
  parser.add_argument("--warmup", type=int, default=20)
117
  parser.add_argument("--iterations", type=int, default=100)
 
119
  args = parser.parse_args()
120
 
121
  helpers = load_helpers()
122
+ ops = (
123
+ helpers.load_source_ops()
124
+ if args.backend == "source"
125
+ else helpers.load_installed_ops(args.artifact)
126
+ )
127
  shapes = {
128
  "small_m16_n128_k128": (16, 128, 128),
129
  "small_m32_n256_k256": (32, 256, 256),
 
136
  results.extend(bench_case(helpers, ops, name, shape, args.warmup, args.iterations))
137
  payload = {
138
  "mode": args.mode,
139
+ "backend": args.backend,
140
  "device": torch.cuda.get_device_name(),
141
  "torch": torch.__version__,
142
  "results": [asdict(item) for item in results],
build/torch211-cxx11-cu128-x86_64-linux/__init__.py CHANGED
@@ -23,7 +23,7 @@ def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
23
  )
24
 
25
 
26
- @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
27
  def _linear_fake(
28
  a_packed: torch.Tensor,
29
  b_packed: torch.Tensor,
@@ -36,6 +36,19 @@ def _linear_fake(
36
  return None
37
 
38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
40
  def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
41
  return None
@@ -70,7 +83,7 @@ def dequantize_fp4_sfa_fp16(
70
  return out
71
 
72
 
73
- def fp4_w4a16_linear_bf16(
74
  a_packed: torch.Tensor,
75
  b_packed: torch.Tensor,
76
  sfa: torch.Tensor,
@@ -81,6 +94,29 @@ def fp4_w4a16_linear_bf16(
81
  ) -> torch.Tensor:
82
  if out is None:
83
  out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
84
- ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
85
  return out
86
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
23
  )
24
 
25
 
26
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
27
  def _linear_fake(
28
  a_packed: torch.Tensor,
29
  b_packed: torch.Tensor,
 
36
  return None
37
 
38
 
39
+ @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
40
+ def _legacy_linear_fake(
41
+ a_packed: torch.Tensor,
42
+ b_packed: torch.Tensor,
43
+ sfa: torch.Tensor,
44
+ sfb: torch.Tensor,
45
+ out: torch.Tensor,
46
+ alpha: float = 1.0,
47
+ variant: int = 0,
48
+ ) -> None:
49
+ return None
50
+
51
+
52
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
53
  def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
54
  return None
 
83
  return out
84
 
85
 
86
+ def nvfp4_gemm_bf16(
87
  a_packed: torch.Tensor,
88
  b_packed: torch.Tensor,
89
  sfa: torch.Tensor,
 
94
  ) -> torch.Tensor:
95
  if out is None:
96
  out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
97
+ ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
98
  return out
99
 
100
+
101
+ def fp4_w4a16_linear_bf16(
102
+ a_packed: torch.Tensor,
103
+ b_packed: torch.Tensor,
104
+ sfa: torch.Tensor,
105
+ sfb: torch.Tensor,
106
+ alpha: float = 1.0,
107
+ out: torch.Tensor | None = None,
108
+ variant: int = 0,
109
+ ) -> torch.Tensor:
110
+ """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
111
+ return nvfp4_gemm_bf16(
112
+ a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
113
+ )
114
+
115
+
116
+ __all__ = [
117
+ "dequantize_fp4_sfa_fp16",
118
+ "fp4_w4a16_linear_bf16",
119
+ "nvfp4_gemm_bf16",
120
+ "quantize_fp4_sfa_fp16",
121
+ "sfa_size_bytes",
122
+ ]
build/torch211-cxx11-cu128-x86_64-linux/{_fp4_gemm_cuda_7e2e5b9.abi3.so → _fp4_gemm_cuda_752e924.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:d031c5d46ab5e14b13f68565ad53dff18a7f9ee76296b8158f626b11b50f61ef
3
- size 671608
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:2b2591080622003318c532d593c7f0738803d71d75953ee7331c3fab57e48e3a
3
+ size 671784
build/torch211-cxx11-cu128-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_7e2e5b9
3
- ops = torch.ops._fp4_gemm_cuda_7e2e5b9
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_7e2e5b9::{op_name}"
 
1
  import torch
2
+ from . import _fp4_gemm_cuda_752e924
3
+ ops = torch.ops._fp4_gemm_cuda_752e924
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp4_gemm_cuda_752e924::{op_name}"
build/torch211-cxx11-cu128-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp4-gemm",
3
- "id": "_fp4_gemm_cuda_7e2e5b9",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
@@ -13,10 +13,21 @@
13
  "digest": {
14
  "algorithm": "sha256",
15
  "files": {
16
- "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
17
- "_fp4_gemm_cuda_7e2e5b9.abi3.so": "0DHF1Gq14UsT9oVlrVPf8Yp/nudilrgVj2JrEbUPYe8=",
18
- "_ops.py": "UnZlP4anh/kfneg4T/V+ueYOufY4kzFv6zxKMKlaRuA=",
19
  "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
20
  }
 
 
 
 
 
 
 
 
 
 
 
21
  }
22
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_752e924",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
13
  "digest": {
14
  "algorithm": "sha256",
15
  "files": {
16
+ "__init__.py": "12G0sUObMOVtJ0cD3xF6s+xKU/n1ciO3m7nh1TZkS7w=",
17
+ "_fp4_gemm_cuda_752e924.abi3.so": "KyWRCAYiADMYxTLVk8fwc4gD1x11lT7nMxw/q1fkjjo=",
18
+ "_ops.py": "u0gVG4RRqNtjB3WknduXgbVTAnrWkWVGrWIfvCSaaXA=",
19
  "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
20
  }
21
+ },
22
+ "provenance": {
23
+ "kernel-builder": {
24
+ "version": "0.17.0-dev0",
25
+ "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
26
+ "dirty": false
27
+ },
28
+ "kernel": {
29
+ "sha": "752e9241351caea5ae5c13bf8b51f1daab967bee",
30
+ "dirty": false
31
+ }
32
  }
33
  }
build/torch211-cxx11-cu130-x86_64-linux/__init__.py CHANGED
@@ -23,7 +23,7 @@ def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
23
  )
24
 
25
 
26
- @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
27
  def _linear_fake(
28
  a_packed: torch.Tensor,
29
  b_packed: torch.Tensor,
@@ -36,6 +36,19 @@ def _linear_fake(
36
  return None
37
 
38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
40
  def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
41
  return None
@@ -70,7 +83,7 @@ def dequantize_fp4_sfa_fp16(
70
  return out
71
 
72
 
73
- def fp4_w4a16_linear_bf16(
74
  a_packed: torch.Tensor,
75
  b_packed: torch.Tensor,
76
  sfa: torch.Tensor,
@@ -81,6 +94,29 @@ def fp4_w4a16_linear_bf16(
81
  ) -> torch.Tensor:
82
  if out is None:
83
  out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
84
- ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
85
  return out
86
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
23
  )
24
 
25
 
26
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
27
  def _linear_fake(
28
  a_packed: torch.Tensor,
29
  b_packed: torch.Tensor,
 
36
  return None
37
 
38
 
39
+ @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
40
+ def _legacy_linear_fake(
41
+ a_packed: torch.Tensor,
42
+ b_packed: torch.Tensor,
43
+ sfa: torch.Tensor,
44
+ sfb: torch.Tensor,
45
+ out: torch.Tensor,
46
+ alpha: float = 1.0,
47
+ variant: int = 0,
48
+ ) -> None:
49
+ return None
50
+
51
+
52
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
53
  def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
54
  return None
 
83
  return out
84
 
85
 
86
+ def nvfp4_gemm_bf16(
87
  a_packed: torch.Tensor,
88
  b_packed: torch.Tensor,
89
  sfa: torch.Tensor,
 
94
  ) -> torch.Tensor:
95
  if out is None:
96
  out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
97
+ ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
98
  return out
99
 
100
+
101
+ def fp4_w4a16_linear_bf16(
102
+ a_packed: torch.Tensor,
103
+ b_packed: torch.Tensor,
104
+ sfa: torch.Tensor,
105
+ sfb: torch.Tensor,
106
+ alpha: float = 1.0,
107
+ out: torch.Tensor | None = None,
108
+ variant: int = 0,
109
+ ) -> torch.Tensor:
110
+ """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
111
+ return nvfp4_gemm_bf16(
112
+ a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
113
+ )
114
+
115
+
116
+ __all__ = [
117
+ "dequantize_fp4_sfa_fp16",
118
+ "fp4_w4a16_linear_bf16",
119
+ "nvfp4_gemm_bf16",
120
+ "quantize_fp4_sfa_fp16",
121
+ "sfa_size_bytes",
122
+ ]
build/torch211-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_7e2e5b9.abi3.so → _fp4_gemm_cuda_752e924.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:e211cafb6cf95d3ba32015492d7de6bdb29cb6df0cff26400146ed1099ce1cb3
3
- size 714176
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:97398c771ed76118703b7e99ed1f757deb2ea0e65512c6548455008b74dd955b
3
+ size 714360
build/torch211-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_7e2e5b9
3
- ops = torch.ops._fp4_gemm_cuda_7e2e5b9
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_7e2e5b9::{op_name}"
 
1
  import torch
2
+ from . import _fp4_gemm_cuda_752e924
3
+ ops = torch.ops._fp4_gemm_cuda_752e924
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp4_gemm_cuda_752e924::{op_name}"
build/torch211-cxx11-cu130-x86_64-linux/metadata.json CHANGED
@@ -1,6 +1,6 @@
1
  {
2
  "name": "fp4-gemm",
3
- "id": "_fp4_gemm_cuda_7e2e5b9",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
@@ -13,10 +13,21 @@
13
  "digest": {
14
  "algorithm": "sha256",
15
  "files": {
16
- "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
17
- "_fp4_gemm_cuda_7e2e5b9.abi3.so": "4hHK+2z5XTujIBVJLX3mvbKctt8M/yZAAUbtEJnOHLM=",
18
- "_ops.py": "UnZlP4anh/kfneg4T/V+ueYOufY4kzFv6zxKMKlaRuA=",
19
  "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
20
  }
 
 
 
 
 
 
 
 
 
 
 
21
  }
22
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_752e924",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
13
  "digest": {
14
  "algorithm": "sha256",
15
  "files": {
16
+ "__init__.py": "12G0sUObMOVtJ0cD3xF6s+xKU/n1ciO3m7nh1TZkS7w=",
17
+ "_fp4_gemm_cuda_752e924.abi3.so": "lzmMdx7XYRhwO36Z7R91fesuoOZVEsZUhFUAi3TdlVs=",
18
+ "_ops.py": "u0gVG4RRqNtjB3WknduXgbVTAnrWkWVGrWIfvCSaaXA=",
19
  "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
20
  }
21
+ },
22
+ "provenance": {
23
+ "kernel-builder": {
24
+ "version": "0.17.0-dev0",
25
+ "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
26
+ "dirty": false
27
+ },
28
+ "kernel": {
29
+ "sha": "752e9241351caea5ae5c13bf8b51f1daab967bee",
30
+ "dirty": false
31
+ }
32
  }
33
  }
build/torch212-cxx11-cu130-x86_64-linux/__init__.py CHANGED
@@ -23,7 +23,7 @@ def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
23
  )
24
 
25
 
26
- @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
27
  def _linear_fake(
28
  a_packed: torch.Tensor,
29
  b_packed: torch.Tensor,
@@ -36,6 +36,19 @@ def _linear_fake(
36
  return None
37
 
38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
40
  def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
41
  return None
@@ -70,7 +83,7 @@ def dequantize_fp4_sfa_fp16(
70
  return out
71
 
72
 
73
- def fp4_w4a16_linear_bf16(
74
  a_packed: torch.Tensor,
75
  b_packed: torch.Tensor,
76
  sfa: torch.Tensor,
@@ -81,6 +94,29 @@ def fp4_w4a16_linear_bf16(
81
  ) -> torch.Tensor:
82
  if out is None:
83
  out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
84
- ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
85
  return out
86
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
23
  )
24
 
25
 
26
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
27
  def _linear_fake(
28
  a_packed: torch.Tensor,
29
  b_packed: torch.Tensor,
 
36
  return None
37
 
38
 
39
+ @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
40
+ def _legacy_linear_fake(
41
+ a_packed: torch.Tensor,
42
+ b_packed: torch.Tensor,
43
+ sfa: torch.Tensor,
44
+ sfb: torch.Tensor,
45
+ out: torch.Tensor,
46
+ alpha: float = 1.0,
47
+ variant: int = 0,
48
+ ) -> None:
49
+ return None
50
+
51
+
52
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
53
  def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
54
  return None
 
83
  return out
84
 
85
 
86
+ def nvfp4_gemm_bf16(
87
  a_packed: torch.Tensor,
88
  b_packed: torch.Tensor,
89
  sfa: torch.Tensor,
 
94
  ) -> torch.Tensor:
95
  if out is None:
96
  out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
97
+ ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
98
  return out
99
 
100
+
101
+ def fp4_w4a16_linear_bf16(
102
+ a_packed: torch.Tensor,
103
+ b_packed: torch.Tensor,
104
+ sfa: torch.Tensor,
105
+ sfb: torch.Tensor,
106
+ alpha: float = 1.0,
107
+ out: torch.Tensor | None = None,
108
+ variant: int = 0,
109
+ ) -> torch.Tensor:
110
+ """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
111
+ return nvfp4_gemm_bf16(
112
+ a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
113
+ )
114
+
115
+
116
+ __all__ = [
117
+ "dequantize_fp4_sfa_fp16",
118
+ "fp4_w4a16_linear_bf16",
119
+ "nvfp4_gemm_bf16",
120
+ "quantize_fp4_sfa_fp16",
121
+ "sfa_size_bytes",
122
+ ]
build/torch212-cxx11-cu130-x86_64-linux/{_fp4_gemm_cuda_7e2e5b9.abi3.so → _fp4_gemm_cuda_752e924.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:63e8e675f69f6cdb3cf1c0a78fe0032f72638db29e8e4844501a2204cca77e08
3
- size 724584
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:4a5f761dd4108c0aad7cd2ef8875be4fb74ebff498e4b8f456a56c1f54e2689a
3
+ size 724760
build/torch212-cxx11-cu130-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_7e2e5b9
3
- ops = torch.ops._fp4_gemm_cuda_7e2e5b9
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_7e2e5b9::{op_name}"
 
1
  import torch
2
+ from . import _fp4_gemm_cuda_752e924
3
+ ops = torch.ops._fp4_gemm_cuda_752e924
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp4_gemm_cuda_752e924::{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_7e2e5b9",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
@@ -13,10 +13,21 @@
13
  "digest": {
14
  "algorithm": "sha256",
15
  "files": {
16
- "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
17
- "_fp4_gemm_cuda_7e2e5b9.abi3.so": "Y+jmdfafbNs88cCnj+ADL3JjjbKejkhEUBoiBMynfgg=",
18
- "_ops.py": "UnZlP4anh/kfneg4T/V+ueYOufY4kzFv6zxKMKlaRuA=",
19
  "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
20
  }
 
 
 
 
 
 
 
 
 
 
 
21
  }
22
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_752e924",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
13
  "digest": {
14
  "algorithm": "sha256",
15
  "files": {
16
+ "__init__.py": "12G0sUObMOVtJ0cD3xF6s+xKU/n1ciO3m7nh1TZkS7w=",
17
+ "_fp4_gemm_cuda_752e924.abi3.so": "Sl92HdQQjAqtfNLviHW+T7dOv/SY5Lj0VqVsH1TiaJo=",
18
+ "_ops.py": "u0gVG4RRqNtjB3WknduXgbVTAnrWkWVGrWIfvCSaaXA=",
19
  "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
20
  }
21
+ },
22
+ "provenance": {
23
+ "kernel-builder": {
24
+ "version": "0.17.0-dev0",
25
+ "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
26
+ "dirty": false
27
+ },
28
+ "kernel": {
29
+ "sha": "752e9241351caea5ae5c13bf8b51f1daab967bee",
30
+ "dirty": false
31
+ }
32
  }
33
  }
build/torch212-cxx11-cu132-x86_64-linux/__init__.py CHANGED
@@ -23,7 +23,7 @@ def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
23
  )
24
 
25
 
26
- @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
27
  def _linear_fake(
28
  a_packed: torch.Tensor,
29
  b_packed: torch.Tensor,
@@ -36,6 +36,19 @@ def _linear_fake(
36
  return None
37
 
38
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
40
  def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
41
  return None
@@ -70,7 +83,7 @@ def dequantize_fp4_sfa_fp16(
70
  return out
71
 
72
 
73
- def fp4_w4a16_linear_bf16(
74
  a_packed: torch.Tensor,
75
  b_packed: torch.Tensor,
76
  sfa: torch.Tensor,
@@ -81,6 +94,29 @@ def fp4_w4a16_linear_bf16(
81
  ) -> torch.Tensor:
82
  if out is None:
83
  out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
84
- ops.fp4_w4a16_linear_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
85
  return out
86
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
23
  )
24
 
25
 
26
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
27
  def _linear_fake(
28
  a_packed: torch.Tensor,
29
  b_packed: torch.Tensor,
 
36
  return None
37
 
38
 
39
+ @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
40
+ def _legacy_linear_fake(
41
+ a_packed: torch.Tensor,
42
+ b_packed: torch.Tensor,
43
+ sfa: torch.Tensor,
44
+ sfb: torch.Tensor,
45
+ out: torch.Tensor,
46
+ alpha: float = 1.0,
47
+ variant: int = 0,
48
+ ) -> None:
49
+ return None
50
+
51
+
52
  @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
53
  def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
54
  return None
 
83
  return out
84
 
85
 
86
+ def nvfp4_gemm_bf16(
87
  a_packed: torch.Tensor,
88
  b_packed: torch.Tensor,
89
  sfa: torch.Tensor,
 
94
  ) -> torch.Tensor:
95
  if out is None:
96
  out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
97
+ ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
98
  return out
99
 
100
+
101
+ def fp4_w4a16_linear_bf16(
102
+ a_packed: torch.Tensor,
103
+ b_packed: torch.Tensor,
104
+ sfa: torch.Tensor,
105
+ sfb: torch.Tensor,
106
+ alpha: float = 1.0,
107
+ out: torch.Tensor | None = None,
108
+ variant: int = 0,
109
+ ) -> torch.Tensor:
110
+ """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
111
+ return nvfp4_gemm_bf16(
112
+ a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
113
+ )
114
+
115
+
116
+ __all__ = [
117
+ "dequantize_fp4_sfa_fp16",
118
+ "fp4_w4a16_linear_bf16",
119
+ "nvfp4_gemm_bf16",
120
+ "quantize_fp4_sfa_fp16",
121
+ "sfa_size_bytes",
122
+ ]
build/torch212-cxx11-cu132-x86_64-linux/{_fp4_gemm_cuda_7e2e5b9.abi3.so → _fp4_gemm_cuda_752e924.abi3.so} RENAMED
@@ -1,3 +1,3 @@
1
  version https://git-lfs.github.com/spec/v1
2
- oid sha256:fd5e66f4faeddfb45f42d8467508f13b2e065fb5d3c52a8404dc6a765a1b891f
3
- size 724632
 
1
  version https://git-lfs.github.com/spec/v1
2
+ oid sha256:99fe652b5e4b685c15570441145e6aae1175055697a3ffa17430b85ea1ea0b87
3
+ size 724816
build/torch212-cxx11-cu132-x86_64-linux/_ops.py CHANGED
@@ -1,9 +1,9 @@
1
  import torch
2
- from . import _fp4_gemm_cuda_7e2e5b9
3
- ops = torch.ops._fp4_gemm_cuda_7e2e5b9
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
- return f"_fp4_gemm_cuda_7e2e5b9::{op_name}"
 
1
  import torch
2
+ from . import _fp4_gemm_cuda_752e924
3
+ ops = torch.ops._fp4_gemm_cuda_752e924
4
 
5
  def add_op_namespace_prefix(op_name: str):
6
  """
7
  Prefix op by namespace.
8
  """
9
+ return f"_fp4_gemm_cuda_752e924::{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_7e2e5b9",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
@@ -13,10 +13,21 @@
13
  "digest": {
14
  "algorithm": "sha256",
15
  "files": {
16
- "__init__.py": "WSLdLQD92ulRtOhSLoAhWgFYN6oucbbkTg19GAaltBs=",
17
- "_fp4_gemm_cuda_7e2e5b9.abi3.so": "/V5m9Prt37RfQthGdQjxOy4GX7XTxSqEBNxqdlobiR8=",
18
- "_ops.py": "UnZlP4anh/kfneg4T/V+ueYOufY4kzFv6zxKMKlaRuA=",
19
  "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
20
  }
 
 
 
 
 
 
 
 
 
 
 
21
  }
22
  }
 
1
  {
2
  "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_752e924",
4
  "version": 1,
5
  "license": "Apache-2.0",
6
  "python-depends": [],
 
13
  "digest": {
14
  "algorithm": "sha256",
15
  "files": {
16
+ "__init__.py": "12G0sUObMOVtJ0cD3xF6s+xKU/n1ciO3m7nh1TZkS7w=",
17
+ "_fp4_gemm_cuda_752e924.abi3.so": "mf5lK15LaFwVVwRBFF5qrhF1BVaXo/+hdDC4XqHqC4c=",
18
+ "_ops.py": "u0gVG4RRqNtjB3WknduXgbVTAnrWkWVGrWIfvCSaaXA=",
19
  "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
20
  }
21
+ },
22
+ "provenance": {
23
+ "kernel-builder": {
24
+ "version": "0.17.0-dev0",
25
+ "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
26
+ "dirty": false
27
+ },
28
+ "kernel": {
29
+ "sha": "752e9241351caea5ae5c13bf8b51f1daab967bee",
30
+ "dirty": false
31
+ }
32
  }
33
  }
build/torch213-cxx11-cu130-x86_64-linux/__init__.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """FlashRT FP4 GEMM kernels."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import torch
6
+
7
+ from ._ops import add_op_namespace_prefix, ops
8
+
9
+
10
+ def sfa_size_bytes(rows: int, dim: int) -> int:
11
+ if rows <= 0 or dim <= 0 or dim % 16 != 0:
12
+ raise ValueError("rows must be positive and dim must be positive/divisible by 16")
13
+ n_blocks = dim // 16
14
+ n_row_super = (rows + 127) // 128
15
+ n_col_super = (n_blocks + 3) // 4
16
+ return n_row_super * n_col_super * 512
17
+
18
+
19
+ def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
20
+ return (
21
+ torch.empty((rows, dim // 2), device=device, dtype=torch.uint8),
22
+ torch.empty((sfa_size_bytes(rows, dim),), device=device, dtype=torch.uint8),
23
+ )
24
+
25
+
26
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
27
+ def _linear_fake(
28
+ a_packed: torch.Tensor,
29
+ b_packed: torch.Tensor,
30
+ sfa: torch.Tensor,
31
+ sfb: torch.Tensor,
32
+ out: torch.Tensor,
33
+ alpha: float = 1.0,
34
+ variant: int = 0,
35
+ ) -> None:
36
+ return None
37
+
38
+
39
+ @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
40
+ def _legacy_linear_fake(
41
+ a_packed: torch.Tensor,
42
+ b_packed: torch.Tensor,
43
+ sfa: torch.Tensor,
44
+ sfb: torch.Tensor,
45
+ out: torch.Tensor,
46
+ alpha: float = 1.0,
47
+ variant: int = 0,
48
+ ) -> None:
49
+ return None
50
+
51
+
52
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
53
+ def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
54
+ return None
55
+
56
+
57
+ @torch.library.register_fake(add_op_namespace_prefix("dequantize_fp4_sfa_fp16"))
58
+ def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor, is_sfb: bool = False) -> None:
59
+ return None
60
+
61
+
62
+ def quantize_fp4_sfa_fp16(
63
+ x: torch.Tensor,
64
+ packed: torch.Tensor | None = None,
65
+ sfa: torch.Tensor | None = None,
66
+ is_sfb: bool = False,
67
+ ):
68
+ if packed is None or sfa is None:
69
+ packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
70
+ ops.quantize_fp4_sfa_fp16(x, packed, sfa, bool(is_sfb))
71
+ return packed, sfa
72
+
73
+
74
+ def dequantize_fp4_sfa_fp16(
75
+ packed: torch.Tensor,
76
+ sfa: torch.Tensor,
77
+ out: torch.Tensor | None = None,
78
+ is_sfb: bool = False,
79
+ ) -> torch.Tensor:
80
+ if out is None:
81
+ out = torch.empty((packed.shape[0], packed.shape[1] * 2), device=packed.device, dtype=torch.float16)
82
+ ops.dequantize_fp4_sfa_fp16(packed, sfa, out, bool(is_sfb))
83
+ return out
84
+
85
+
86
+ def nvfp4_gemm_bf16(
87
+ a_packed: torch.Tensor,
88
+ b_packed: torch.Tensor,
89
+ sfa: torch.Tensor,
90
+ sfb: torch.Tensor,
91
+ alpha: float = 1.0,
92
+ out: torch.Tensor | None = None,
93
+ variant: int = 0,
94
+ ) -> torch.Tensor:
95
+ if out is None:
96
+ out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
97
+ ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
98
+ return out
99
+
100
+
101
+ def fp4_w4a16_linear_bf16(
102
+ a_packed: torch.Tensor,
103
+ b_packed: torch.Tensor,
104
+ sfa: torch.Tensor,
105
+ sfb: torch.Tensor,
106
+ alpha: float = 1.0,
107
+ out: torch.Tensor | None = None,
108
+ variant: int = 0,
109
+ ) -> torch.Tensor:
110
+ """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
111
+ return nvfp4_gemm_bf16(
112
+ a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
113
+ )
114
+
115
+
116
+ __all__ = [
117
+ "dequantize_fp4_sfa_fp16",
118
+ "fp4_w4a16_linear_bf16",
119
+ "nvfp4_gemm_bf16",
120
+ "quantize_fp4_sfa_fp16",
121
+ "sfa_size_bytes",
122
+ ]
build/torch213-cxx11-cu130-x86_64-linux/_fp4_gemm_cuda_752e924.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:5bc5b957118e7c4d272ea8eff37b23d864fbf0327d579bb7d79ad05c6e6facdf
3
+ size 724608
build/torch213-cxx11-cu130-x86_64-linux/_ops.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from . import _fp4_gemm_cuda_752e924
3
+ ops = torch.ops._fp4_gemm_cuda_752e924
4
+
5
+ def add_op_namespace_prefix(op_name: str):
6
+ """
7
+ Prefix op by namespace.
8
+ """
9
+ return f"_fp4_gemm_cuda_752e924::{op_name}"
build/torch213-cxx11-cu130-x86_64-linux/fp4_gemm/__init__.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import ctypes
2
+ import importlib.util
3
+ import sys
4
+ from pathlib import Path
5
+ from types import ModuleType
6
+
7
+
8
+ def _import_from_path(file_path: Path) -> ModuleType:
9
+ # We cannot use the module name as-is, after adding it to `sys.modules`,
10
+ # it would also be used for other imports. So, we make a module name that
11
+ # depends on the path for it to be unique using the hex-encoded hash of
12
+ # the path.
13
+ path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
14
+ module_name = path_hash
15
+ spec = importlib.util.spec_from_file_location(module_name, file_path)
16
+ if spec is None:
17
+ raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
18
+ module = importlib.util.module_from_spec(spec)
19
+ if module is None:
20
+ raise ImportError(f"Cannot load module {module_name} from spec")
21
+ sys.modules[module_name] = module
22
+ spec.loader.exec_module(module) # type: ignore
23
+ return module
24
+
25
+
26
+ globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
build/torch213-cxx11-cu130-x86_64-linux/metadata.json ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_752e924",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "python-depends": [],
7
+ "backend": {
8
+ "type": "cuda",
9
+ "archs": [
10
+ "12.0a"
11
+ ]
12
+ },
13
+ "digest": {
14
+ "algorithm": "sha256",
15
+ "files": {
16
+ "__init__.py": "12G0sUObMOVtJ0cD3xF6s+xKU/n1ciO3m7nh1TZkS7w=",
17
+ "_fp4_gemm_cuda_752e924.abi3.so": "W8W5VxGOfE0nLqjv83sj2GT78DJ9V5u315rQXG5vrN8=",
18
+ "_ops.py": "u0gVG4RRqNtjB3WknduXgbVTAnrWkWVGrWIfvCSaaXA=",
19
+ "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
20
+ }
21
+ },
22
+ "provenance": {
23
+ "kernel-builder": {
24
+ "version": "0.17.0-dev0",
25
+ "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
26
+ "dirty": false
27
+ },
28
+ "kernel": {
29
+ "sha": "752e9241351caea5ae5c13bf8b51f1daab967bee",
30
+ "dirty": false
31
+ }
32
+ }
33
+ }
build/torch213-cxx11-cu132-x86_64-linux/__init__.py ADDED
@@ -0,0 +1,122 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """FlashRT FP4 GEMM kernels."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import torch
6
+
7
+ from ._ops import add_op_namespace_prefix, ops
8
+
9
+
10
+ def sfa_size_bytes(rows: int, dim: int) -> int:
11
+ if rows <= 0 or dim <= 0 or dim % 16 != 0:
12
+ raise ValueError("rows must be positive and dim must be positive/divisible by 16")
13
+ n_blocks = dim // 16
14
+ n_row_super = (rows + 127) // 128
15
+ n_col_super = (n_blocks + 3) // 4
16
+ return n_row_super * n_col_super * 512
17
+
18
+
19
+ def _alloc_fp4(rows: int, dim: int, device: torch.device | str):
20
+ return (
21
+ torch.empty((rows, dim // 2), device=device, dtype=torch.uint8),
22
+ torch.empty((sfa_size_bytes(rows, dim),), device=device, dtype=torch.uint8),
23
+ )
24
+
25
+
26
+ @torch.library.register_fake(add_op_namespace_prefix("nvfp4_gemm_bf16"))
27
+ def _linear_fake(
28
+ a_packed: torch.Tensor,
29
+ b_packed: torch.Tensor,
30
+ sfa: torch.Tensor,
31
+ sfb: torch.Tensor,
32
+ out: torch.Tensor,
33
+ alpha: float = 1.0,
34
+ variant: int = 0,
35
+ ) -> None:
36
+ return None
37
+
38
+
39
+ @torch.library.register_fake(add_op_namespace_prefix("fp4_w4a16_linear_bf16"))
40
+ def _legacy_linear_fake(
41
+ a_packed: torch.Tensor,
42
+ b_packed: torch.Tensor,
43
+ sfa: torch.Tensor,
44
+ sfb: torch.Tensor,
45
+ out: torch.Tensor,
46
+ alpha: float = 1.0,
47
+ variant: int = 0,
48
+ ) -> None:
49
+ return None
50
+
51
+
52
+ @torch.library.register_fake(add_op_namespace_prefix("quantize_fp4_sfa_fp16"))
53
+ def _quant_fake(x: torch.Tensor, packed: torch.Tensor, sfa: torch.Tensor, is_sfb: bool = False) -> None:
54
+ return None
55
+
56
+
57
+ @torch.library.register_fake(add_op_namespace_prefix("dequantize_fp4_sfa_fp16"))
58
+ def _dequant_fake(packed: torch.Tensor, sfa: torch.Tensor, out: torch.Tensor, is_sfb: bool = False) -> None:
59
+ return None
60
+
61
+
62
+ def quantize_fp4_sfa_fp16(
63
+ x: torch.Tensor,
64
+ packed: torch.Tensor | None = None,
65
+ sfa: torch.Tensor | None = None,
66
+ is_sfb: bool = False,
67
+ ):
68
+ if packed is None or sfa is None:
69
+ packed, sfa = _alloc_fp4(x.shape[0], x.shape[1], x.device)
70
+ ops.quantize_fp4_sfa_fp16(x, packed, sfa, bool(is_sfb))
71
+ return packed, sfa
72
+
73
+
74
+ def dequantize_fp4_sfa_fp16(
75
+ packed: torch.Tensor,
76
+ sfa: torch.Tensor,
77
+ out: torch.Tensor | None = None,
78
+ is_sfb: bool = False,
79
+ ) -> torch.Tensor:
80
+ if out is None:
81
+ out = torch.empty((packed.shape[0], packed.shape[1] * 2), device=packed.device, dtype=torch.float16)
82
+ ops.dequantize_fp4_sfa_fp16(packed, sfa, out, bool(is_sfb))
83
+ return out
84
+
85
+
86
+ def nvfp4_gemm_bf16(
87
+ a_packed: torch.Tensor,
88
+ b_packed: torch.Tensor,
89
+ sfa: torch.Tensor,
90
+ sfb: torch.Tensor,
91
+ alpha: float = 1.0,
92
+ out: torch.Tensor | None = None,
93
+ variant: int = 0,
94
+ ) -> torch.Tensor:
95
+ if out is None:
96
+ out = torch.empty((a_packed.shape[0], b_packed.shape[0]), device=a_packed.device, dtype=torch.bfloat16)
97
+ ops.nvfp4_gemm_bf16(a_packed, b_packed, sfa, sfb, out, float(alpha), int(variant))
98
+ return out
99
+
100
+
101
+ def fp4_w4a16_linear_bf16(
102
+ a_packed: torch.Tensor,
103
+ b_packed: torch.Tensor,
104
+ sfa: torch.Tensor,
105
+ sfb: torch.Tensor,
106
+ alpha: float = 1.0,
107
+ out: torch.Tensor | None = None,
108
+ variant: int = 0,
109
+ ) -> torch.Tensor:
110
+ """Compatibility alias for :func:`nvfp4_gemm_bf16`."""
111
+ return nvfp4_gemm_bf16(
112
+ a_packed, b_packed, sfa, sfb, alpha=alpha, out=out, variant=variant
113
+ )
114
+
115
+
116
+ __all__ = [
117
+ "dequantize_fp4_sfa_fp16",
118
+ "fp4_w4a16_linear_bf16",
119
+ "nvfp4_gemm_bf16",
120
+ "quantize_fp4_sfa_fp16",
121
+ "sfa_size_bytes",
122
+ ]
build/torch213-cxx11-cu132-x86_64-linux/_fp4_gemm_cuda_752e924.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:660955d1fb41a89b8822862c7b20caba35f289f6e068f9e83eab13810f82b49b
3
+ size 724656
build/torch213-cxx11-cu132-x86_64-linux/_ops.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from . import _fp4_gemm_cuda_752e924
3
+ ops = torch.ops._fp4_gemm_cuda_752e924
4
+
5
+ def add_op_namespace_prefix(op_name: str):
6
+ """
7
+ Prefix op by namespace.
8
+ """
9
+ return f"_fp4_gemm_cuda_752e924::{op_name}"
build/torch213-cxx11-cu132-x86_64-linux/fp4_gemm/__init__.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import ctypes
2
+ import importlib.util
3
+ import sys
4
+ from pathlib import Path
5
+ from types import ModuleType
6
+
7
+
8
+ def _import_from_path(file_path: Path) -> ModuleType:
9
+ # We cannot use the module name as-is, after adding it to `sys.modules`,
10
+ # it would also be used for other imports. So, we make a module name that
11
+ # depends on the path for it to be unique using the hex-encoded hash of
12
+ # the path.
13
+ path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
14
+ module_name = path_hash
15
+ spec = importlib.util.spec_from_file_location(module_name, file_path)
16
+ if spec is None:
17
+ raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
18
+ module = importlib.util.module_from_spec(spec)
19
+ if module is None:
20
+ raise ImportError(f"Cannot load module {module_name} from spec")
21
+ sys.modules[module_name] = module
22
+ spec.loader.exec_module(module) # type: ignore
23
+ return module
24
+
25
+
26
+ globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
build/torch213-cxx11-cu132-x86_64-linux/metadata.json ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "fp4-gemm",
3
+ "id": "_fp4_gemm_cuda_752e924",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "python-depends": [],
7
+ "backend": {
8
+ "type": "cuda",
9
+ "archs": [
10
+ "12.0a"
11
+ ]
12
+ },
13
+ "digest": {
14
+ "algorithm": "sha256",
15
+ "files": {
16
+ "__init__.py": "12G0sUObMOVtJ0cD3xF6s+xKU/n1ciO3m7nh1TZkS7w=",
17
+ "_fp4_gemm_cuda_752e924.abi3.so": "ZglV0ftBqJuIIoYseyDKujXyifbgaPnoPqsTgQ+CtJs=",
18
+ "_ops.py": "u0gVG4RRqNtjB3WknduXgbVTAnrWkWVGrWIfvCSaaXA=",
19
+ "fp4_gemm/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY="
20
+ }
21
+ },
22
+ "provenance": {
23
+ "kernel-builder": {
24
+ "version": "0.17.0-dev0",
25
+ "sha": "19aaa6421e674e9fecc352bbae6eab81d19a6bf4",
26
+ "dirty": false
27
+ },
28
+ "kernel": {
29
+ "sha": "752e9241351caea5ae5c13bf8b51f1daab967bee",
30
+ "dirty": false
31
+ }
32
+ }
33
+ }