Uploaded using `kernel-builder`.
Browse files- build/torch211-cxx11-cu128-x86_64-linux/__init__.py +99 -0
- build/torch211-cxx11-cu128-x86_64-linux/_grouped_moe_gemv_cuda_1683349.abi3.so +3 -0
- build/torch211-cxx11-cu128-x86_64-linux/_ops.py +9 -0
- build/torch211-cxx11-cu128-x86_64-linux/grouped_moe_gemv/__init__.py +26 -0
- build/torch211-cxx11-cu128-x86_64-linux/metadata.json +13 -0
- build/torch211-cxx11-cu130-x86_64-linux/__init__.py +99 -0
- build/torch211-cxx11-cu130-x86_64-linux/_grouped_moe_gemv_cuda_1683349.abi3.so +3 -0
- build/torch211-cxx11-cu130-x86_64-linux/_ops.py +9 -0
- build/torch211-cxx11-cu130-x86_64-linux/grouped_moe_gemv/__init__.py +26 -0
- build/torch211-cxx11-cu130-x86_64-linux/metadata.json +13 -0
- build/torch212-cxx11-cu130-x86_64-linux/__init__.py +99 -0
- build/torch212-cxx11-cu130-x86_64-linux/_grouped_moe_gemv_cuda_1683349.abi3.so +3 -0
- build/torch212-cxx11-cu130-x86_64-linux/_ops.py +9 -0
- build/torch212-cxx11-cu130-x86_64-linux/grouped_moe_gemv/__init__.py +26 -0
- build/torch212-cxx11-cu130-x86_64-linux/metadata.json +13 -0
- build/torch212-cxx11-cu132-x86_64-linux/__init__.py +99 -0
- build/torch212-cxx11-cu132-x86_64-linux/_grouped_moe_gemv_cuda_1683349.abi3.so +3 -0
- build/torch212-cxx11-cu132-x86_64-linux/_ops.py +9 -0
- build/torch212-cxx11-cu132-x86_64-linux/grouped_moe_gemv/__init__.py +26 -0
- build/torch212-cxx11-cu132-x86_64-linux/metadata.json +13 -0
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 |
+
}
|