Hugging Face's logo

Kernels:
Atlas-Inference
/
nvfp4-dense-gemm

Kernel card Files
xet
 Community
AzeezIsh commited on
Commit
9ce7aeb
·
verified ·
1 Parent(s): 4c5264e

Add w4a16_gemv_sw decode and w4a16_dequant prefill to v1

Browse files
Files changed (12) hide show
  1. build/torch210-cxx11-cu130-aarch64-linux/__init__.py +40 -11
  2. build/torch210-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py +40 -11
  3. build/torch211-cxx11-cu130-aarch64-linux/__init__.py +40 -11
  4. build/torch211-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py +40 -11
  5. build/torch212-cxx11-cu130-aarch64-linux/__init__.py +40 -11
  6. build/torch212-cxx11-cu130-aarch64-linux/__pycache__/__init__.cpython-312.pyc +0 -0
  7. build/torch212-cxx11-cu130-aarch64-linux/__pycache__/_ops.cpython-312.pyc +0 -0
  8. build/torch212-cxx11-cu130-aarch64-linux/_atlas_w4a16_gemv.abi3.so +3 -0
  9. build/torch212-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py +40 -11
  10. build/torch212-cxx11-cu132-aarch64-linux/__init__.py +40 -11
  11. build/torch212-cxx11-cu132-aarch64-linux/_atlas_w4a16_gemv.abi3.so +3 -0
  12. build/torch212-cxx11-cu132-aarch64-linux/nvfp4_dense_gemm/__init__.py +40 -11
build/torch210-cxx11-cu130-aarch64-linux/__init__.py CHANGED
@@ -1,14 +1,23 @@
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
@@ -29,6 +38,17 @@ class PackedWeight:
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)
@@ -61,10 +81,15 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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,
@@ -74,8 +99,17 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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
 
@@ -84,7 +118,6 @@ 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
 
@@ -92,7 +125,6 @@ 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
 
@@ -100,17 +132,14 @@ 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
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for 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 os
9
  import torch
10
+ import torch.nn.functional as F
11
 
12
  from ._ops import ops
13
 
14
+ _gemv_so = os.path.join(os.path.dirname(__file__), "_atlas_w4a16_gemv.abi3.so")
15
+ if os.path.exists(_gemv_so):
16
+ torch.ops.load_library(_gemv_so)
17
+ _gemv_ops = torch.ops._atlas_w4a16_gemv
18
+ else:
19
+ _gemv_ops = None
20
+
21
  FLOAT8_E4M3_MAX = 448.0
22
  FLOAT4_E2M1_MAX = 6.0
23
  SF_VEC_SIZE = 16
 
38
  sf_rowmajor: torch.Tensor | None = None
39
 
40
 
41
+ def _get_scale2(global_scale: torch.Tensor) -> float:
42
+ s2 = getattr(global_scale, "_atlas_scale2", None)
43
+ if s2 is None:
44
+ s2 = 1.0 / float(global_scale.item())
45
+ try:
46
+ setattr(global_scale, "_atlas_scale2", s2)
47
+ except Exception:
48
+ pass
49
+ return s2
50
+
51
+
52
  def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
53
  n, k = wd.shape
54
  g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
 
81
  qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
82
  sf_m, sf_n = swizzled_sf_shape(n, k)
83
  sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
84
+ gs_t = torch.tensor([gs], device=device, dtype=torch.float32)
85
+ try:
86
+ setattr(gs_t, "_atlas_scale2", scale2)
87
+ except Exception:
88
+ pass
89
  return PackedWeight(
90
  qweight=qweight,
91
  sf=sf,
92
+ global_scale=gs_t,
93
  n=n,
94
  k=k,
95
  swiglu=swiglu,
 
99
 
100
  def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
101
  x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
102
+ m_rows = x2.size(0)
103
+ scale2 = _get_scale2(pw.global_scale)
104
+ if _gemv_ops is not None:
105
+ if m_rows <= 8:
106
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
107
+ _gemv_ops.w4a16_gemv_sw(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
108
+ return out.reshape(*x.shape[:-1], pw.n)
109
+ w_bf16 = torch.empty((pw.n, pw.k), device=x.device, dtype=torch.bfloat16)
110
+ _gemv_ops.w4a16_dequant(pw.qweight, pw.sf_rowmajor, scale2, w_bf16)
111
+ return F.linear(x2, w_bf16).reshape(*x.shape[:-1], pw.n)
112
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
113
  ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
114
  return out.reshape(*x.shape[:-1], pw.n)
115
 
 
118
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
119
  scale2: float, C: torch.Tensor,
120
  ) -> None:
 
121
  ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
122
 
123
 
 
125
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
126
  scale2: float, C: torch.Tensor,
127
  ) -> None:
 
128
  ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
129
 
130
 
 
132
  B_packed: torch.Tensor, B_scale: torch.Tensor,
133
  scale2: float, B_fp8_out: torch.Tensor,
134
  ) -> None:
 
135
  ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
136
 
137
 
138
  def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
 
139
  ops.fp8_gemm_t(A, B_fp8, C)
140
 
141
 
142
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
 
143
  ops.bf16_to_fp8(src, dst)
144
 
145
 
build/torch210-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py CHANGED
@@ -1,14 +1,23 @@
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
@@ -29,6 +38,17 @@ class PackedWeight:
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)
@@ -61,10 +81,15 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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,
@@ -74,8 +99,17 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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
 
@@ -84,7 +118,6 @@ 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
 
@@ -92,7 +125,6 @@ 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
 
@@ -100,17 +132,14 @@ 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
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for 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 os
9
  import torch
10
+ import torch.nn.functional as F
11
 
12
  from ._ops import ops
13
 
14
+ _gemv_so = os.path.join(os.path.dirname(__file__), "_atlas_w4a16_gemv.abi3.so")
15
+ if os.path.exists(_gemv_so):
16
+ torch.ops.load_library(_gemv_so)
17
+ _gemv_ops = torch.ops._atlas_w4a16_gemv
18
+ else:
19
+ _gemv_ops = None
20
+
21
  FLOAT8_E4M3_MAX = 448.0
22
  FLOAT4_E2M1_MAX = 6.0
23
  SF_VEC_SIZE = 16
 
38
  sf_rowmajor: torch.Tensor | None = None
39
 
40
 
41
+ def _get_scale2(global_scale: torch.Tensor) -> float:
42
+ s2 = getattr(global_scale, "_atlas_scale2", None)
43
+ if s2 is None:
44
+ s2 = 1.0 / float(global_scale.item())
45
+ try:
46
+ setattr(global_scale, "_atlas_scale2", s2)
47
+ except Exception:
48
+ pass
49
+ return s2
50
+
51
+
52
  def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
53
  n, k = wd.shape
54
  g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
 
81
  qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
82
  sf_m, sf_n = swizzled_sf_shape(n, k)
83
  sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
84
+ gs_t = torch.tensor([gs], device=device, dtype=torch.float32)
85
+ try:
86
+ setattr(gs_t, "_atlas_scale2", scale2)
87
+ except Exception:
88
+ pass
89
  return PackedWeight(
90
  qweight=qweight,
91
  sf=sf,
92
+ global_scale=gs_t,
93
  n=n,
94
  k=k,
95
  swiglu=swiglu,
 
99
 
100
  def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
101
  x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
102
+ m_rows = x2.size(0)
103
+ scale2 = _get_scale2(pw.global_scale)
104
+ if _gemv_ops is not None:
105
+ if m_rows <= 8:
106
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
107
+ _gemv_ops.w4a16_gemv_sw(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
108
+ return out.reshape(*x.shape[:-1], pw.n)
109
+ w_bf16 = torch.empty((pw.n, pw.k), device=x.device, dtype=torch.bfloat16)
110
+ _gemv_ops.w4a16_dequant(pw.qweight, pw.sf_rowmajor, scale2, w_bf16)
111
+ return F.linear(x2, w_bf16).reshape(*x.shape[:-1], pw.n)
112
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
113
  ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
114
  return out.reshape(*x.shape[:-1], pw.n)
115
 
 
118
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
119
  scale2: float, C: torch.Tensor,
120
  ) -> None:
 
121
  ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
122
 
123
 
 
125
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
126
  scale2: float, C: torch.Tensor,
127
  ) -> None:
 
128
  ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
129
 
130
 
 
132
  B_packed: torch.Tensor, B_scale: torch.Tensor,
133
  scale2: float, B_fp8_out: torch.Tensor,
134
  ) -> None:
 
135
  ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
136
 
137
 
138
  def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
 
139
  ops.fp8_gemm_t(A, B_fp8, C)
140
 
141
 
142
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
 
143
  ops.bf16_to_fp8(src, dst)
144
 
145
 
build/torch211-cxx11-cu130-aarch64-linux/__init__.py CHANGED
@@ -1,14 +1,23 @@
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
@@ -29,6 +38,17 @@ class PackedWeight:
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)
@@ -61,10 +81,15 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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,
@@ -74,8 +99,17 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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
 
@@ -84,7 +118,6 @@ 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
 
@@ -92,7 +125,6 @@ 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
 
@@ -100,17 +132,14 @@ 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
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for 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 os
9
  import torch
10
+ import torch.nn.functional as F
11
 
12
  from ._ops import ops
13
 
14
+ _gemv_so = os.path.join(os.path.dirname(__file__), "_atlas_w4a16_gemv.abi3.so")
15
+ if os.path.exists(_gemv_so):
16
+ torch.ops.load_library(_gemv_so)
17
+ _gemv_ops = torch.ops._atlas_w4a16_gemv
18
+ else:
19
+ _gemv_ops = None
20
+
21
  FLOAT8_E4M3_MAX = 448.0
22
  FLOAT4_E2M1_MAX = 6.0
23
  SF_VEC_SIZE = 16
 
38
  sf_rowmajor: torch.Tensor | None = None
39
 
40
 
41
+ def _get_scale2(global_scale: torch.Tensor) -> float:
42
+ s2 = getattr(global_scale, "_atlas_scale2", None)
43
+ if s2 is None:
44
+ s2 = 1.0 / float(global_scale.item())
45
+ try:
46
+ setattr(global_scale, "_atlas_scale2", s2)
47
+ except Exception:
48
+ pass
49
+ return s2
50
+
51
+
52
  def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
53
  n, k = wd.shape
54
  g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
 
81
  qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
82
  sf_m, sf_n = swizzled_sf_shape(n, k)
83
  sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
84
+ gs_t = torch.tensor([gs], device=device, dtype=torch.float32)
85
+ try:
86
+ setattr(gs_t, "_atlas_scale2", scale2)
87
+ except Exception:
88
+ pass
89
  return PackedWeight(
90
  qweight=qweight,
91
  sf=sf,
92
+ global_scale=gs_t,
93
  n=n,
94
  k=k,
95
  swiglu=swiglu,
 
99
 
100
  def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
101
  x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
102
+ m_rows = x2.size(0)
103
+ scale2 = _get_scale2(pw.global_scale)
104
+ if _gemv_ops is not None:
105
+ if m_rows <= 8:
106
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
107
+ _gemv_ops.w4a16_gemv_sw(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
108
+ return out.reshape(*x.shape[:-1], pw.n)
109
+ w_bf16 = torch.empty((pw.n, pw.k), device=x.device, dtype=torch.bfloat16)
110
+ _gemv_ops.w4a16_dequant(pw.qweight, pw.sf_rowmajor, scale2, w_bf16)
111
+ return F.linear(x2, w_bf16).reshape(*x.shape[:-1], pw.n)
112
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
113
  ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
114
  return out.reshape(*x.shape[:-1], pw.n)
115
 
 
118
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
119
  scale2: float, C: torch.Tensor,
120
  ) -> None:
 
121
  ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
122
 
123
 
 
125
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
126
  scale2: float, C: torch.Tensor,
127
  ) -> None:
 
128
  ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
129
 
130
 
 
132
  B_packed: torch.Tensor, B_scale: torch.Tensor,
133
  scale2: float, B_fp8_out: torch.Tensor,
134
  ) -> None:
 
135
  ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
136
 
137
 
138
  def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
 
139
  ops.fp8_gemm_t(A, B_fp8, C)
140
 
141
 
142
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
 
143
  ops.bf16_to_fp8(src, dst)
144
 
145
 
build/torch211-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py CHANGED
@@ -1,14 +1,23 @@
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
@@ -29,6 +38,17 @@ class PackedWeight:
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)
@@ -61,10 +81,15 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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,
@@ -74,8 +99,17 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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
 
@@ -84,7 +118,6 @@ 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
 
@@ -92,7 +125,6 @@ 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
 
@@ -100,17 +132,14 @@ 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
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for 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 os
9
  import torch
10
+ import torch.nn.functional as F
11
 
12
  from ._ops import ops
13
 
14
+ _gemv_so = os.path.join(os.path.dirname(__file__), "_atlas_w4a16_gemv.abi3.so")
15
+ if os.path.exists(_gemv_so):
16
+ torch.ops.load_library(_gemv_so)
17
+ _gemv_ops = torch.ops._atlas_w4a16_gemv
18
+ else:
19
+ _gemv_ops = None
20
+
21
  FLOAT8_E4M3_MAX = 448.0
22
  FLOAT4_E2M1_MAX = 6.0
23
  SF_VEC_SIZE = 16
 
38
  sf_rowmajor: torch.Tensor | None = None
39
 
40
 
41
+ def _get_scale2(global_scale: torch.Tensor) -> float:
42
+ s2 = getattr(global_scale, "_atlas_scale2", None)
43
+ if s2 is None:
44
+ s2 = 1.0 / float(global_scale.item())
45
+ try:
46
+ setattr(global_scale, "_atlas_scale2", s2)
47
+ except Exception:
48
+ pass
49
+ return s2
50
+
51
+
52
  def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
53
  n, k = wd.shape
54
  g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
 
81
  qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
82
  sf_m, sf_n = swizzled_sf_shape(n, k)
83
  sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
84
+ gs_t = torch.tensor([gs], device=device, dtype=torch.float32)
85
+ try:
86
+ setattr(gs_t, "_atlas_scale2", scale2)
87
+ except Exception:
88
+ pass
89
  return PackedWeight(
90
  qweight=qweight,
91
  sf=sf,
92
+ global_scale=gs_t,
93
  n=n,
94
  k=k,
95
  swiglu=swiglu,
 
99
 
100
  def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
101
  x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
102
+ m_rows = x2.size(0)
103
+ scale2 = _get_scale2(pw.global_scale)
104
+ if _gemv_ops is not None:
105
+ if m_rows <= 8:
106
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
107
+ _gemv_ops.w4a16_gemv_sw(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
108
+ return out.reshape(*x.shape[:-1], pw.n)
109
+ w_bf16 = torch.empty((pw.n, pw.k), device=x.device, dtype=torch.bfloat16)
110
+ _gemv_ops.w4a16_dequant(pw.qweight, pw.sf_rowmajor, scale2, w_bf16)
111
+ return F.linear(x2, w_bf16).reshape(*x.shape[:-1], pw.n)
112
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
113
  ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
114
  return out.reshape(*x.shape[:-1], pw.n)
115
 
 
118
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
119
  scale2: float, C: torch.Tensor,
120
  ) -> None:
 
121
  ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
122
 
123
 
 
125
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
126
  scale2: float, C: torch.Tensor,
127
  ) -> None:
 
128
  ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
129
 
130
 
 
132
  B_packed: torch.Tensor, B_scale: torch.Tensor,
133
  scale2: float, B_fp8_out: torch.Tensor,
134
  ) -> None:
 
135
  ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
136
 
137
 
138
  def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
 
139
  ops.fp8_gemm_t(A, B_fp8, C)
140
 
141
 
142
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
 
143
  ops.bf16_to_fp8(src, dst)
144
 
145
 
build/torch212-cxx11-cu130-aarch64-linux/__init__.py CHANGED
@@ -1,14 +1,23 @@
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
@@ -29,6 +38,17 @@ class PackedWeight:
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)
@@ -61,10 +81,15 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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,
@@ -74,8 +99,17 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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
 
@@ -84,7 +118,6 @@ 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
 
@@ -92,7 +125,6 @@ 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
 
@@ -100,17 +132,14 @@ 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
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for 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 os
9
  import torch
10
+ import torch.nn.functional as F
11
 
12
  from ._ops import ops
13
 
14
+ _gemv_so = os.path.join(os.path.dirname(__file__), "_atlas_w4a16_gemv.abi3.so")
15
+ if os.path.exists(_gemv_so):
16
+ torch.ops.load_library(_gemv_so)
17
+ _gemv_ops = torch.ops._atlas_w4a16_gemv
18
+ else:
19
+ _gemv_ops = None
20
+
21
  FLOAT8_E4M3_MAX = 448.0
22
  FLOAT4_E2M1_MAX = 6.0
23
  SF_VEC_SIZE = 16
 
38
  sf_rowmajor: torch.Tensor | None = None
39
 
40
 
41
+ def _get_scale2(global_scale: torch.Tensor) -> float:
42
+ s2 = getattr(global_scale, "_atlas_scale2", None)
43
+ if s2 is None:
44
+ s2 = 1.0 / float(global_scale.item())
45
+ try:
46
+ setattr(global_scale, "_atlas_scale2", s2)
47
+ except Exception:
48
+ pass
49
+ return s2
50
+
51
+
52
  def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
53
  n, k = wd.shape
54
  g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
 
81
  qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
82
  sf_m, sf_n = swizzled_sf_shape(n, k)
83
  sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
84
+ gs_t = torch.tensor([gs], device=device, dtype=torch.float32)
85
+ try:
86
+ setattr(gs_t, "_atlas_scale2", scale2)
87
+ except Exception:
88
+ pass
89
  return PackedWeight(
90
  qweight=qweight,
91
  sf=sf,
92
+ global_scale=gs_t,
93
  n=n,
94
  k=k,
95
  swiglu=swiglu,
 
99
 
100
  def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
101
  x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
102
+ m_rows = x2.size(0)
103
+ scale2 = _get_scale2(pw.global_scale)
104
+ if _gemv_ops is not None:
105
+ if m_rows <= 8:
106
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
107
+ _gemv_ops.w4a16_gemv_sw(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
108
+ return out.reshape(*x.shape[:-1], pw.n)
109
+ w_bf16 = torch.empty((pw.n, pw.k), device=x.device, dtype=torch.bfloat16)
110
+ _gemv_ops.w4a16_dequant(pw.qweight, pw.sf_rowmajor, scale2, w_bf16)
111
+ return F.linear(x2, w_bf16).reshape(*x.shape[:-1], pw.n)
112
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
113
  ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
114
  return out.reshape(*x.shape[:-1], pw.n)
115
 
 
118
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
119
  scale2: float, C: torch.Tensor,
120
  ) -> None:
 
121
  ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
122
 
123
 
 
125
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
126
  scale2: float, C: torch.Tensor,
127
  ) -> None:
 
128
  ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
129
 
130
 
 
132
  B_packed: torch.Tensor, B_scale: torch.Tensor,
133
  scale2: float, B_fp8_out: torch.Tensor,
134
  ) -> None:
 
135
  ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
136
 
137
 
138
  def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
 
139
  ops.fp8_gemm_t(A, B_fp8, C)
140
 
141
 
142
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
 
143
  ops.bf16_to_fp8(src, dst)
144
 
145
 
build/torch212-cxx11-cu130-aarch64-linux/__pycache__/__init__.cpython-312.pyc ADDED
Binary file (9.58 kB). View file
 
build/torch212-cxx11-cu130-aarch64-linux/__pycache__/_ops.cpython-312.pyc ADDED
Binary file (644 Bytes). View file
 
build/torch212-cxx11-cu130-aarch64-linux/_atlas_w4a16_gemv.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:06039b9894e062e1bcc52210c2555513cdeb2de5da160dcf538166ff9f9df87f
3
+ size 1416920
build/torch212-cxx11-cu130-aarch64-linux/nvfp4_dense_gemm/__init__.py CHANGED
@@ -1,14 +1,23 @@
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
@@ -29,6 +38,17 @@ class PackedWeight:
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)
@@ -61,10 +81,15 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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,
@@ -74,8 +99,17 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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
 
@@ -84,7 +118,6 @@ 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
 
@@ -92,7 +125,6 @@ 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
 
@@ -100,17 +132,14 @@ 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
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for 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 os
9
  import torch
10
+ import torch.nn.functional as F
11
 
12
  from ._ops import ops
13
 
14
+ _gemv_so = os.path.join(os.path.dirname(__file__), "_atlas_w4a16_gemv.abi3.so")
15
+ if os.path.exists(_gemv_so):
16
+ torch.ops.load_library(_gemv_so)
17
+ _gemv_ops = torch.ops._atlas_w4a16_gemv
18
+ else:
19
+ _gemv_ops = None
20
+
21
  FLOAT8_E4M3_MAX = 448.0
22
  FLOAT4_E2M1_MAX = 6.0
23
  SF_VEC_SIZE = 16
 
38
  sf_rowmajor: torch.Tensor | None = None
39
 
40
 
41
+ def _get_scale2(global_scale: torch.Tensor) -> float:
42
+ s2 = getattr(global_scale, "_atlas_scale2", None)
43
+ if s2 is None:
44
+ s2 = 1.0 / float(global_scale.item())
45
+ try:
46
+ setattr(global_scale, "_atlas_scale2", s2)
47
+ except Exception:
48
+ pass
49
+ return s2
50
+
51
+
52
  def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
53
  n, k = wd.shape
54
  g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
 
81
  qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
82
  sf_m, sf_n = swizzled_sf_shape(n, k)
83
  sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
84
+ gs_t = torch.tensor([gs], device=device, dtype=torch.float32)
85
+ try:
86
+ setattr(gs_t, "_atlas_scale2", scale2)
87
+ except Exception:
88
+ pass
89
  return PackedWeight(
90
  qweight=qweight,
91
  sf=sf,
92
+ global_scale=gs_t,
93
  n=n,
94
  k=k,
95
  swiglu=swiglu,
 
99
 
100
  def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
101
  x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
102
+ m_rows = x2.size(0)
103
+ scale2 = _get_scale2(pw.global_scale)
104
+ if _gemv_ops is not None:
105
+ if m_rows <= 8:
106
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
107
+ _gemv_ops.w4a16_gemv_sw(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
108
+ return out.reshape(*x.shape[:-1], pw.n)
109
+ w_bf16 = torch.empty((pw.n, pw.k), device=x.device, dtype=torch.bfloat16)
110
+ _gemv_ops.w4a16_dequant(pw.qweight, pw.sf_rowmajor, scale2, w_bf16)
111
+ return F.linear(x2, w_bf16).reshape(*x.shape[:-1], pw.n)
112
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
113
  ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
114
  return out.reshape(*x.shape[:-1], pw.n)
115
 
 
118
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
119
  scale2: float, C: torch.Tensor,
120
  ) -> None:
 
121
  ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
122
 
123
 
 
125
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
126
  scale2: float, C: torch.Tensor,
127
  ) -> None:
 
128
  ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
129
 
130
 
 
132
  B_packed: torch.Tensor, B_scale: torch.Tensor,
133
  scale2: float, B_fp8_out: torch.Tensor,
134
  ) -> None:
 
135
  ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
136
 
137
 
138
  def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
 
139
  ops.fp8_gemm_t(A, B_fp8, C)
140
 
141
 
142
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
 
143
  ops.bf16_to_fp8(src, dst)
144
 
145
 
build/torch212-cxx11-cu132-aarch64-linux/__init__.py CHANGED
@@ -1,14 +1,23 @@
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
@@ -29,6 +38,17 @@ class PackedWeight:
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)
@@ -61,10 +81,15 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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,
@@ -74,8 +99,17 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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
 
@@ -84,7 +118,6 @@ 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
 
@@ -92,7 +125,6 @@ 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
 
@@ -100,17 +132,14 @@ 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
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for 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 os
9
  import torch
10
+ import torch.nn.functional as F
11
 
12
  from ._ops import ops
13
 
14
+ _gemv_so = os.path.join(os.path.dirname(__file__), "_atlas_w4a16_gemv.abi3.so")
15
+ if os.path.exists(_gemv_so):
16
+ torch.ops.load_library(_gemv_so)
17
+ _gemv_ops = torch.ops._atlas_w4a16_gemv
18
+ else:
19
+ _gemv_ops = None
20
+
21
  FLOAT8_E4M3_MAX = 448.0
22
  FLOAT4_E2M1_MAX = 6.0
23
  SF_VEC_SIZE = 16
 
38
  sf_rowmajor: torch.Tensor | None = None
39
 
40
 
41
+ def _get_scale2(global_scale: torch.Tensor) -> float:
42
+ s2 = getattr(global_scale, "_atlas_scale2", None)
43
+ if s2 is None:
44
+ s2 = 1.0 / float(global_scale.item())
45
+ try:
46
+ setattr(global_scale, "_atlas_scale2", s2)
47
+ except Exception:
48
+ pass
49
+ return s2
50
+
51
+
52
  def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
53
  n, k = wd.shape
54
  g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
 
81
  qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
82
  sf_m, sf_n = swizzled_sf_shape(n, k)
83
  sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
84
+ gs_t = torch.tensor([gs], device=device, dtype=torch.float32)
85
+ try:
86
+ setattr(gs_t, "_atlas_scale2", scale2)
87
+ except Exception:
88
+ pass
89
  return PackedWeight(
90
  qweight=qweight,
91
  sf=sf,
92
+ global_scale=gs_t,
93
  n=n,
94
  k=k,
95
  swiglu=swiglu,
 
99
 
100
  def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
101
  x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
102
+ m_rows = x2.size(0)
103
+ scale2 = _get_scale2(pw.global_scale)
104
+ if _gemv_ops is not None:
105
+ if m_rows <= 8:
106
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
107
+ _gemv_ops.w4a16_gemv_sw(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
108
+ return out.reshape(*x.shape[:-1], pw.n)
109
+ w_bf16 = torch.empty((pw.n, pw.k), device=x.device, dtype=torch.bfloat16)
110
+ _gemv_ops.w4a16_dequant(pw.qweight, pw.sf_rowmajor, scale2, w_bf16)
111
+ return F.linear(x2, w_bf16).reshape(*x.shape[:-1], pw.n)
112
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
113
  ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
114
  return out.reshape(*x.shape[:-1], pw.n)
115
 
 
118
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
119
  scale2: float, C: torch.Tensor,
120
  ) -> None:
 
121
  ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
122
 
123
 
 
125
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
126
  scale2: float, C: torch.Tensor,
127
  ) -> None:
 
128
  ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
129
 
130
 
 
132
  B_packed: torch.Tensor, B_scale: torch.Tensor,
133
  scale2: float, B_fp8_out: torch.Tensor,
134
  ) -> None:
 
135
  ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
136
 
137
 
138
  def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
 
139
  ops.fp8_gemm_t(A, B_fp8, C)
140
 
141
 
142
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
 
143
  ops.bf16_to_fp8(src, dst)
144
 
145
 
build/torch212-cxx11-cu132-aarch64-linux/_atlas_w4a16_gemv.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:06039b9894e062e1bcc52210c2555513cdeb2de5da160dcf538166ff9f9df87f
3
+ size 1416920
build/torch212-cxx11-cu132-aarch64-linux/nvfp4_dense_gemm/__init__.py CHANGED
@@ -1,14 +1,23 @@
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
@@ -29,6 +38,17 @@ class PackedWeight:
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)
@@ -61,10 +81,15 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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,
@@ -74,8 +99,17 @@ def pack(w: torch.Tensor, device: str | torch.device = "cuda", swiglu: bool = Fa
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
 
@@ -84,7 +118,6 @@ 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
 
@@ -92,7 +125,6 @@ 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
 
@@ -100,17 +132,14 @@ 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
 
 
1
  # SPDX-License-Identifier: AGPL-3.0-only
2
+ """Atlas dense NVFP4/FP8 GEMM kernels for 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 os
9
  import torch
10
+ import torch.nn.functional as F
11
 
12
  from ._ops import ops
13
 
14
+ _gemv_so = os.path.join(os.path.dirname(__file__), "_atlas_w4a16_gemv.abi3.so")
15
+ if os.path.exists(_gemv_so):
16
+ torch.ops.load_library(_gemv_so)
17
+ _gemv_ops = torch.ops._atlas_w4a16_gemv
18
+ else:
19
+ _gemv_ops = None
20
+
21
  FLOAT8_E4M3_MAX = 448.0
22
  FLOAT4_E2M1_MAX = 6.0
23
  SF_VEC_SIZE = 16
 
38
  sf_rowmajor: torch.Tensor | None = None
39
 
40
 
41
+ def _get_scale2(global_scale: torch.Tensor) -> float:
42
+ s2 = getattr(global_scale, "_atlas_scale2", None)
43
+ if s2 is None:
44
+ s2 = 1.0 / float(global_scale.item())
45
+ try:
46
+ setattr(global_scale, "_atlas_scale2", s2)
47
+ except Exception:
48
+ pass
49
+ return s2
50
+
51
+
52
  def _pack_e2m1_torch(wd: torch.Tensor, scale2: float) -> tuple[torch.Tensor, torch.Tensor]:
53
  n, k = wd.shape
54
  g = wd.float().reshape(n, k // SF_VEC_SIZE, SF_VEC_SIZE)
 
81
  qweight, sf_rowmajor = _pack_e2m1_torch(wd, scale2)
82
  sf_m, sf_n = swizzled_sf_shape(n, k)
83
  sf = torch.zeros((sf_m, sf_n), device=device, dtype=torch.int32)
84
+ gs_t = torch.tensor([gs], device=device, dtype=torch.float32)
85
+ try:
86
+ setattr(gs_t, "_atlas_scale2", scale2)
87
+ except Exception:
88
+ pass
89
  return PackedWeight(
90
  qweight=qweight,
91
  sf=sf,
92
+ global_scale=gs_t,
93
  n=n,
94
  k=k,
95
  swiglu=swiglu,
 
99
 
100
  def gemm(pw: PackedWeight, x: torch.Tensor) -> torch.Tensor:
101
  x2 = x.reshape(-1, pw.k).to(torch.bfloat16).contiguous()
102
+ m_rows = x2.size(0)
103
+ scale2 = _get_scale2(pw.global_scale)
104
+ if _gemv_ops is not None:
105
+ if m_rows <= 8:
106
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
107
+ _gemv_ops.w4a16_gemv_sw(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
108
+ return out.reshape(*x.shape[:-1], pw.n)
109
+ w_bf16 = torch.empty((pw.n, pw.k), device=x.device, dtype=torch.bfloat16)
110
+ _gemv_ops.w4a16_dequant(pw.qweight, pw.sf_rowmajor, scale2, w_bf16)
111
+ return F.linear(x2, w_bf16).reshape(*x.shape[:-1], pw.n)
112
+ out = torch.empty((m_rows, pw.n), device=x.device, dtype=torch.bfloat16)
113
  ops.w4a16_gemm(x2, pw.qweight, pw.sf_rowmajor, scale2, out)
114
  return out.reshape(*x.shape[:-1], pw.n)
115
 
 
118
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
119
  scale2: float, C: torch.Tensor,
120
  ) -> None:
 
121
  ops.w4a16_gemm(A, B_packed, B_scale, scale2, C)
122
 
123
 
 
125
  A: torch.Tensor, B_packed: torch.Tensor, B_scale: torch.Tensor,
126
  scale2: float, C: torch.Tensor,
127
  ) -> None:
 
128
  ops.w4a16_gemm_t(A, B_packed, B_scale, scale2, C)
129
 
130
 
 
132
  B_packed: torch.Tensor, B_scale: torch.Tensor,
133
  scale2: float, B_fp8_out: torch.Tensor,
134
  ) -> None:
 
135
  ops.predequant_nvfp4_to_fp8(B_packed, B_scale, scale2, B_fp8_out)
136
 
137
 
138
  def fp8_gemm_t(A: torch.Tensor, B_fp8: torch.Tensor, C: torch.Tensor) -> None:
 
139
  ops.fp8_gemm_t(A, B_fp8, C)
140
 
141
 
142
  def bf16_to_fp8(src: torch.Tensor, dst: torch.Tensor) -> None:
 
143
  ops.bf16_to_fp8(src, dst)
144
 
145