Hugging Face's logo

Kernels:
Atlas-Inference
/
nvfp4-dense-gemm

Kernel card Files
xet
 Community
Azeez commited on
Commit
4c5264e
·
1 Parent(s): b43edbb

Add transformers NVFP4Kernel entry points (PackedWeight, pack, gemm, swizzled_sf_shape) for v1

Browse files
Files changed (12) hide show
  1. build/torch210-cxx11-cu130-aarch64-linux/__init__.py +86 -16
  2. build/torch210-cxx11-cu130-aarch64-linux/metadata.json +1 -1
  3. build/torch210-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py +127 -26
  4. build/torch211-cxx11-cu130-aarch64-linux/__init__.py +86 -16
  5. build/torch211-cxx11-cu130-aarch64-linux/metadata.json +1 -1
  6. build/torch211-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py +127 -26
  7. build/torch212-cxx11-cu130-aarch64-linux/__init__.py +86 -16
  8. build/torch212-cxx11-cu130-aarch64-linux/metadata.json +1 -1
  9. build/torch212-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py +127 -26
  10. build/torch212-cxx11-cu132-aarch64-linux/__init__.py +86 -16
  11. build/torch212-cxx11-cu132-aarch64-linux/metadata.json +1 -1
  12. build/torch212-cxx11-cu132-aarch64-linux/nvfp4_dense_gemm/__init__.py +127 -26
build/torch210-cxx11-cu130-aarch64-linux/__init__.py CHANGED
@@ -1,26 +1,83 @@
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
- """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121).
3
-
4
- Powers the attention Q/K/V/O projections and the dense FFN
5
- (intermediate_size=17408, hidden=5120). Two precisions are supported:
6
-
7
- - W4A16 (NVFP4 weight, BF16 activation): direct path
8
- - FP8 W/A: via ``predequant_nvfp4_to_fp8`` + ``fp8_gemm_t`` for
9
- KV-cache scenarios where the weight is loaded as NVFP4 but the
10
- GEMM runs in FP8 against an FP8-packed activation.
11
  """
12
 
 
13
  import torch
14
 
15
  from ._ops import ops
16
 
17
- __all__ = [
18
- "w4a16_gemm",
19
- "w4a16_gemm_t",
20
- "predequant_nvfp4_to_fp8",
21
- "fp8_gemm_t",
22
- "bf16_to_fp8",
23
- ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
 
25
 
26
  def w4a16_gemm(
@@ -55,3 +112,16 @@ def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
55
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
56
  """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
57
  ops.bf16_to_fp8(src, dst)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121),
3
+ with Hugging Face ``transformers`` ``NVFP4Config`` entry points
4
+ (``PackedWeight``, ``pack``, ``gemm``, ``swizzled_sf_shape``).
 
 
 
 
 
 
5
  """
6
 
7
+ from dataclasses import dataclass
8
  import torch
9
 
10
  from ._ops import ops
11
 
12
+ FLOAT8_E4M3_MAX = 448.0
13
+ FLOAT4_E2M1_MAX = 6.0
14
+ SF_VEC_SIZE = 16
15
+
16
+
17
+ def swizzled_sf_shape(rows: int, cols: int) -> tuple[int, int]:
18
+ return (rows + 127) // 128 * 128, (cols // SF_VEC_SIZE + 3) // 4
19
+
20
+
21
+ @dataclass
22
+ class PackedWeight:
23
+ qweight: torch.Tensor
24
+ sf: torch.Tensor
25
+ global_scale: torch.Tensor
26
+ n: int
27
+ k: int
28
+ swiglu: bool = False
29
+ sf_rowmajor: torch.Tensor | None = None
30
+
31
+
32
+ def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
33
+ n, k = wd.shape
34
+ g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
35
+ block_amax = g.abs().amax(dim=-1)
36
+ raw_sf = (block_amax / (FLOAT4_E2M1_MAX * max(scale2, 1e-12))).clamp(max=FLOAT8_E4M3_MAX)
37
+ sf_fp8 = raw_sf.to(torch.float8_e4m3fn)
38
+ sf_f32 = sf_fp8.float() * max(scale2, 1e-12)
39
+ scaled = torch.where(sf_f32.unsqueeze(-1) > 0, g / sf_f32.unsqueeze(-1), torch.zeros_like(g))
40
+ ax = scaled.abs()
41
+ nib = torch.zeros_like(ax, dtype=torch.uint8)
42
+ nib = torch.where(ax >= 0.25, torch.tensor(1, device=wd.device, dtype=torch.uint8), nib)
43
+ nib = torch.where(ax >= 0.75, torch.tensor(2, device=wd.device, dtype=torch.uint8), nib)
44
+ nib = torch.where(ax >= 1.25, torch.tensor(3, device=wd.device, dtype=torch.uint8), nib)
45
+ nib = torch.where(ax >= 1.75, torch.tensor(4, device=wd.device, dtype=torch.uint8), nib)
46
+ nib = torch.where(ax >= 2.50, torch.tensor(5, device=wd.device, dtype=torch.uint8), nib)
47
+ nib = torch.where(ax >= 3.50, torch.tensor(6, device=wd.device, dtype=torch.uint8), nib)
48
+ nib = torch.where(ax >= 5.00, torch.tensor(7, device=wd.device, dtype=torch.uint8), nib)
49
+ sign = torch.where(scaled < 0, torch.tensor(8, device=wd.device, dtype=torch.uint8), torch.tensor(0, device=wd.device, dtype=torch.uint8))
50
+ codes = (nib | sign).reshape(n, k // 2, 2)
51
+ qweight = (codes[:, :, 0] | (codes[:, :, 1] << 4)).contiguous()
52
+ return qweight, sf_fp8.view(torch.uint8).contiguous()
53
+
54
+
55
+ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = False) -> PackedWeight:
56
+ wd = w.detach().to(device, torch.bfloat16).contiguous()
57
+ n, k = wd.shape
58
+ amax = wd.abs().amax().to(torch.float32).clamp(min=1e-12)
59
+ scale2 = (amax / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX)).item()
60
+ gs = 1.0 / max(scale2, 1e-12)
61
+ qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
62
+ sf_m, sf_n = swizzled_sf_shape(n, k)
63
+ sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
64
+ return PackedWeight(
65
+ qweight=qweight,
66
+ sf=sf,
67
+ global_scale=torch.tensor([gs], device=device, dtype=torch.float32),
68
+ n=n,
69
+ k=k,
70
+ swiglu=swiglu,
71
+ sf_rowmajor=sf_rowmajor,
72
+ )
73
+
74
+
75
+ def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
76
+ x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
77
+ out = torch.empty((x2.size(0), pw.n), device=x.device, dtype=torch.bfloat16)
78
+ scale2 = 1.0 / pw.global_scale.item()
79
+ ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
80
+ return out.reshape(*x.shape[:-1], pw.n)
81
 
82
 
83
  def w4a16_gemm(
 
112
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
113
  """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
114
  ops.bf16_to_fp8(src, dst)
115
+
116
+
117
+ __all__ = [
118
+ "PackedWeight",
119
+ "pack",
120
+ "gemm",
121
+ "swizzled_sf_shape",
122
+ "w4a16_gemm",
123
+ "w4a16_gemm_t",
124
+ "predequant_nvfp4_to_fp8",
125
+ "fp8_gemm_t",
126
+ "bf16_to_fp8",
127
+ ]
build/torch210-cxx11-cu130-aarch64-linux/metadata.json CHANGED
@@ -1,7 +1,7 @@
1
  {
2
  "name": "nvfp4-dense-gemm",
3
  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
4
- "version": 0,
5
  "license": "AGPL-3.0-only",
6
  "python-depends": [],
7
  "backend": {
 
1
  {
2
  "name": "nvfp4-dense-gemm",
3
  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
4
+ "version": 1,
5
  "license": "AGPL-3.0-only",
6
  "python-depends": [],
7
  "backend": {
build/torch210-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py CHANGED
@@ -1,26 +1,127 @@
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")))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121),
3
+ with Hugging Face ``transformers`` ``NVFP4Config`` entry points
4
+ (``PackedWeight``, ``pack``, ``gemm``, ``swizzled_sf_shape``).
5
+ """
6
+
7
+ from dataclasses import dataclass
8
+ import torch
9
+
10
+ from ._ops import ops
11
+
12
+ FLOAT8_E4M3_MAX = 448.0
13
+ FLOAT4_E2M1_MAX = 6.0
14
+ SF_VEC_SIZE = 16
15
+
16
+
17
+ def swizzled_sf_shape(rows: int, cols: int) -> tuple[int, int]:
18
+ return (rows + 127) // 128 * 128, (cols // SF_VEC_SIZE + 3) // 4
19
+
20
+
21
+ @dataclass
22
+ class PackedWeight:
23
+ qweight: torch.Tensor
24
+ sf: torch.Tensor
25
+ global_scale: torch.Tensor
26
+ n: int
27
+ k: int
28
+ swiglu: bool = False
29
+ sf_rowmajor: torch.Tensor | None = None
30
+
31
+
32
+ def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
33
+ n, k = wd.shape
34
+ g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
35
+ block_amax = g.abs().amax(dim=-1)
36
+ raw_sf = (block_amax / (FLOAT4_E2M1_MAX * max(scale2, 1e-12))).clamp(max=FLOAT8_E4M3_MAX)
37
+ sf_fp8 = raw_sf.to(torch.float8_e4m3fn)
38
+ sf_f32 = sf_fp8.float() * max(scale2, 1e-12)
39
+ scaled = torch.where(sf_f32.unsqueeze(-1) > 0, g / sf_f32.unsqueeze(-1), torch.zeros_like(g))
40
+ ax = scaled.abs()
41
+ nib = torch.zeros_like(ax, dtype=torch.uint8)
42
+ nib = torch.where(ax >= 0.25, torch.tensor(1, device=wd.device, dtype=torch.uint8), nib)
43
+ nib = torch.where(ax >= 0.75, torch.tensor(2, device=wd.device, dtype=torch.uint8), nib)
44
+ nib = torch.where(ax >= 1.25, torch.tensor(3, device=wd.device, dtype=torch.uint8), nib)
45
+ nib = torch.where(ax >= 1.75, torch.tensor(4, device=wd.device, dtype=torch.uint8), nib)
46
+ nib = torch.where(ax >= 2.50, torch.tensor(5, device=wd.device, dtype=torch.uint8), nib)
47
+ nib = torch.where(ax >= 3.50, torch.tensor(6, device=wd.device, dtype=torch.uint8), nib)
48
+ nib = torch.where(ax >= 5.00, torch.tensor(7, device=wd.device, dtype=torch.uint8), nib)
49
+ sign = torch.where(scaled < 0, torch.tensor(8, device=wd.device, dtype=torch.uint8), torch.tensor(0, device=wd.device, dtype=torch.uint8))
50
+ codes = (nib | sign).reshape(n, k // 2, 2)
51
+ qweight = (codes[:, :, 0] | (codes[:, :, 1] << 4)).contiguous()
52
+ return qweight, sf_fp8.view(torch.uint8).contiguous()
53
+
54
+
55
+ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = False) -> PackedWeight:
56
+ wd = w.detach().to(device, torch.bfloat16).contiguous()
57
+ n, k = wd.shape
58
+ amax = wd.abs().amax().to(torch.float32).clamp(min=1e-12)
59
+ scale2 = (amax / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX)).item()
60
+ gs = 1.0 / max(scale2, 1e-12)
61
+ qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
62
+ sf_m, sf_n = swizzled_sf_shape(n, k)
63
+ sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
64
+ return PackedWeight(
65
+ qweight=qweight,
66
+ sf=sf,
67
+ global_scale=torch.tensor([gs], device=device, dtype=torch.float32),
68
+ n=n,
69
+ k=k,
70
+ swiglu=swiglu,
71
+ sf_rowmajor=sf_rowmajor,
72
+ )
73
+
74
+
75
+ def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
76
+ x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
77
+ out = torch.empty((x2.size(0), pw.n), device=x.device, dtype=torch.bfloat16)
78
+ scale2 = 1.0 / pw.global_scale.item()
79
+ ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
80
+ return out.reshape(*x.shape[:-1], pw.n)
81
+
82
+
83
+ def w4a16_gemm(
84
+ A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
85
+ scale2: float, C: torch.Tensor,
86
+ ) -> None:
87
+ """C = A @ B^T where B is NVFP4. Standard layout."""
88
+ ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
89
+
90
+
91
+ def w4a16_gemm_t(
92
+ A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
93
+ scale2: float, C: torch.Tensor,
94
+ ) -> None:
95
+ """Transposed-B variant (N-major) — typically faster for [N>K] shapes."""
96
+ ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
97
+
98
+
99
+ def predequant_nvfp4_to_fp8(
100
+ B_packed: torch.Tensor, B_scale: torch.Tensor,
101
+ scale2: float, B_fp8_out: torch.Tensor,
102
+ ) -> None:
103
+ """Materialize an NVFP4-packed weight into a dense FP8 E4M3 tensor."""
104
+ ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
105
+
106
+
107
+ def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
108
+ """C = A @ B^T where both A is BF16 and B is FP8 E4M3."""
109
+ ops.fp8_gemm_t(A, B_fp8, C)
110
+
111
+
112
+ def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
113
+ """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
114
+ ops.bf16_to_fp8(src, dst)
115
+
116
+
117
+ __all__ = [
118
+ "PackedWeight",
119
+ "pack",
120
+ "gemm",
121
+ "swizzled_sf_shape",
122
+ "w4a16_gemm",
123
+ "w4a16_gemm_t",
124
+ "predequant_nvfp4_to_fp8",
125
+ "fp8_gemm_t",
126
+ "bf16_to_fp8",
127
+ ]
build/torch211-cxx11-cu130-aarch64-linux/__init__.py CHANGED
@@ -1,26 +1,83 @@
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
- """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121).
3
-
4
- Powers the attention Q/K/V/O projections and the dense FFN
5
- (intermediate_size=17408, hidden=5120). Two precisions are supported:
6
-
7
- - W4A16 (NVFP4 weight, BF16 activation): direct path
8
- - FP8 W/A: via ``predequant_nvfp4_to_fp8`` + ``fp8_gemm_t`` for
9
- KV-cache scenarios where the weight is loaded as NVFP4 but the
10
- GEMM runs in FP8 against an FP8-packed activation.
11
  """
12
 
 
13
  import torch
14
 
15
  from ._ops import ops
16
 
17
- __all__ = [
18
- "w4a16_gemm",
19
- "w4a16_gemm_t",
20
- "predequant_nvfp4_to_fp8",
21
- "fp8_gemm_t",
22
- "bf16_to_fp8",
23
- ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
 
25
 
26
  def w4a16_gemm(
@@ -55,3 +112,16 @@ def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
55
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
56
  """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
57
  ops.bf16_to_fp8(src, dst)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121),
3
+ with Hugging Face ``transformers`` ``NVFP4Config`` entry points
4
+ (``PackedWeight``, ``pack``, ``gemm``, ``swizzled_sf_shape``).
 
 
 
 
 
 
5
  """
6
 
7
+ from dataclasses import dataclass
8
  import torch
9
 
10
  from ._ops import ops
11
 
12
+ FLOAT8_E4M3_MAX = 448.0
13
+ FLOAT4_E2M1_MAX = 6.0
14
+ SF_VEC_SIZE = 16
15
+
16
+
17
+ def swizzled_sf_shape(rows: int, cols: int) -> tuple[int, int]:
18
+ return (rows + 127) // 128 * 128, (cols // SF_VEC_SIZE + 3) // 4
19
+
20
+
21
+ @dataclass
22
+ class PackedWeight:
23
+ qweight: torch.Tensor
24
+ sf: torch.Tensor
25
+ global_scale: torch.Tensor
26
+ n: int
27
+ k: int
28
+ swiglu: bool = False
29
+ sf_rowmajor: torch.Tensor | None = None
30
+
31
+
32
+ def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
33
+ n, k = wd.shape
34
+ g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
35
+ block_amax = g.abs().amax(dim=-1)
36
+ raw_sf = (block_amax / (FLOAT4_E2M1_MAX * max(scale2, 1e-12))).clamp(max=FLOAT8_E4M3_MAX)
37
+ sf_fp8 = raw_sf.to(torch.float8_e4m3fn)
38
+ sf_f32 = sf_fp8.float() * max(scale2, 1e-12)
39
+ scaled = torch.where(sf_f32.unsqueeze(-1) > 0, g / sf_f32.unsqueeze(-1), torch.zeros_like(g))
40
+ ax = scaled.abs()
41
+ nib = torch.zeros_like(ax, dtype=torch.uint8)
42
+ nib = torch.where(ax >= 0.25, torch.tensor(1, device=wd.device, dtype=torch.uint8), nib)
43
+ nib = torch.where(ax >= 0.75, torch.tensor(2, device=wd.device, dtype=torch.uint8), nib)
44
+ nib = torch.where(ax >= 1.25, torch.tensor(3, device=wd.device, dtype=torch.uint8), nib)
45
+ nib = torch.where(ax >= 1.75, torch.tensor(4, device=wd.device, dtype=torch.uint8), nib)
46
+ nib = torch.where(ax >= 2.50, torch.tensor(5, device=wd.device, dtype=torch.uint8), nib)
47
+ nib = torch.where(ax >= 3.50, torch.tensor(6, device=wd.device, dtype=torch.uint8), nib)
48
+ nib = torch.where(ax >= 5.00, torch.tensor(7, device=wd.device, dtype=torch.uint8), nib)
49
+ sign = torch.where(scaled < 0, torch.tensor(8, device=wd.device, dtype=torch.uint8), torch.tensor(0, device=wd.device, dtype=torch.uint8))
50
+ codes = (nib | sign).reshape(n, k // 2, 2)
51
+ qweight = (codes[:, :, 0] | (codes[:, :, 1] << 4)).contiguous()
52
+ return qweight, sf_fp8.view(torch.uint8).contiguous()
53
+
54
+
55
+ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = False) -> PackedWeight:
56
+ wd = w.detach().to(device, torch.bfloat16).contiguous()
57
+ n, k = wd.shape
58
+ amax = wd.abs().amax().to(torch.float32).clamp(min=1e-12)
59
+ scale2 = (amax / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX)).item()
60
+ gs = 1.0 / max(scale2, 1e-12)
61
+ qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
62
+ sf_m, sf_n = swizzled_sf_shape(n, k)
63
+ sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
64
+ return PackedWeight(
65
+ qweight=qweight,
66
+ sf=sf,
67
+ global_scale=torch.tensor([gs], device=device, dtype=torch.float32),
68
+ n=n,
69
+ k=k,
70
+ swiglu=swiglu,
71
+ sf_rowmajor=sf_rowmajor,
72
+ )
73
+
74
+
75
+ def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
76
+ x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
77
+ out = torch.empty((x2.size(0), pw.n), device=x.device, dtype=torch.bfloat16)
78
+ scale2 = 1.0 / pw.global_scale.item()
79
+ ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
80
+ return out.reshape(*x.shape[:-1], pw.n)
81
 
82
 
83
  def w4a16_gemm(
 
112
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
113
  """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
114
  ops.bf16_to_fp8(src, dst)
115
+
116
+
117
+ __all__ = [
118
+ "PackedWeight",
119
+ "pack",
120
+ "gemm",
121
+ "swizzled_sf_shape",
122
+ "w4a16_gemm",
123
+ "w4a16_gemm_t",
124
+ "predequant_nvfp4_to_fp8",
125
+ "fp8_gemm_t",
126
+ "bf16_to_fp8",
127
+ ]
build/torch211-cxx11-cu130-aarch64-linux/metadata.json CHANGED
@@ -1,7 +1,7 @@
1
  {
2
  "name": "nvfp4-dense-gemm",
3
  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
4
- "version": 0,
5
  "license": "AGPL-3.0-only",
6
  "python-depends": [],
7
  "backend": {
 
1
  {
2
  "name": "nvfp4-dense-gemm",
3
  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
4
+ "version": 1,
5
  "license": "AGPL-3.0-only",
6
  "python-depends": [],
7
  "backend": {
build/torch211-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py CHANGED
@@ -1,26 +1,127 @@
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")))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121),
3
+ with Hugging Face ``transformers`` ``NVFP4Config`` entry points
4
+ (``PackedWeight``, ``pack``, ``gemm``, ``swizzled_sf_shape``).
5
+ """
6
+
7
+ from dataclasses import dataclass
8
+ import torch
9
+
10
+ from ._ops import ops
11
+
12
+ FLOAT8_E4M3_MAX = 448.0
13
+ FLOAT4_E2M1_MAX = 6.0
14
+ SF_VEC_SIZE = 16
15
+
16
+
17
+ def swizzled_sf_shape(rows: int, cols: int) -> tuple[int, int]:
18
+ return (rows + 127) // 128 * 128, (cols // SF_VEC_SIZE + 3) // 4
19
+
20
+
21
+ @dataclass
22
+ class PackedWeight:
23
+ qweight: torch.Tensor
24
+ sf: torch.Tensor
25
+ global_scale: torch.Tensor
26
+ n: int
27
+ k: int
28
+ swiglu: bool = False
29
+ sf_rowmajor: torch.Tensor | None = None
30
+
31
+
32
+ def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
33
+ n, k = wd.shape
34
+ g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
35
+ block_amax = g.abs().amax(dim=-1)
36
+ raw_sf = (block_amax / (FLOAT4_E2M1_MAX * max(scale2, 1e-12))).clamp(max=FLOAT8_E4M3_MAX)
37
+ sf_fp8 = raw_sf.to(torch.float8_e4m3fn)
38
+ sf_f32 = sf_fp8.float() * max(scale2, 1e-12)
39
+ scaled = torch.where(sf_f32.unsqueeze(-1) > 0, g / sf_f32.unsqueeze(-1), torch.zeros_like(g))
40
+ ax = scaled.abs()
41
+ nib = torch.zeros_like(ax, dtype=torch.uint8)
42
+ nib = torch.where(ax >= 0.25, torch.tensor(1, device=wd.device, dtype=torch.uint8), nib)
43
+ nib = torch.where(ax >= 0.75, torch.tensor(2, device=wd.device, dtype=torch.uint8), nib)
44
+ nib = torch.where(ax >= 1.25, torch.tensor(3, device=wd.device, dtype=torch.uint8), nib)
45
+ nib = torch.where(ax >= 1.75, torch.tensor(4, device=wd.device, dtype=torch.uint8), nib)
46
+ nib = torch.where(ax >= 2.50, torch.tensor(5, device=wd.device, dtype=torch.uint8), nib)
47
+ nib = torch.where(ax >= 3.50, torch.tensor(6, device=wd.device, dtype=torch.uint8), nib)
48
+ nib = torch.where(ax >= 5.00, torch.tensor(7, device=wd.device, dtype=torch.uint8), nib)
49
+ sign = torch.where(scaled < 0, torch.tensor(8, device=wd.device, dtype=torch.uint8), torch.tensor(0, device=wd.device, dtype=torch.uint8))
50
+ codes = (nib | sign).reshape(n, k // 2, 2)
51
+ qweight = (codes[:, :, 0] | (codes[:, :, 1] << 4)).contiguous()
52
+ return qweight, sf_fp8.view(torch.uint8).contiguous()
53
+
54
+
55
+ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = False) -> PackedWeight:
56
+ wd = w.detach().to(device, torch.bfloat16).contiguous()
57
+ n, k = wd.shape
58
+ amax = wd.abs().amax().to(torch.float32).clamp(min=1e-12)
59
+ scale2 = (amax / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX)).item()
60
+ gs = 1.0 / max(scale2, 1e-12)
61
+ qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
62
+ sf_m, sf_n = swizzled_sf_shape(n, k)
63
+ sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
64
+ return PackedWeight(
65
+ qweight=qweight,
66
+ sf=sf,
67
+ global_scale=torch.tensor([gs], device=device, dtype=torch.float32),
68
+ n=n,
69
+ k=k,
70
+ swiglu=swiglu,
71
+ sf_rowmajor=sf_rowmajor,
72
+ )
73
+
74
+
75
+ def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
76
+ x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
77
+ out = torch.empty((x2.size(0), pw.n), device=x.device, dtype=torch.bfloat16)
78
+ scale2 = 1.0 / pw.global_scale.item()
79
+ ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
80
+ return out.reshape(*x.shape[:-1], pw.n)
81
+
82
+
83
+ def w4a16_gemm(
84
+ A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
85
+ scale2: float, C: torch.Tensor,
86
+ ) -> None:
87
+ """C = A @ B^T where B is NVFP4. Standard layout."""
88
+ ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
89
+
90
+
91
+ def w4a16_gemm_t(
92
+ A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
93
+ scale2: float, C: torch.Tensor,
94
+ ) -> None:
95
+ """Transposed-B variant (N-major) — typically faster for [N>K] shapes."""
96
+ ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
97
+
98
+
99
+ def predequant_nvfp4_to_fp8(
100
+ B_packed: torch.Tensor, B_scale: torch.Tensor,
101
+ scale2: float, B_fp8_out: torch.Tensor,
102
+ ) -> None:
103
+ """Materialize an NVFP4-packed weight into a dense FP8 E4M3 tensor."""
104
+ ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
105
+
106
+
107
+ def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
108
+ """C = A @ B^T where both A is BF16 and B is FP8 E4M3."""
109
+ ops.fp8_gemm_t(A, B_fp8, C)
110
+
111
+
112
+ def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
113
+ """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
114
+ ops.bf16_to_fp8(src, dst)
115
+
116
+
117
+ __all__ = [
118
+ "PackedWeight",
119
+ "pack",
120
+ "gemm",
121
+ "swizzled_sf_shape",
122
+ "w4a16_gemm",
123
+ "w4a16_gemm_t",
124
+ "predequant_nvfp4_to_fp8",
125
+ "fp8_gemm_t",
126
+ "bf16_to_fp8",
127
+ ]
build/torch212-cxx11-cu130-aarch64-linux/__init__.py CHANGED
@@ -1,26 +1,83 @@
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
- """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121).
3
-
4
- Powers the attention Q/K/V/O projections and the dense FFN
5
- (intermediate_size=17408, hidden=5120). Two precisions are supported:
6
-
7
- - W4A16 (NVFP4 weight, BF16 activation): direct path
8
- - FP8 W/A: via ``predequant_nvfp4_to_fp8`` + ``fp8_gemm_t`` for
9
- KV-cache scenarios where the weight is loaded as NVFP4 but the
10
- GEMM runs in FP8 against an FP8-packed activation.
11
  """
12
 
 
13
  import torch
14
 
15
  from ._ops import ops
16
 
17
- __all__ = [
18
- "w4a16_gemm",
19
- "w4a16_gemm_t",
20
- "predequant_nvfp4_to_fp8",
21
- "fp8_gemm_t",
22
- "bf16_to_fp8",
23
- ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
 
25
 
26
  def w4a16_gemm(
@@ -55,3 +112,16 @@ def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
55
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
56
  """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
57
  ops.bf16_to_fp8(src, dst)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121),
3
+ with Hugging Face ``transformers`` ``NVFP4Config`` entry points
4
+ (``PackedWeight``, ``pack``, ``gemm``, ``swizzled_sf_shape``).
 
 
 
 
 
 
5
  """
6
 
7
+ from dataclasses import dataclass
8
  import torch
9
 
10
  from ._ops import ops
11
 
12
+ FLOAT8_E4M3_MAX = 448.0
13
+ FLOAT4_E2M1_MAX = 6.0
14
+ SF_VEC_SIZE = 16
15
+
16
+
17
+ def swizzled_sf_shape(rows: int, cols: int) -> tuple[int, int]:
18
+ return (rows + 127) // 128 * 128, (cols // SF_VEC_SIZE + 3) // 4
19
+
20
+
21
+ @dataclass
22
+ class PackedWeight:
23
+ qweight: torch.Tensor
24
+ sf: torch.Tensor
25
+ global_scale: torch.Tensor
26
+ n: int
27
+ k: int
28
+ swiglu: bool = False
29
+ sf_rowmajor: torch.Tensor | None = None
30
+
31
+
32
+ def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
33
+ n, k = wd.shape
34
+ g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
35
+ block_amax = g.abs().amax(dim=-1)
36
+ raw_sf = (block_amax / (FLOAT4_E2M1_MAX * max(scale2, 1e-12))).clamp(max=FLOAT8_E4M3_MAX)
37
+ sf_fp8 = raw_sf.to(torch.float8_e4m3fn)
38
+ sf_f32 = sf_fp8.float() * max(scale2, 1e-12)
39
+ scaled = torch.where(sf_f32.unsqueeze(-1) > 0, g / sf_f32.unsqueeze(-1), torch.zeros_like(g))
40
+ ax = scaled.abs()
41
+ nib = torch.zeros_like(ax, dtype=torch.uint8)
42
+ nib = torch.where(ax >= 0.25, torch.tensor(1, device=wd.device, dtype=torch.uint8), nib)
43
+ nib = torch.where(ax >= 0.75, torch.tensor(2, device=wd.device, dtype=torch.uint8), nib)
44
+ nib = torch.where(ax >= 1.25, torch.tensor(3, device=wd.device, dtype=torch.uint8), nib)
45
+ nib = torch.where(ax >= 1.75, torch.tensor(4, device=wd.device, dtype=torch.uint8), nib)
46
+ nib = torch.where(ax >= 2.50, torch.tensor(5, device=wd.device, dtype=torch.uint8), nib)
47
+ nib = torch.where(ax >= 3.50, torch.tensor(6, device=wd.device, dtype=torch.uint8), nib)
48
+ nib = torch.where(ax >= 5.00, torch.tensor(7, device=wd.device, dtype=torch.uint8), nib)
49
+ sign = torch.where(scaled < 0, torch.tensor(8, device=wd.device, dtype=torch.uint8), torch.tensor(0, device=wd.device, dtype=torch.uint8))
50
+ codes = (nib | sign).reshape(n, k // 2, 2)
51
+ qweight = (codes[:, :, 0] | (codes[:, :, 1] << 4)).contiguous()
52
+ return qweight, sf_fp8.view(torch.uint8).contiguous()
53
+
54
+
55
+ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = False) -> PackedWeight:
56
+ wd = w.detach().to(device, torch.bfloat16).contiguous()
57
+ n, k = wd.shape
58
+ amax = wd.abs().amax().to(torch.float32).clamp(min=1e-12)
59
+ scale2 = (amax / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX)).item()
60
+ gs = 1.0 / max(scale2, 1e-12)
61
+ qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
62
+ sf_m, sf_n = swizzled_sf_shape(n, k)
63
+ sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
64
+ return PackedWeight(
65
+ qweight=qweight,
66
+ sf=sf,
67
+ global_scale=torch.tensor([gs], device=device, dtype=torch.float32),
68
+ n=n,
69
+ k=k,
70
+ swiglu=swiglu,
71
+ sf_rowmajor=sf_rowmajor,
72
+ )
73
+
74
+
75
+ def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
76
+ x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
77
+ out = torch.empty((x2.size(0), pw.n), device=x.device, dtype=torch.bfloat16)
78
+ scale2 = 1.0 / pw.global_scale.item()
79
+ ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
80
+ return out.reshape(*x.shape[:-1], pw.n)
81
 
82
 
83
  def w4a16_gemm(
 
112
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
113
  """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
114
  ops.bf16_to_fp8(src, dst)
115
+
116
+
117
+ __all__ = [
118
+ "PackedWeight",
119
+ "pack",
120
+ "gemm",
121
+ "swizzled_sf_shape",
122
+ "w4a16_gemm",
123
+ "w4a16_gemm_t",
124
+ "predequant_nvfp4_to_fp8",
125
+ "fp8_gemm_t",
126
+ "bf16_to_fp8",
127
+ ]
build/torch212-cxx11-cu130-aarch64-linux/metadata.json CHANGED
@@ -1,7 +1,7 @@
1
  {
2
  "name": "nvfp4-dense-gemm",
3
  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
4
- "version": 0,
5
  "license": "AGPL-3.0-only",
6
  "python-depends": [],
7
  "backend": {
 
1
  {
2
  "name": "nvfp4-dense-gemm",
3
  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
4
+ "version": 1,
5
  "license": "AGPL-3.0-only",
6
  "python-depends": [],
7
  "backend": {
build/torch212-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py CHANGED
@@ -1,26 +1,127 @@
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")))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121),
3
+ with Hugging Face ``transformers`` ``NVFP4Config`` entry points
4
+ (``PackedWeight``, ``pack``, ``gemm``, ``swizzled_sf_shape``).
5
+ """
6
+
7
+ from dataclasses import dataclass
8
+ import torch
9
+
10
+ from ._ops import ops
11
+
12
+ FLOAT8_E4M3_MAX = 448.0
13
+ FLOAT4_E2M1_MAX = 6.0
14
+ SF_VEC_SIZE = 16
15
+
16
+
17
+ def swizzled_sf_shape(rows: int, cols: int) -> tuple[int, int]:
18
+ return (rows + 127) // 128 * 128, (cols // SF_VEC_SIZE + 3) // 4
19
+
20
+
21
+ @dataclass
22
+ class PackedWeight:
23
+ qweight: torch.Tensor
24
+ sf: torch.Tensor
25
+ global_scale: torch.Tensor
26
+ n: int
27
+ k: int
28
+ swiglu: bool = False
29
+ sf_rowmajor: torch.Tensor | None = None
30
+
31
+
32
+ def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
33
+ n, k = wd.shape
34
+ g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
35
+ block_amax = g.abs().amax(dim=-1)
36
+ raw_sf = (block_amax / (FLOAT4_E2M1_MAX * max(scale2, 1e-12))).clamp(max=FLOAT8_E4M3_MAX)
37
+ sf_fp8 = raw_sf.to(torch.float8_e4m3fn)
38
+ sf_f32 = sf_fp8.float() * max(scale2, 1e-12)
39
+ scaled = torch.where(sf_f32.unsqueeze(-1) > 0, g / sf_f32.unsqueeze(-1), torch.zeros_like(g))
40
+ ax = scaled.abs()
41
+ nib = torch.zeros_like(ax, dtype=torch.uint8)
42
+ nib = torch.where(ax >= 0.25, torch.tensor(1, device=wd.device, dtype=torch.uint8), nib)
43
+ nib = torch.where(ax >= 0.75, torch.tensor(2, device=wd.device, dtype=torch.uint8), nib)
44
+ nib = torch.where(ax >= 1.25, torch.tensor(3, device=wd.device, dtype=torch.uint8), nib)
45
+ nib = torch.where(ax >= 1.75, torch.tensor(4, device=wd.device, dtype=torch.uint8), nib)
46
+ nib = torch.where(ax >= 2.50, torch.tensor(5, device=wd.device, dtype=torch.uint8), nib)
47
+ nib = torch.where(ax >= 3.50, torch.tensor(6, device=wd.device, dtype=torch.uint8), nib)
48
+ nib = torch.where(ax >= 5.00, torch.tensor(7, device=wd.device, dtype=torch.uint8), nib)
49
+ sign = torch.where(scaled < 0, torch.tensor(8, device=wd.device, dtype=torch.uint8), torch.tensor(0, device=wd.device, dtype=torch.uint8))
50
+ codes = (nib | sign).reshape(n, k // 2, 2)
51
+ qweight = (codes[:, :, 0] | (codes[:, :, 1] << 4)).contiguous()
52
+ return qweight, sf_fp8.view(torch.uint8).contiguous()
53
+
54
+
55
+ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = False) -> PackedWeight:
56
+ wd = w.detach().to(device, torch.bfloat16).contiguous()
57
+ n, k = wd.shape
58
+ amax = wd.abs().amax().to(torch.float32).clamp(min=1e-12)
59
+ scale2 = (amax / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX)).item()
60
+ gs = 1.0 / max(scale2, 1e-12)
61
+ qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
62
+ sf_m, sf_n = swizzled_sf_shape(n, k)
63
+ sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
64
+ return PackedWeight(
65
+ qweight=qweight,
66
+ sf=sf,
67
+ global_scale=torch.tensor([gs], device=device, dtype=torch.float32),
68
+ n=n,
69
+ k=k,
70
+ swiglu=swiglu,
71
+ sf_rowmajor=sf_rowmajor,
72
+ )
73
+
74
+
75
+ def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
76
+ x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
77
+ out = torch.empty((x2.size(0), pw.n), device=x.device, dtype=torch.bfloat16)
78
+ scale2 = 1.0 / pw.global_scale.item()
79
+ ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
80
+ return out.reshape(*x.shape[:-1], pw.n)
81
+
82
+
83
+ def w4a16_gemm(
84
+ A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
85
+ scale2: float, C: torch.Tensor,
86
+ ) -> None:
87
+ """C = A @ B^T where B is NVFP4. Standard layout."""
88
+ ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
89
+
90
+
91
+ def w4a16_gemm_t(
92
+ A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
93
+ scale2: float, C: torch.Tensor,
94
+ ) -> None:
95
+ """Transposed-B variant (N-major) — typically faster for [N>K] shapes."""
96
+ ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
97
+
98
+
99
+ def predequant_nvfp4_to_fp8(
100
+ B_packed: torch.Tensor, B_scale: torch.Tensor,
101
+ scale2: float, B_fp8_out: torch.Tensor,
102
+ ) -> None:
103
+ """Materialize an NVFP4-packed weight into a dense FP8 E4M3 tensor."""
104
+ ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
105
+
106
+
107
+ def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
108
+ """C = A @ B^T where both A is BF16 and B is FP8 E4M3."""
109
+ ops.fp8_gemm_t(A, B_fp8, C)
110
+
111
+
112
+ def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
113
+ """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
114
+ ops.bf16_to_fp8(src, dst)
115
+
116
+
117
+ __all__ = [
118
+ "PackedWeight",
119
+ "pack",
120
+ "gemm",
121
+ "swizzled_sf_shape",
122
+ "w4a16_gemm",
123
+ "w4a16_gemm_t",
124
+ "predequant_nvfp4_to_fp8",
125
+ "fp8_gemm_t",
126
+ "bf16_to_fp8",
127
+ ]
build/torch212-cxx11-cu132-aarch64-linux/__init__.py CHANGED
@@ -1,26 +1,83 @@
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
- """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121).
3
-
4
- Powers the attention Q/K/V/O projections and the dense FFN
5
- (intermediate_size=17408, hidden=5120). Two precisions are supported:
6
-
7
- - W4A16 (NVFP4 weight, BF16 activation): direct path
8
- - FP8 W/A: via ``predequant_nvfp4_to_fp8`` + ``fp8_gemm_t`` for
9
- KV-cache scenarios where the weight is loaded as NVFP4 but the
10
- GEMM runs in FP8 against an FP8-packed activation.
11
  """
12
 
 
13
  import torch
14
 
15
  from ._ops import ops
16
 
17
- __all__ = [
18
- "w4a16_gemm",
19
- "w4a16_gemm_t",
20
- "predequant_nvfp4_to_fp8",
21
- "fp8_gemm_t",
22
- "bf16_to_fp8",
23
- ]
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
24
 
25
 
26
  def w4a16_gemm(
@@ -55,3 +112,16 @@ def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
55
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
56
  """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
57
  ops.bf16_to_fp8(src, dst)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121),
3
+ with Hugging Face ``transformers`` ``NVFP4Config`` entry points
4
+ (``PackedWeight``, ``pack``, ``gemm``, ``swizzled_sf_shape``).
 
 
 
 
 
 
5
  """
6
 
7
+ from dataclasses import dataclass
8
  import torch
9
 
10
  from ._ops import ops
11
 
12
+ FLOAT8_E4M3_MAX = 448.0
13
+ FLOAT4_E2M1_MAX = 6.0
14
+ SF_VEC_SIZE = 16
15
+
16
+
17
+ def swizzled_sf_shape(rows: int, cols: int) -> tuple[int, int]:
18
+ return (rows + 127) // 128 * 128, (cols // SF_VEC_SIZE + 3) // 4
19
+
20
+
21
+ @dataclass
22
+ class PackedWeight:
23
+ qweight: torch.Tensor
24
+ sf: torch.Tensor
25
+ global_scale: torch.Tensor
26
+ n: int
27
+ k: int
28
+ swiglu: bool = False
29
+ sf_rowmajor: torch.Tensor | None = None
30
+
31
+
32
+ def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
33
+ n, k = wd.shape
34
+ g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
35
+ block_amax = g.abs().amax(dim=-1)
36
+ raw_sf = (block_amax / (FLOAT4_E2M1_MAX * max(scale2, 1e-12))).clamp(max=FLOAT8_E4M3_MAX)
37
+ sf_fp8 = raw_sf.to(torch.float8_e4m3fn)
38
+ sf_f32 = sf_fp8.float() * max(scale2, 1e-12)
39
+ scaled = torch.where(sf_f32.unsqueeze(-1) > 0, g / sf_f32.unsqueeze(-1), torch.zeros_like(g))
40
+ ax = scaled.abs()
41
+ nib = torch.zeros_like(ax, dtype=torch.uint8)
42
+ nib = torch.where(ax >= 0.25, torch.tensor(1, device=wd.device, dtype=torch.uint8), nib)
43
+ nib = torch.where(ax >= 0.75, torch.tensor(2, device=wd.device, dtype=torch.uint8), nib)
44
+ nib = torch.where(ax >= 1.25, torch.tensor(3, device=wd.device, dtype=torch.uint8), nib)
45
+ nib = torch.where(ax >= 1.75, torch.tensor(4, device=wd.device, dtype=torch.uint8), nib)
46
+ nib = torch.where(ax >= 2.50, torch.tensor(5, device=wd.device, dtype=torch.uint8), nib)
47
+ nib = torch.where(ax >= 3.50, torch.tensor(6, device=wd.device, dtype=torch.uint8), nib)
48
+ nib = torch.where(ax >= 5.00, torch.tensor(7, device=wd.device, dtype=torch.uint8), nib)
49
+ sign = torch.where(scaled < 0, torch.tensor(8, device=wd.device, dtype=torch.uint8), torch.tensor(0, device=wd.device, dtype=torch.uint8))
50
+ codes = (nib | sign).reshape(n, k // 2, 2)
51
+ qweight = (codes[:, :, 0] | (codes[:, :, 1] << 4)).contiguous()
52
+ return qweight, sf_fp8.view(torch.uint8).contiguous()
53
+
54
+
55
+ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = False) -> PackedWeight:
56
+ wd = w.detach().to(device, torch.bfloat16).contiguous()
57
+ n, k = wd.shape
58
+ amax = wd.abs().amax().to(torch.float32).clamp(min=1e-12)
59
+ scale2 = (amax / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX)).item()
60
+ gs = 1.0 / max(scale2, 1e-12)
61
+ qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
62
+ sf_m, sf_n = swizzled_sf_shape(n, k)
63
+ sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
64
+ return PackedWeight(
65
+ qweight=qweight,
66
+ sf=sf,
67
+ global_scale=torch.tensor([gs], device=device, dtype=torch.float32),
68
+ n=n,
69
+ k=k,
70
+ swiglu=swiglu,
71
+ sf_rowmajor=sf_rowmajor,
72
+ )
73
+
74
+
75
+ def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
76
+ x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
77
+ out = torch.empty((x2.size(0), pw.n), device=x.device, dtype=torch.bfloat16)
78
+ scale2 = 1.0 / pw.global_scale.item()
79
+ ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
80
+ return out.reshape(*x.shape[:-1], pw.n)
81
 
82
 
83
  def w4a16_gemm(
 
112
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
113
  """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
114
  ops.bf16_to_fp8(src, dst)
115
+
116
+
117
+ __all__ = [
118
+ "PackedWeight",
119
+ "pack",
120
+ "gemm",
121
+ "swizzled_sf_shape",
122
+ "w4a16_gemm",
123
+ "w4a16_gemm_t",
124
+ "predequant_nvfp4_to_fp8",
125
+ "fp8_gemm_t",
126
+ "bf16_to_fp8",
127
+ ]
build/torch212-cxx11-cu132-aarch64-linux/metadata.json CHANGED
@@ -1,7 +1,7 @@
1
  {
2
  "name": "nvfp4-dense-gemm",
3
  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
4
- "version": 0,
5
  "license": "AGPL-3.0-only",
6
  "python-depends": [],
7
  "backend": {
 
1
  {
2
  "name": "nvfp4-dense-gemm",
3
  "id": "_nvfp4_dense_gemm_cuda_bcba8ad",
4
+ "version": 1,
5
  "license": "AGPL-3.0-only",
6
  "python-depends": [],
7
  "backend": {
build/torch212-cxx11-cu132-aarch64-linux/nvfp4_dense_gemm/__init__.py CHANGED
@@ -1,26 +1,127 @@
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")))
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for Qwen3.6-27B on GB10 (SM121),
3
+ with Hugging Face ``transformers`` ``NVFP4Config`` entry points
4
+ (``PackedWeight``, ``pack``, ``gemm``, ``swizzled_sf_shape``).
5
+ """
6
+
7
+ from dataclasses import dataclass
8
+ import torch
9
+
10
+ from ._ops import ops
11
+
12
+ FLOAT8_E4M3_MAX = 448.0
13
+ FLOAT4_E2M1_MAX = 6.0
14
+ SF_VEC_SIZE = 16
15
+
16
+
17
+ def swizzled_sf_shape(rows: int, cols: int) -> tuple[int, int]:
18
+ return (rows + 127) // 128 * 128, (cols // SF_VEC_SIZE + 3) // 4
19
+
20
+
21
+ @dataclass
22
+ class PackedWeight:
23
+ qweight: torch.Tensor
24
+ sf: torch.Tensor
25
+ global_scale: torch.Tensor
26
+ n: int
27
+ k: int
28
+ swiglu: bool = False
29
+ sf_rowmajor: torch.Tensor | None = None
30
+
31
+
32
+ def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
33
+ n, k = wd.shape
34
+ g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
35
+ block_amax = g.abs().amax(dim=-1)
36
+ raw_sf = (block_amax / (FLOAT4_E2M1_MAX * max(scale2, 1e-12))).clamp(max=FLOAT8_E4M3_MAX)
37
+ sf_fp8 = raw_sf.to(torch.float8_e4m3fn)
38
+ sf_f32 = sf_fp8.float() * max(scale2, 1e-12)
39
+ scaled = torch.where(sf_f32.unsqueeze(-1) > 0, g / sf_f32.unsqueeze(-1), torch.zeros_like(g))
40
+ ax = scaled.abs()
41
+ nib = torch.zeros_like(ax, dtype=torch.uint8)
42
+ nib = torch.where(ax >= 0.25, torch.tensor(1, device=wd.device, dtype=torch.uint8), nib)
43
+ nib = torch.where(ax >= 0.75, torch.tensor(2, device=wd.device, dtype=torch.uint8), nib)
44
+ nib = torch.where(ax >= 1.25, torch.tensor(3, device=wd.device, dtype=torch.uint8), nib)
45
+ nib = torch.where(ax >= 1.75, torch.tensor(4, device=wd.device, dtype=torch.uint8), nib)
46
+ nib = torch.where(ax >= 2.50, torch.tensor(5, device=wd.device, dtype=torch.uint8), nib)
47
+ nib = torch.where(ax >= 3.50, torch.tensor(6, device=wd.device, dtype=torch.uint8), nib)
48
+ nib = torch.where(ax >= 5.00, torch.tensor(7, device=wd.device, dtype=torch.uint8), nib)
49
+ sign = torch.where(scaled < 0, torch.tensor(8, device=wd.device, dtype=torch.uint8), torch.tensor(0, device=wd.device, dtype=torch.uint8))
50
+ codes = (nib | sign).reshape(n, k // 2, 2)
51
+ qweight = (codes[:, :, 0] | (codes[:, :, 1] << 4)).contiguous()
52
+ return qweight, sf_fp8.view(torch.uint8).contiguous()
53
+
54
+
55
+ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = False) -> PackedWeight:
56
+ wd = w.detach().to(device, torch.bfloat16).contiguous()
57
+ n, k = wd.shape
58
+ amax = wd.abs().amax().to(torch.float32).clamp(min=1e-12)
59
+ scale2 = (amax / (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX)).item()
60
+ gs = 1.0 / max(scale2, 1e-12)
61
+ qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
62
+ sf_m, sf_n = swizzled_sf_shape(n, k)
63
+ sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
64
+ return PackedWeight(
65
+ qweight=qweight,
66
+ sf=sf,
67
+ global_scale=torch.tensor([gs], device=device, dtype=torch.float32),
68
+ n=n,
69
+ k=k,
70
+ swiglu=swiglu,
71
+ sf_rowmajor=sf_rowmajor,
72
+ )
73
+
74
+
75
+ def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
76
+ x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
77
+ out = torch.empty((x2.size(0), pw.n), device=x.device, dtype=torch.bfloat16)
78
+ scale2 = 1.0 / pw.global_scale.item()
79
+ ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
80
+ return out.reshape(*x.shape[:-1], pw.n)
81
+
82
+
83
+ def w4a16_gemm(
84
+ A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
85
+ scale2: float, C: torch.Tensor,
86
+ ) -> None:
87
+ """C = A @ B^T where B is NVFP4. Standard layout."""
88
+ ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
89
+
90
+
91
+ def w4a16_gemm_t(
92
+ A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
93
+ scale2: float, C: torch.Tensor,
94
+ ) -> None:
95
+ """Transposed-B variant (N-major) — typically faster for [N>K] shapes."""
96
+ ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
97
+
98
+
99
+ def predequant_nvfp4_to_fp8(
100
+ B_packed: torch.Tensor, B_scale: torch.Tensor,
101
+ scale2: float, B_fp8_out: torch.Tensor,
102
+ ) -> None:
103
+ """Materialize an NVFP4-packed weight into a dense FP8 E4M3 tensor."""
104
+ ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
105
+
106
+
107
+ def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
108
+ """C = A @ B^T where both A is BF16 and B is FP8 E4M3."""
109
+ ops.fp8_gemm_t(A, B_fp8, C)
110
+
111
+
112
+ def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
113
+ """Pair-wise BF16 → FP8 E4M3 quantization (no scaling, uses PTX cvt)."""
114
+ ops.bf16_to_fp8(src, dst)
115
+
116
+
117
+ __all__ = [
118
+ "PackedWeight",
119
+ "pack",
120
+ "gemm",
121
+ "swizzled_sf_shape",
122
+ "w4a16_gemm",
123
+ "w4a16_gemm_t",
124
+ "predequant_nvfp4_to_fp8",
125
+ "fp8_gemm_t",
126
+ "bf16_to_fp8",
127
+ ]