liangsu9988 commited on
Commit
16ecd37
·
verified ·
1 Parent(s): aab237a

Uploaded using `kernel-builder`.

Browse files
build/torch211-cxx11-cu128-x86_64-linux/__init__.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """FlashRT grouped MoE GEMV kernels."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Optional
6
+
7
+ import torch
8
+
9
+ from ._ops import add_op_namespace_prefix, ops
10
+
11
+
12
+ @torch.library.register_fake(add_op_namespace_prefix("w4a16_decode_gemv_bf16"))
13
+ def _w4a16_decode_gemv_fake(
14
+ x_bf16: torch.Tensor,
15
+ weight_packed: torch.Tensor,
16
+ sfb: torch.Tensor,
17
+ alpha: float,
18
+ out: torch.Tensor,
19
+ ) -> None:
20
+ k = x_bf16.shape[0] if x_bf16.dim() == 1 else x_bf16.shape[1]
21
+ if weight_packed.dim() != 2 or weight_packed.shape[1] != k // 2 or out.shape != (weight_packed.shape[0],):
22
+ raise RuntimeError("expected x (K,) or (1,K), weight_packed (N,K/2), out (N,)")
23
+ return None
24
+
25
+
26
+ @torch.library.register_fake(add_op_namespace_prefix("grouped_w4a16_gemv_bf16"))
27
+ def _grouped_w4a16_gemv_fake(
28
+ activations: torch.Tensor,
29
+ weight_stack: torch.Tensor,
30
+ sfb_stack: torch.Tensor,
31
+ alpha_stack: torch.Tensor,
32
+ expert_idx: torch.Tensor,
33
+ w_stride: int,
34
+ sfb_stride: int,
35
+ out: torch.Tensor,
36
+ ) -> None:
37
+ if activations.dim() != 2 or out.dim() != 2 or out.shape[0] != activations.shape[0]:
38
+ raise RuntimeError("expected activations (slots,K), out (slots,N)")
39
+ if expert_idx.shape != (activations.shape[0],):
40
+ raise RuntimeError("expert_idx must have shape (slots,)")
41
+ return None
42
+
43
+
44
+ def w4a16_decode_gemv_bf16(
45
+ x_bf16: torch.Tensor,
46
+ weight_packed: torch.Tensor,
47
+ sfb: torch.Tensor,
48
+ *,
49
+ alpha: float = 1.0,
50
+ out: Optional[torch.Tensor] = None,
51
+ ) -> torch.Tensor:
52
+ if out is None:
53
+ out = torch.empty((weight_packed.shape[0],), device=x_bf16.device, dtype=torch.bfloat16)
54
+ ops.w4a16_decode_gemv_bf16(x_bf16, weight_packed, sfb, float(alpha), out)
55
+ return out
56
+
57
+
58
+ def grouped_w4a16_gemv_bf16(
59
+ activations: torch.Tensor,
60
+ weight_stack: torch.Tensor,
61
+ sfb_stack: torch.Tensor,
62
+ alpha_stack: torch.Tensor,
63
+ expert_idx: torch.Tensor,
64
+ *,
65
+ n: int,
66
+ w_stride: Optional[int] = None,
67
+ sfb_stride: Optional[int] = None,
68
+ out: Optional[torch.Tensor] = None,
69
+ ) -> torch.Tensor:
70
+ """Run one W4A16 GEMV per routed slot.
71
+
72
+ `weight_stack` is a flat expert stack. `w_stride` and `sfb_stride` are byte
73
+ strides between experts; by default `w_stride = n * K / 2`.
74
+ """
75
+
76
+ k = activations.shape[1]
77
+ if out is None:
78
+ out = torch.empty((activations.shape[0], int(n)), device=activations.device, dtype=torch.bfloat16)
79
+ if w_stride is None:
80
+ w_stride = int(n) * k // 2
81
+ if sfb_stride is None:
82
+ raise RuntimeError("sfb_stride must be provided because swizzled SF size is layout-dependent")
83
+ ops.grouped_w4a16_gemv_bf16(
84
+ activations,
85
+ weight_stack,
86
+ sfb_stack,
87
+ alpha_stack,
88
+ expert_idx,
89
+ int(w_stride),
90
+ int(sfb_stride),
91
+ out,
92
+ )
93
+ return out
94
+
95
+
96
+ __all__ = [
97
+ "grouped_w4a16_gemv_bf16",
98
+ "w4a16_decode_gemv_bf16",
99
+ ]
build/torch211-cxx11-cu128-x86_64-linux/_grouped_moe_gemv_cuda_1683349.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:e1393d4970efc36091e142b61f250e978caf3744de67fd93e872867c6dc1595f
3
+ size 173904
build/torch211-cxx11-cu128-x86_64-linux/_ops.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from . import _grouped_moe_gemv_cuda_1683349
3
+ ops = torch.ops._grouped_moe_gemv_cuda_1683349
4
+
5
+ def add_op_namespace_prefix(op_name: str):
6
+ """
7
+ Prefix op by namespace.
8
+ """
9
+ return f"_grouped_moe_gemv_cuda_1683349::{op_name}"
build/torch211-cxx11-cu128-x86_64-linux/grouped_moe_gemv/__init__.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import ctypes
2
+ import importlib.util
3
+ import sys
4
+ from pathlib import Path
5
+ from types import ModuleType
6
+
7
+
8
+ def _import_from_path(file_path: Path) -> ModuleType:
9
+ # We cannot use the module name as-is, after adding it to `sys.modules`,
10
+ # it would also be used for other imports. So, we make a module name that
11
+ # depends on the path for it to be unique using the hex-encoded hash of
12
+ # the path.
13
+ path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
14
+ module_name = path_hash
15
+ spec = importlib.util.spec_from_file_location(module_name, file_path)
16
+ if spec is None:
17
+ raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
18
+ module = importlib.util.module_from_spec(spec)
19
+ if module is None:
20
+ raise ImportError(f"Cannot load module {module_name} from spec")
21
+ sys.modules[module_name] = module
22
+ spec.loader.exec_module(module) # type: ignore
23
+ return module
24
+
25
+
26
+ globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
build/torch211-cxx11-cu128-x86_64-linux/metadata.json ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "grouped-moe-gemv",
3
+ "id": "_grouped_moe_gemv_cuda_1683349",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "python-depends": [],
7
+ "backend": {
8
+ "type": "cuda",
9
+ "archs": [
10
+ "12.0a"
11
+ ]
12
+ }
13
+ }
build/torch211-cxx11-cu130-x86_64-linux/__init__.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """FlashRT grouped MoE GEMV kernels."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Optional
6
+
7
+ import torch
8
+
9
+ from ._ops import add_op_namespace_prefix, ops
10
+
11
+
12
+ @torch.library.register_fake(add_op_namespace_prefix("w4a16_decode_gemv_bf16"))
13
+ def _w4a16_decode_gemv_fake(
14
+ x_bf16: torch.Tensor,
15
+ weight_packed: torch.Tensor,
16
+ sfb: torch.Tensor,
17
+ alpha: float,
18
+ out: torch.Tensor,
19
+ ) -> None:
20
+ k = x_bf16.shape[0] if x_bf16.dim() == 1 else x_bf16.shape[1]
21
+ if weight_packed.dim() != 2 or weight_packed.shape[1] != k // 2 or out.shape != (weight_packed.shape[0],):
22
+ raise RuntimeError("expected x (K,) or (1,K), weight_packed (N,K/2), out (N,)")
23
+ return None
24
+
25
+
26
+ @torch.library.register_fake(add_op_namespace_prefix("grouped_w4a16_gemv_bf16"))
27
+ def _grouped_w4a16_gemv_fake(
28
+ activations: torch.Tensor,
29
+ weight_stack: torch.Tensor,
30
+ sfb_stack: torch.Tensor,
31
+ alpha_stack: torch.Tensor,
32
+ expert_idx: torch.Tensor,
33
+ w_stride: int,
34
+ sfb_stride: int,
35
+ out: torch.Tensor,
36
+ ) -> None:
37
+ if activations.dim() != 2 or out.dim() != 2 or out.shape[0] != activations.shape[0]:
38
+ raise RuntimeError("expected activations (slots,K), out (slots,N)")
39
+ if expert_idx.shape != (activations.shape[0],):
40
+ raise RuntimeError("expert_idx must have shape (slots,)")
41
+ return None
42
+
43
+
44
+ def w4a16_decode_gemv_bf16(
45
+ x_bf16: torch.Tensor,
46
+ weight_packed: torch.Tensor,
47
+ sfb: torch.Tensor,
48
+ *,
49
+ alpha: float = 1.0,
50
+ out: Optional[torch.Tensor] = None,
51
+ ) -> torch.Tensor:
52
+ if out is None:
53
+ out = torch.empty((weight_packed.shape[0],), device=x_bf16.device, dtype=torch.bfloat16)
54
+ ops.w4a16_decode_gemv_bf16(x_bf16, weight_packed, sfb, float(alpha), out)
55
+ return out
56
+
57
+
58
+ def grouped_w4a16_gemv_bf16(
59
+ activations: torch.Tensor,
60
+ weight_stack: torch.Tensor,
61
+ sfb_stack: torch.Tensor,
62
+ alpha_stack: torch.Tensor,
63
+ expert_idx: torch.Tensor,
64
+ *,
65
+ n: int,
66
+ w_stride: Optional[int] = None,
67
+ sfb_stride: Optional[int] = None,
68
+ out: Optional[torch.Tensor] = None,
69
+ ) -> torch.Tensor:
70
+ """Run one W4A16 GEMV per routed slot.
71
+
72
+ `weight_stack` is a flat expert stack. `w_stride` and `sfb_stride` are byte
73
+ strides between experts; by default `w_stride = n * K / 2`.
74
+ """
75
+
76
+ k = activations.shape[1]
77
+ if out is None:
78
+ out = torch.empty((activations.shape[0], int(n)), device=activations.device, dtype=torch.bfloat16)
79
+ if w_stride is None:
80
+ w_stride = int(n) * k // 2
81
+ if sfb_stride is None:
82
+ raise RuntimeError("sfb_stride must be provided because swizzled SF size is layout-dependent")
83
+ ops.grouped_w4a16_gemv_bf16(
84
+ activations,
85
+ weight_stack,
86
+ sfb_stack,
87
+ alpha_stack,
88
+ expert_idx,
89
+ int(w_stride),
90
+ int(sfb_stride),
91
+ out,
92
+ )
93
+ return out
94
+
95
+
96
+ __all__ = [
97
+ "grouped_w4a16_gemv_bf16",
98
+ "w4a16_decode_gemv_bf16",
99
+ ]
build/torch211-cxx11-cu130-x86_64-linux/_grouped_moe_gemv_cuda_1683349.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d4d5f5b0bb26387a375b49bd17940ec6650e627d34434a448445b68f02cbc870
3
+ size 163912
build/torch211-cxx11-cu130-x86_64-linux/_ops.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from . import _grouped_moe_gemv_cuda_1683349
3
+ ops = torch.ops._grouped_moe_gemv_cuda_1683349
4
+
5
+ def add_op_namespace_prefix(op_name: str):
6
+ """
7
+ Prefix op by namespace.
8
+ """
9
+ return f"_grouped_moe_gemv_cuda_1683349::{op_name}"
build/torch211-cxx11-cu130-x86_64-linux/grouped_moe_gemv/__init__.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import ctypes
2
+ import importlib.util
3
+ import sys
4
+ from pathlib import Path
5
+ from types import ModuleType
6
+
7
+
8
+ def _import_from_path(file_path: Path) -> ModuleType:
9
+ # We cannot use the module name as-is, after adding it to `sys.modules`,
10
+ # it would also be used for other imports. So, we make a module name that
11
+ # depends on the path for it to be unique using the hex-encoded hash of
12
+ # the path.
13
+ path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
14
+ module_name = path_hash
15
+ spec = importlib.util.spec_from_file_location(module_name, file_path)
16
+ if spec is None:
17
+ raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
18
+ module = importlib.util.module_from_spec(spec)
19
+ if module is None:
20
+ raise ImportError(f"Cannot load module {module_name} from spec")
21
+ sys.modules[module_name] = module
22
+ spec.loader.exec_module(module) # type: ignore
23
+ return module
24
+
25
+
26
+ globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
build/torch211-cxx11-cu130-x86_64-linux/metadata.json ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "grouped-moe-gemv",
3
+ "id": "_grouped_moe_gemv_cuda_1683349",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "python-depends": [],
7
+ "backend": {
8
+ "type": "cuda",
9
+ "archs": [
10
+ "12.0a"
11
+ ]
12
+ }
13
+ }
build/torch212-cxx11-cu130-x86_64-linux/__init__.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """FlashRT grouped MoE GEMV kernels."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Optional
6
+
7
+ import torch
8
+
9
+ from ._ops import add_op_namespace_prefix, ops
10
+
11
+
12
+ @torch.library.register_fake(add_op_namespace_prefix("w4a16_decode_gemv_bf16"))
13
+ def _w4a16_decode_gemv_fake(
14
+ x_bf16: torch.Tensor,
15
+ weight_packed: torch.Tensor,
16
+ sfb: torch.Tensor,
17
+ alpha: float,
18
+ out: torch.Tensor,
19
+ ) -> None:
20
+ k = x_bf16.shape[0] if x_bf16.dim() == 1 else x_bf16.shape[1]
21
+ if weight_packed.dim() != 2 or weight_packed.shape[1] != k // 2 or out.shape != (weight_packed.shape[0],):
22
+ raise RuntimeError("expected x (K,) or (1,K), weight_packed (N,K/2), out (N,)")
23
+ return None
24
+
25
+
26
+ @torch.library.register_fake(add_op_namespace_prefix("grouped_w4a16_gemv_bf16"))
27
+ def _grouped_w4a16_gemv_fake(
28
+ activations: torch.Tensor,
29
+ weight_stack: torch.Tensor,
30
+ sfb_stack: torch.Tensor,
31
+ alpha_stack: torch.Tensor,
32
+ expert_idx: torch.Tensor,
33
+ w_stride: int,
34
+ sfb_stride: int,
35
+ out: torch.Tensor,
36
+ ) -> None:
37
+ if activations.dim() != 2 or out.dim() != 2 or out.shape[0] != activations.shape[0]:
38
+ raise RuntimeError("expected activations (slots,K), out (slots,N)")
39
+ if expert_idx.shape != (activations.shape[0],):
40
+ raise RuntimeError("expert_idx must have shape (slots,)")
41
+ return None
42
+
43
+
44
+ def w4a16_decode_gemv_bf16(
45
+ x_bf16: torch.Tensor,
46
+ weight_packed: torch.Tensor,
47
+ sfb: torch.Tensor,
48
+ *,
49
+ alpha: float = 1.0,
50
+ out: Optional[torch.Tensor] = None,
51
+ ) -> torch.Tensor:
52
+ if out is None:
53
+ out = torch.empty((weight_packed.shape[0],), device=x_bf16.device, dtype=torch.bfloat16)
54
+ ops.w4a16_decode_gemv_bf16(x_bf16, weight_packed, sfb, float(alpha), out)
55
+ return out
56
+
57
+
58
+ def grouped_w4a16_gemv_bf16(
59
+ activations: torch.Tensor,
60
+ weight_stack: torch.Tensor,
61
+ sfb_stack: torch.Tensor,
62
+ alpha_stack: torch.Tensor,
63
+ expert_idx: torch.Tensor,
64
+ *,
65
+ n: int,
66
+ w_stride: Optional[int] = None,
67
+ sfb_stride: Optional[int] = None,
68
+ out: Optional[torch.Tensor] = None,
69
+ ) -> torch.Tensor:
70
+ """Run one W4A16 GEMV per routed slot.
71
+
72
+ `weight_stack` is a flat expert stack. `w_stride` and `sfb_stride` are byte
73
+ strides between experts; by default `w_stride = n * K / 2`.
74
+ """
75
+
76
+ k = activations.shape[1]
77
+ if out is None:
78
+ out = torch.empty((activations.shape[0], int(n)), device=activations.device, dtype=torch.bfloat16)
79
+ if w_stride is None:
80
+ w_stride = int(n) * k // 2
81
+ if sfb_stride is None:
82
+ raise RuntimeError("sfb_stride must be provided because swizzled SF size is layout-dependent")
83
+ ops.grouped_w4a16_gemv_bf16(
84
+ activations,
85
+ weight_stack,
86
+ sfb_stack,
87
+ alpha_stack,
88
+ expert_idx,
89
+ int(w_stride),
90
+ int(sfb_stride),
91
+ out,
92
+ )
93
+ return out
94
+
95
+
96
+ __all__ = [
97
+ "grouped_w4a16_gemv_bf16",
98
+ "w4a16_decode_gemv_bf16",
99
+ ]
build/torch212-cxx11-cu130-x86_64-linux/_grouped_moe_gemv_cuda_1683349.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f7b735196f5e18c071afa5416f9eca8db4e2ef9f27189ffab025c52e323457fc
3
+ size 174448
build/torch212-cxx11-cu130-x86_64-linux/_ops.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from . import _grouped_moe_gemv_cuda_1683349
3
+ ops = torch.ops._grouped_moe_gemv_cuda_1683349
4
+
5
+ def add_op_namespace_prefix(op_name: str):
6
+ """
7
+ Prefix op by namespace.
8
+ """
9
+ return f"_grouped_moe_gemv_cuda_1683349::{op_name}"
build/torch212-cxx11-cu130-x86_64-linux/grouped_moe_gemv/__init__.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import ctypes
2
+ import importlib.util
3
+ import sys
4
+ from pathlib import Path
5
+ from types import ModuleType
6
+
7
+
8
+ def _import_from_path(file_path: Path) -> ModuleType:
9
+ # We cannot use the module name as-is, after adding it to `sys.modules`,
10
+ # it would also be used for other imports. So, we make a module name that
11
+ # depends on the path for it to be unique using the hex-encoded hash of
12
+ # the path.
13
+ path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
14
+ module_name = path_hash
15
+ spec = importlib.util.spec_from_file_location(module_name, file_path)
16
+ if spec is None:
17
+ raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
18
+ module = importlib.util.module_from_spec(spec)
19
+ if module is None:
20
+ raise ImportError(f"Cannot load module {module_name} from spec")
21
+ sys.modules[module_name] = module
22
+ spec.loader.exec_module(module) # type: ignore
23
+ return module
24
+
25
+
26
+ globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
build/torch212-cxx11-cu130-x86_64-linux/metadata.json ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "grouped-moe-gemv",
3
+ "id": "_grouped_moe_gemv_cuda_1683349",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "python-depends": [],
7
+ "backend": {
8
+ "type": "cuda",
9
+ "archs": [
10
+ "12.0a"
11
+ ]
12
+ }
13
+ }
build/torch212-cxx11-cu132-x86_64-linux/__init__.py ADDED
@@ -0,0 +1,99 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """FlashRT grouped MoE GEMV kernels."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Optional
6
+
7
+ import torch
8
+
9
+ from ._ops import add_op_namespace_prefix, ops
10
+
11
+
12
+ @torch.library.register_fake(add_op_namespace_prefix("w4a16_decode_gemv_bf16"))
13
+ def _w4a16_decode_gemv_fake(
14
+ x_bf16: torch.Tensor,
15
+ weight_packed: torch.Tensor,
16
+ sfb: torch.Tensor,
17
+ alpha: float,
18
+ out: torch.Tensor,
19
+ ) -> None:
20
+ k = x_bf16.shape[0] if x_bf16.dim() == 1 else x_bf16.shape[1]
21
+ if weight_packed.dim() != 2 or weight_packed.shape[1] != k // 2 or out.shape != (weight_packed.shape[0],):
22
+ raise RuntimeError("expected x (K,) or (1,K), weight_packed (N,K/2), out (N,)")
23
+ return None
24
+
25
+
26
+ @torch.library.register_fake(add_op_namespace_prefix("grouped_w4a16_gemv_bf16"))
27
+ def _grouped_w4a16_gemv_fake(
28
+ activations: torch.Tensor,
29
+ weight_stack: torch.Tensor,
30
+ sfb_stack: torch.Tensor,
31
+ alpha_stack: torch.Tensor,
32
+ expert_idx: torch.Tensor,
33
+ w_stride: int,
34
+ sfb_stride: int,
35
+ out: torch.Tensor,
36
+ ) -> None:
37
+ if activations.dim() != 2 or out.dim() != 2 or out.shape[0] != activations.shape[0]:
38
+ raise RuntimeError("expected activations (slots,K), out (slots,N)")
39
+ if expert_idx.shape != (activations.shape[0],):
40
+ raise RuntimeError("expert_idx must have shape (slots,)")
41
+ return None
42
+
43
+
44
+ def w4a16_decode_gemv_bf16(
45
+ x_bf16: torch.Tensor,
46
+ weight_packed: torch.Tensor,
47
+ sfb: torch.Tensor,
48
+ *,
49
+ alpha: float = 1.0,
50
+ out: Optional[torch.Tensor] = None,
51
+ ) -> torch.Tensor:
52
+ if out is None:
53
+ out = torch.empty((weight_packed.shape[0],), device=x_bf16.device, dtype=torch.bfloat16)
54
+ ops.w4a16_decode_gemv_bf16(x_bf16, weight_packed, sfb, float(alpha), out)
55
+ return out
56
+
57
+
58
+ def grouped_w4a16_gemv_bf16(
59
+ activations: torch.Tensor,
60
+ weight_stack: torch.Tensor,
61
+ sfb_stack: torch.Tensor,
62
+ alpha_stack: torch.Tensor,
63
+ expert_idx: torch.Tensor,
64
+ *,
65
+ n: int,
66
+ w_stride: Optional[int] = None,
67
+ sfb_stride: Optional[int] = None,
68
+ out: Optional[torch.Tensor] = None,
69
+ ) -> torch.Tensor:
70
+ """Run one W4A16 GEMV per routed slot.
71
+
72
+ `weight_stack` is a flat expert stack. `w_stride` and `sfb_stride` are byte
73
+ strides between experts; by default `w_stride = n * K / 2`.
74
+ """
75
+
76
+ k = activations.shape[1]
77
+ if out is None:
78
+ out = torch.empty((activations.shape[0], int(n)), device=activations.device, dtype=torch.bfloat16)
79
+ if w_stride is None:
80
+ w_stride = int(n) * k // 2
81
+ if sfb_stride is None:
82
+ raise RuntimeError("sfb_stride must be provided because swizzled SF size is layout-dependent")
83
+ ops.grouped_w4a16_gemv_bf16(
84
+ activations,
85
+ weight_stack,
86
+ sfb_stack,
87
+ alpha_stack,
88
+ expert_idx,
89
+ int(w_stride),
90
+ int(sfb_stride),
91
+ out,
92
+ )
93
+ return out
94
+
95
+
96
+ __all__ = [
97
+ "grouped_w4a16_gemv_bf16",
98
+ "w4a16_decode_gemv_bf16",
99
+ ]
build/torch212-cxx11-cu132-x86_64-linux/_grouped_moe_gemv_cuda_1683349.abi3.so ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d1e7570b56c4dddb95cd4c18015c02517beedc8eee8678fced97250cdc9bb2cc
3
+ size 174448
build/torch212-cxx11-cu132-x86_64-linux/_ops.py ADDED
@@ -0,0 +1,9 @@
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from . import _grouped_moe_gemv_cuda_1683349
3
+ ops = torch.ops._grouped_moe_gemv_cuda_1683349
4
+
5
+ def add_op_namespace_prefix(op_name: str):
6
+ """
7
+ Prefix op by namespace.
8
+ """
9
+ return f"_grouped_moe_gemv_cuda_1683349::{op_name}"
build/torch212-cxx11-cu132-x86_64-linux/grouped_moe_gemv/__init__.py ADDED
@@ -0,0 +1,26 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import ctypes
2
+ import importlib.util
3
+ import sys
4
+ from pathlib import Path
5
+ from types import ModuleType
6
+
7
+
8
+ def _import_from_path(file_path: Path) -> ModuleType:
9
+ # We cannot use the module name as-is, after adding it to `sys.modules`,
10
+ # it would also be used for other imports. So, we make a module name that
11
+ # depends on the path for it to be unique using the hex-encoded hash of
12
+ # the path.
13
+ path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
14
+ module_name = path_hash
15
+ spec = importlib.util.spec_from_file_location(module_name, file_path)
16
+ if spec is None:
17
+ raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
18
+ module = importlib.util.module_from_spec(spec)
19
+ if module is None:
20
+ raise ImportError(f"Cannot load module {module_name} from spec")
21
+ sys.modules[module_name] = module
22
+ spec.loader.exec_module(module) # type: ignore
23
+ return module
24
+
25
+
26
+ globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))
build/torch212-cxx11-cu132-x86_64-linux/metadata.json ADDED
@@ -0,0 +1,13 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "name": "grouped-moe-gemv",
3
+ "id": "_grouped_moe_gemv_cuda_1683349",
4
+ "version": 1,
5
+ "license": "Apache-2.0",
6
+ "python-depends": [],
7
+ "backend": {
8
+ "type": "cuda",
9
+ "archs": [
10
+ "12.0a"
11
+ ]
12
+ }
13
+ }