Remove QServe and FBGEMM FP8 quantization (#31109)
This commit is contained in:
Brayden Zhong
@@ -332,7 +332,7 @@ Please consult the documentation below and [server_args.py](https://github.com/s
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--quantization`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>The quantization method.</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`None`</td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>awq</code>, <code>fp8</code>, <code>mxfp8</code>, <code>gptq</code>, <code>marlin</code>, <code>gptq_marlin</code>, <code>awq_marlin</code>, <code>bitsandbytes</code>, <code>gguf</code>, <code>modelopt</code>, <code>modelopt_fp8</code>, <code>modelopt_fp4</code>, <code>nvfp4_online</code>, <code>modelopt_mixed</code>, <code>petit_nvfp4</code>, <code>w8a8_int8</code>, <code>w8a8_fp8</code>, <code>moe_wna16</code>, <code>qoq</code>, <code>w4afp8</code>, <code>mxfp4</code>, <code>auto-round</code>, <code>compressed-tensors</code>, <code>modelslim</code>, <code>quark</code>, <code>quark_int4fp8_moe</code>, <code>quark_mxfp4</code>, <code>mlx_q4</code>, <code>mlx_q8</code>, <code>unquant</code></td>
|
||||
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>awq</code>, <code>fp8</code>, <code>mxfp8</code>, <code>gptq</code>, <code>marlin</code>, <code>gptq_marlin</code>, <code>awq_marlin</code>, <code>bitsandbytes</code>, <code>gguf</code>, <code>modelopt</code>, <code>modelopt_fp8</code>, <code>modelopt_fp4</code>, <code>nvfp4_online</code>, <code>modelopt_mixed</code>, <code>petit_nvfp4</code>, <code>w8a8_int8</code>, <code>w8a8_fp8</code>, <code>moe_wna16</code>, <code>w4afp8</code>, <code>mxfp4</code>, <code>auto-round</code>, <code>compressed-tensors</code>, <code>modelslim</code>, <code>quark</code>, <code>quark_int4fp8_moe</code>, <code>quark_mxfp4</code>, <code>mlx_q4</code>, <code>mlx_q8</code>, <code>unquant</code></td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--quantization-param-path`</td>
|
||||
|
||||
@@ -53,7 +53,6 @@ WEIGHT_LOADER_V2_SUPPORTED = [
|
||||
"GPTQMarlin24LinearMethod",
|
||||
"TPUInt8LinearMethod",
|
||||
"GPTQLinearMethod",
|
||||
"FBGEMMFp8LinearMethod",
|
||||
"ModelOptFp8LinearMethod",
|
||||
"ModelOptFp4LinearMethod",
|
||||
"ComfyUIFp4LinearMethod",
|
||||
|
||||
@@ -1329,7 +1329,6 @@ class ModelConfig:
|
||||
"fp8",
|
||||
"compressed_tensors",
|
||||
"compressed-tensors",
|
||||
"fbgemm_fp8",
|
||||
"w8a8_fp8",
|
||||
"petit_nvfp4",
|
||||
"quark",
|
||||
@@ -1350,14 +1349,12 @@ class ModelConfig:
|
||||
"gptq_marlin_24",
|
||||
"gptq_marlin",
|
||||
"awq_marlin",
|
||||
"fbgemm_fp8",
|
||||
"compressed_tensors",
|
||||
"compressed-tensors",
|
||||
"experts_int8",
|
||||
"w8a8_int8",
|
||||
"w8a8_fp8",
|
||||
"moe_wna16",
|
||||
"qoq",
|
||||
"w4afp8",
|
||||
"petit_nvfp4",
|
||||
"quark",
|
||||
|
||||
@@ -66,7 +66,6 @@ WEIGHT_LOADER_V2_SUPPORTED = [
|
||||
"GPTQMarlin24LinearMethod",
|
||||
"TPUInt8LinearMethod",
|
||||
"GPTQLinearMethod",
|
||||
"FBGEMMFp8LinearMethod",
|
||||
"GPTQLinearAscendMethod",
|
||||
"GPTQLinearIntelAMXMethod",
|
||||
"GPTQMoEAscendMethod",
|
||||
|
||||
@@ -27,7 +27,6 @@ from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import
|
||||
CompressedTensorsConfig,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8 import Fp8Config
|
||||
from sglang.srt.layers.quantization.fpgemm_fp8 import FBGEMMFp8Config
|
||||
from sglang.srt.layers.quantization.gguf import GGUFConfig
|
||||
from sglang.srt.layers.quantization.gptq import (
|
||||
CPUGPTQConfig,
|
||||
@@ -49,7 +48,6 @@ from sglang.srt.layers.quantization.npu_mxfp4 import Mxfp4W4A8Config
|
||||
from sglang.srt.layers.quantization.npu_mxfp4_w4a4 import Mxfp4W4A4Config
|
||||
from sglang.srt.layers.quantization.nvfp4_online import NvFp4OnlineConfig
|
||||
from sglang.srt.layers.quantization.petit import PetitNvFp4Config
|
||||
from sglang.srt.layers.quantization.qoq import QoQConfig
|
||||
from sglang.srt.layers.quantization.quark.quark import QuarkConfig
|
||||
from sglang.srt.layers.quantization.quark_int4fp8_moe import QuarkInt4Fp8Config
|
||||
from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config
|
||||
@@ -91,10 +89,8 @@ BASE_QUANTIZATION_METHODS: Dict[str, Type[QuantizationConfig]] = {
|
||||
"gptq_marlin": GPTQMarlinConfig,
|
||||
"moe_wna16": MoeWNA16Config,
|
||||
"compressed-tensors": CompressedTensorsConfig,
|
||||
"qoq": QoQConfig,
|
||||
"w4afp8": W4AFp8Config,
|
||||
"petit_nvfp4": PetitNvFp4Config,
|
||||
"fbgemm_fp8": FBGEMMFp8Config,
|
||||
"quark": QuarkConfig,
|
||||
"quark_mxfp4": QuarkConfig,
|
||||
"auto-round": AutoRoundConfig,
|
||||
|
||||
@@ -1,202 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any, List, Optional
|
||||
|
||||
import torch
|
||||
from torch.nn import Module
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
|
||||
from sglang.srt.layers.linear import LinearBase
|
||||
from sglang.srt.layers.parameter import ChannelQuantScaleParameter, ModelWeightParameter
|
||||
from sglang.srt.layers.quantization.base_config import (
|
||||
LinearMethodBase,
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
apply_fp8_linear,
|
||||
can_auto_enable_marlin_fp8,
|
||||
cutlass_fp8_supported,
|
||||
normalize_e4m3fn_to_e4m3fnuz,
|
||||
)
|
||||
from sglang.srt.layers.quantization.marlin_utils_fp8 import (
|
||||
apply_fp8_marlin_linear,
|
||||
prepare_fp8_layer_for_marlin,
|
||||
)
|
||||
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
|
||||
from sglang.srt.layers.quantization.utils import is_layer_skipped
|
||||
from sglang.srt.utils import get_bool_env_var, is_cuda
|
||||
|
||||
_is_cuda = is_cuda()
|
||||
_is_fp8_fnuz = is_fp8_fnuz()
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class FBGEMMFp8Config(QuantizationConfig):
|
||||
"""Config class for FBGEMM Fp8."""
|
||||
|
||||
def __init__(self, ignore_list: list[str], input_scale_ub: float):
|
||||
super().__init__()
|
||||
self.ignore_list = ignore_list if ignore_list else []
|
||||
self.input_scale_ub = input_scale_ub
|
||||
|
||||
# For GPUs that lack FP8 hardware suspport, we can leverage the Marlin
|
||||
# kernel for fast weight-only FP8 quantization
|
||||
# self.use_marlin = not marlin_fp8_supported()
|
||||
self.use_marlin = False
|
||||
if _is_cuda:
|
||||
force_marlin = get_bool_env_var("SGLANG_FORCE_FP8_MARLIN")
|
||||
auto_enable = can_auto_enable_marlin_fp8()
|
||||
self.use_marlin = force_marlin or auto_enable
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> str:
|
||||
return "fbgemm_fp8"
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> list[torch.dtype]:
|
||||
return [torch.bfloat16, torch.float16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 80
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> list[str]:
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: dict[str, Any]) -> FBGEMMFp8Config:
|
||||
ignore_list = cls.get_from_keys(config, ["modules_to_not_convert"])
|
||||
input_scale_ub = cls.get_from_keys(config, ["activation_scale_ub"])
|
||||
return cls(ignore_list=ignore_list, input_scale_ub=input_scale_ub)
|
||||
|
||||
def get_quant_method(
|
||||
self, layer: torch.nn.Module, prefix: str
|
||||
) -> Optional[QuantizeMethodBase]:
|
||||
if isinstance(layer, LinearBase):
|
||||
if is_layer_skipped(
|
||||
prefix=prefix,
|
||||
ignored_layers=self.ignore_list,
|
||||
fused_mapping=self.packed_modules_mapping,
|
||||
):
|
||||
return UnquantizedLinearMethod()
|
||||
return FBGEMMFp8LinearMethod(self)
|
||||
return None
|
||||
|
||||
def get_scaled_act_names(self) -> List[str]:
|
||||
return []
|
||||
|
||||
|
||||
class FBGEMMFp8LinearMethod(LinearMethodBase):
|
||||
|
||||
def __init__(self, quant_config: FBGEMMFp8Config):
|
||||
self.quant_config = quant_config
|
||||
# self.fp8_linear = Fp8LinearOp(
|
||||
# act_quant_static=False, act_quant_group_shape=GroupShape.PER_TOKEN)
|
||||
self.out_dtype = torch.get_default_dtype()
|
||||
self.cutlass_fp8_supported = cutlass_fp8_supported()
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: list[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
# maybe_create_device_identity()
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
del input_size, output_size
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
|
||||
layer.logical_widths = output_partition_sizes
|
||||
|
||||
layer.input_size_per_partition = input_size_per_partition
|
||||
layer.output_size_per_partition = output_size_per_partition
|
||||
layer.orig_dtype = params_dtype
|
||||
|
||||
# WEIGHT
|
||||
weight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition,
|
||||
dtype=torch.float8_e4m3fn,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("weight", weight)
|
||||
|
||||
# WEIGHT SCALE
|
||||
weight_scale = ChannelQuantScaleParameter(
|
||||
data=torch.empty((sum(output_partition_sizes), 1), dtype=torch.float32),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
weight_scale[:] = torch.finfo(torch.float32).min
|
||||
layer.register_parameter("weight_scale", weight_scale)
|
||||
|
||||
# INPUT SCALE UPPER BOUND
|
||||
input_scale_ub = torch.nn.Parameter(
|
||||
torch.tensor((self.quant_config.input_scale_ub), dtype=torch.float32),
|
||||
requires_grad=False,
|
||||
)
|
||||
layer.input_scale_ub = input_scale_ub
|
||||
|
||||
def process_weights_after_loading(self, layer: Module) -> None:
|
||||
# required by torch.compile
|
||||
layer.weight_scale = Parameter(layer.weight_scale.data, requires_grad=False)
|
||||
layer.weight = Parameter(layer.weight.data, requires_grad=False)
|
||||
|
||||
weight = layer.weight
|
||||
|
||||
if _is_fp8_fnuz:
|
||||
weight, weight_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
|
||||
weight=weight, weight_scale=layer.weight_scale, input_scale=None
|
||||
)
|
||||
if input_scale is not None:
|
||||
layer.input_scale = Parameter(input_scale, requires_grad=False)
|
||||
layer.weight_scale = Parameter(weight_scale, requires_grad=False)
|
||||
|
||||
layer.weight = Parameter(weight.t(), requires_grad=False)
|
||||
if self.quant_config.use_marlin:
|
||||
prepare_fp8_layer_for_marlin(layer)
|
||||
# Activations not quantized for marlin.
|
||||
del layer.input_scale_ub
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
|
||||
if self.quant_config.use_marlin:
|
||||
return apply_fp8_marlin_linear(
|
||||
input=x,
|
||||
weight=layer.weight,
|
||||
weight_scale=layer.weight_scale,
|
||||
workspace=layer.workspace,
|
||||
size_n=layer.output_size_per_partition,
|
||||
size_k=layer.input_size_per_partition,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
return apply_fp8_linear(
|
||||
input=x,
|
||||
weight=layer.weight,
|
||||
weight_scale=layer.weight_scale,
|
||||
input_scale=None,
|
||||
input_scale_ub=layer.input_scale_ub,
|
||||
bias=bias,
|
||||
cutlass_fp8_supported=self.cutlass_fp8_supported,
|
||||
use_per_token_if_dynamic=False,
|
||||
)
|
||||
@@ -1,245 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from sglang.kernels.ops.quantization.int8_kernel import per_token_quant_int8
|
||||
from sglang.srt.layers.parameter import (
|
||||
ChannelQuantScaleParameter,
|
||||
GroupQuantScaleParameter,
|
||||
ModelWeightParameter,
|
||||
)
|
||||
from sglang.srt.layers.quantization.base_config import (
|
||||
LinearMethodBase,
|
||||
QuantizationConfig,
|
||||
QuantizeMethodBase,
|
||||
)
|
||||
from sglang.srt.utils import is_cuda
|
||||
|
||||
_is_cuda = is_cuda()
|
||||
if _is_cuda:
|
||||
from sgl_kernel import qserve_w4a8_per_chn_gemm, qserve_w4a8_per_group_gemm
|
||||
|
||||
|
||||
QoQ_SUPPORTED_WEIGHT_BITS = [4]
|
||||
QoQ_SUPPORTED_GROUP_SIZES = [-1, 128]
|
||||
|
||||
|
||||
class QoQConfig(QuantizationConfig):
|
||||
"""Config class for QoQ Quantization.
|
||||
|
||||
- Weight: static, per-channel/group, asymmetric
|
||||
- Activation: dynamic, per-token, symmetric
|
||||
|
||||
Reference: https://arxiv.org/abs/2405.04532
|
||||
https://github.com/mit-han-lab/omniserve
|
||||
"""
|
||||
|
||||
def __init__(self, weight_bits: int, group_size: int) -> None:
|
||||
self.weight_bits = weight_bits
|
||||
self.group_size = group_size
|
||||
|
||||
# Verify
|
||||
if self.weight_bits not in QoQ_SUPPORTED_WEIGHT_BITS:
|
||||
raise ValueError(
|
||||
f"QoQ does not support weight_bits = {self.weight_bits}. "
|
||||
f"Only weight_bits = {QoQ_SUPPORTED_WEIGHT_BITS} "
|
||||
"are supported."
|
||||
)
|
||||
if self.group_size not in QoQ_SUPPORTED_GROUP_SIZES:
|
||||
raise ValueError(
|
||||
f"QoQ does not support group_size = {self.group_size}. "
|
||||
f"Only group_sizes = {QoQ_SUPPORTED_GROUP_SIZES} "
|
||||
"are supported."
|
||||
)
|
||||
|
||||
# 4 bits packed into 8 bit datatype.
|
||||
self.pack_factor = 8 // self.weight_bits
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return "QoQConfig(weight_bits={}, group_size={})".format(
|
||||
self.weight_bits, self.group_size
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_supported_act_dtypes(cls) -> List[torch.dtype]:
|
||||
return [torch.float16]
|
||||
|
||||
@classmethod
|
||||
def get_min_capability(cls) -> int:
|
||||
return 80
|
||||
|
||||
@classmethod
|
||||
def get_name(cls) -> str:
|
||||
return "qoq"
|
||||
|
||||
@classmethod
|
||||
def get_config_filenames(cls) -> List[str]:
|
||||
"""List of filenames to search for in the model directory."""
|
||||
return [
|
||||
"quant_config.json",
|
||||
"quantize_config.json",
|
||||
]
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict[str, Any]) -> QoQConfig:
|
||||
weight_bits = cls.get_from_keys(config, ["wbits"])
|
||||
group_size = cls.get_from_keys(config, ["group_size"])
|
||||
return cls(weight_bits, group_size)
|
||||
|
||||
def get_quant_method(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
prefix: str,
|
||||
) -> Optional[QuantizeMethodBase]:
|
||||
from sglang.srt.layers.linear import LinearBase
|
||||
|
||||
if isinstance(layer, LinearBase):
|
||||
return QoQLinearMethod(self)
|
||||
return None
|
||||
|
||||
def get_scaled_act_names(self) -> List[str]:
|
||||
return []
|
||||
|
||||
|
||||
class QoQLinearMethod(LinearMethodBase):
|
||||
"""Linear method for QoQ.
|
||||
|
||||
Args:
|
||||
quant_config: The QoQ quantization config.
|
||||
"""
|
||||
|
||||
def __init__(self, quant_config: QoQConfig):
|
||||
self.quant_config = quant_config
|
||||
|
||||
def create_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
input_size_per_partition: int,
|
||||
output_partition_sizes: List[int],
|
||||
input_size: int,
|
||||
output_size: int,
|
||||
params_dtype: torch.dtype,
|
||||
**extra_weight_attrs,
|
||||
):
|
||||
|
||||
weight_loader = extra_weight_attrs.get("weight_loader")
|
||||
|
||||
# Validate output_size_per_partition
|
||||
output_size_per_partition = sum(output_partition_sizes)
|
||||
if output_size_per_partition % 32 != 0:
|
||||
raise ValueError(
|
||||
f"Weight output_size_per_partition = "
|
||||
f"{output_size_per_partition} is not divisible by 32."
|
||||
)
|
||||
|
||||
# Validate input_size_per_partition
|
||||
if input_size_per_partition % self.quant_config.pack_factor != 0:
|
||||
raise ValueError(
|
||||
f"Weight input_size_per_partition = "
|
||||
f"{input_size_per_partition} is not divisible by "
|
||||
f"pack_factor = {self.quant_config.pack_factor}."
|
||||
)
|
||||
if (
|
||||
self.quant_config.group_size != -1
|
||||
and input_size_per_partition % self.quant_config.group_size != 0
|
||||
):
|
||||
raise ValueError(
|
||||
f"Weight input_size_per_partition = "
|
||||
f"{input_size_per_partition} is not divisible by "
|
||||
f"group_size = {self.quant_config.group_size}."
|
||||
)
|
||||
|
||||
qweight = ModelWeightParameter(
|
||||
data=torch.empty(
|
||||
output_size_per_partition,
|
||||
input_size_per_partition // self.quant_config.pack_factor,
|
||||
dtype=torch.int8,
|
||||
),
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("qweight", qweight)
|
||||
|
||||
s1_scales = ChannelQuantScaleParameter(
|
||||
data=torch.empty(output_size_per_partition, dtype=torch.float16),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("s1_scales", s1_scales)
|
||||
|
||||
if self.quant_config.group_size == -1:
|
||||
s1_szeros = ChannelQuantScaleParameter(
|
||||
data=torch.empty(output_size_per_partition, dtype=torch.float16),
|
||||
output_dim=0,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("s1_szeros", s1_szeros)
|
||||
else:
|
||||
s2_scales = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
(
|
||||
input_size_per_partition // self.quant_config.group_size,
|
||||
output_size_per_partition,
|
||||
),
|
||||
dtype=torch.int8,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("s2_scales", s2_scales)
|
||||
|
||||
s2_zeros = GroupQuantScaleParameter(
|
||||
data=torch.empty(
|
||||
(
|
||||
input_size_per_partition // self.quant_config.group_size,
|
||||
output_size_per_partition,
|
||||
),
|
||||
dtype=torch.int8,
|
||||
),
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
weight_loader=weight_loader,
|
||||
)
|
||||
layer.register_parameter("s2_zeros", s2_zeros)
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
layer.qweight = Parameter(layer.qweight.data, requires_grad=False)
|
||||
layer.s1_scales = Parameter(layer.s1_scales.data, requires_grad=False)
|
||||
if self.quant_config.group_size == -1:
|
||||
layer.s1_szeros = Parameter(layer.s1_szeros.data, requires_grad=False)
|
||||
else:
|
||||
layer.s2_scales = Parameter(layer.s2_scales.data, requires_grad=False)
|
||||
layer.s2_zeros = Parameter(layer.s2_zeros.data, requires_grad=False)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
):
|
||||
assert x.dtype == torch.float16, "QoQ only supports float16 input now"
|
||||
if self.quant_config.group_size == -1:
|
||||
x_q, x_scale, x_sum = per_token_quant_int8(
|
||||
x, scale_dtype=x.dtype, cal_sum=True
|
||||
)
|
||||
out = qserve_w4a8_per_chn_gemm(
|
||||
x_q, layer.qweight, layer.s1_scales, x_scale, layer.s1_szeros, x_sum
|
||||
)
|
||||
else:
|
||||
x_q, x_scale = per_token_quant_int8(x, scale_dtype=x.dtype)
|
||||
out = qserve_w4a8_per_group_gemm(
|
||||
x_q,
|
||||
layer.qweight,
|
||||
layer.s2_zeros,
|
||||
layer.s2_scales,
|
||||
layer.s1_scales,
|
||||
x_scale,
|
||||
)
|
||||
if bias is not None:
|
||||
out = out + bias
|
||||
return out
|
||||
@@ -169,7 +169,6 @@ QUANTIZATION_CHOICES = [
|
||||
"w8a8_int8", # mentioned in quantization.md documentation, supporting compressed-tensors quant_method.
|
||||
"w8a8_fp8", # mentioned in quantization.md documentation, supporting compressed-tensors quant_method.
|
||||
"moe_wna16", # custom loading logic for gptq/awq checkpoints (likely untested/unused)
|
||||
"qoq",
|
||||
"w4afp8",
|
||||
"mxfp4", # MOE-only.
|
||||
"auto-round",
|
||||
|
||||
@@ -270,8 +270,6 @@ set(SOURCES
|
||||
"csrc/gemm/per_token_group_quant_8bit.cu"
|
||||
"csrc/gemm/per_token_group_quant_8bit_v2.cu"
|
||||
"csrc/gemm/per_token_quant_fp8.cu"
|
||||
"csrc/gemm/qserve_w4a8_per_chn_gemm.cu"
|
||||
"csrc/gemm/qserve_w4a8_per_group_gemm.cu"
|
||||
"csrc/gemm/gptq/gptq_kernel.cu"
|
||||
"csrc/grammar/apply_token_bitmask_inplace_cuda.cu"
|
||||
|
||||
|
||||
@@ -1,214 +0,0 @@
|
||||
import argparse
|
||||
import copy
|
||||
import itertools
|
||||
import os
|
||||
|
||||
import torch
|
||||
import triton
|
||||
from sgl_kernel import (
|
||||
int8_scaled_mm,
|
||||
qserve_w4a8_per_chn_gemm,
|
||||
qserve_w4a8_per_group_gemm,
|
||||
)
|
||||
|
||||
from sglang.utils import is_in_ci
|
||||
|
||||
IS_CI = is_in_ci()
|
||||
|
||||
|
||||
def to_int8(tensor: torch.Tensor) -> torch.Tensor:
|
||||
return torch.round(tensor.clamp(min=-128, max=127)).to(dtype=torch.int8)
|
||||
|
||||
|
||||
WEIGHT_SHAPES = {
|
||||
"meta-llama/Llama-3.1-8B-Instruct": [
|
||||
([4096, 6144], 1),
|
||||
([4096, 4096], 0),
|
||||
([4096, 28672], 1),
|
||||
([14336, 4096], 0),
|
||||
],
|
||||
"meta-llama/Llama-3.3-70B-Instruct": [
|
||||
([8192, 10240], 1),
|
||||
([8192, 8192], 0),
|
||||
([8192, 57344], 1),
|
||||
([28672, 8192], 0),
|
||||
],
|
||||
"mistralai/Mistral-Large-Instruct-2407": [
|
||||
([12288, 14336], 1),
|
||||
([12288, 12288], 0),
|
||||
([12288, 57344], 1),
|
||||
([28672, 12288], 0),
|
||||
],
|
||||
"Qwen/Qwen2.5-7B-Instruct": [
|
||||
([3584, 4608], 1),
|
||||
([3584, 3584], 0),
|
||||
([3584, 37888], 1),
|
||||
([18944, 3584], 0),
|
||||
],
|
||||
"Qwen/Qwen2.5-32B-Instruct": [
|
||||
([5120, 7168], 1),
|
||||
([5120, 5120], 0),
|
||||
([5120, 55296], 1),
|
||||
([27648, 5120], 0),
|
||||
],
|
||||
"Qwen/Qwen2.5-72B-Instruct": [
|
||||
([8192, 10240], 1),
|
||||
([8192, 8192], 0),
|
||||
([8192, 59136], 1),
|
||||
([29568, 8192], 0),
|
||||
],
|
||||
"deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct": [
|
||||
([2048, 3072], 1),
|
||||
([2048, 4096], 1),
|
||||
([2048, 2048], 0),
|
||||
([2048, 576], 0),
|
||||
([2048, 21888], 1),
|
||||
([10944, 2048], 0),
|
||||
([2048, 2816], 1),
|
||||
([1408, 2048], 0),
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
# CI environment uses simplified parameters
|
||||
if IS_CI:
|
||||
batch_sizes = [1, 16] # Simplified for CI
|
||||
else:
|
||||
batch_sizes = [1, 16, 32, 64, 128, 256, 512, 1024, 2048]
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["batch_size"],
|
||||
x_vals=batch_sizes,
|
||||
x_log=False,
|
||||
line_arg="provider",
|
||||
line_vals=["FP16", "W8A8", "Qserve_W4A8_Per_Channel", "Qserve_W4A8_Per_Group"],
|
||||
line_names=["FP16", "W8A8", "Qserve_W4A8_Per_Channel", "Qserve_W4A8_Per_Group"],
|
||||
styles=[("blue", "-"), ("orange", "-"), ("green", "-"), ("red", "-")],
|
||||
ylabel="ms",
|
||||
plot_name="FP16_vs_W8A8_vs_Qserve_W4A8_GEMM",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(batch_size, provider, N, K):
|
||||
M = batch_size
|
||||
# For W8A8
|
||||
a = to_int8(torch.randn((M, K), device="cuda") * 5)
|
||||
b = to_int8(torch.randn((N, K), device="cuda").t() * 5)
|
||||
a_fp16 = a.to(torch.float16)
|
||||
b_fp16 = b.to(torch.float16)
|
||||
scale_a = torch.randn((M,), device="cuda", dtype=torch.float32)
|
||||
scale_b = torch.randn((N,), device="cuda", dtype=torch.float32)
|
||||
|
||||
# For Qserve W4A8 per channel
|
||||
a_qserve_chn = a
|
||||
# two int4s pack into one int8
|
||||
b_qserve_chn = to_int8(torch.randn((N, K // 2), device="cuda") * 5)
|
||||
# b_qserve_chn = b.t().contiguous()
|
||||
scale_a_qserve_chn = scale_a.to(torch.float16)
|
||||
scale_b_qserve_chn = scale_b.to(torch.float16)
|
||||
szero_b_qserve_chn = torch.randn((N,), device="cuda", dtype=torch.float16)
|
||||
a_sum_qserve_chn = torch.randn((M,), device="cuda", dtype=torch.float16)
|
||||
|
||||
# For Qserve W4A8 per group
|
||||
group_size = 128
|
||||
assert K % group_size == 0, "K must be divisible by group_size"
|
||||
a_qserve_group = a
|
||||
# two int4s pack into one int8
|
||||
b_qserve_group = to_int8(torch.randn((N, K // 2), device="cuda") * 5)
|
||||
# b_qserve_group = b.t().contiguous()
|
||||
scale_a_qserve_group = scale_a.to(torch.float16)
|
||||
scale_b_qserve_group = scale_b.to(torch.float16)
|
||||
scale_i8_b_qserve_group = to_int8(
|
||||
torch.randn((K // group_size, N), device="cuda", dtype=torch.float16)
|
||||
)
|
||||
zero_i8_b_qserve_group = to_int8(
|
||||
torch.randn((K // group_size, N), device="cuda", dtype=torch.float16)
|
||||
)
|
||||
|
||||
quantiles = [0.5, 0.2, 0.8]
|
||||
if provider == "FP16":
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: torch.matmul(a_fp16, b_fp16),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
if provider == "W8A8":
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: int8_scaled_mm(a, b, scale_a, scale_b, torch.float16),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
if provider == "Qserve_W4A8_Per_Channel":
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: qserve_w4a8_per_chn_gemm(
|
||||
a_qserve_chn,
|
||||
b_qserve_chn,
|
||||
scale_b_qserve_chn,
|
||||
scale_a_qserve_chn,
|
||||
szero_b_qserve_chn,
|
||||
a_sum_qserve_chn,
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
if provider == "Qserve_W4A8_Per_Group":
|
||||
ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
|
||||
lambda: qserve_w4a8_per_group_gemm(
|
||||
a_qserve_group,
|
||||
b_qserve_group,
|
||||
zero_i8_b_qserve_group,
|
||||
scale_i8_b_qserve_group,
|
||||
scale_b_qserve_group,
|
||||
scale_a_qserve_group,
|
||||
),
|
||||
quantiles=quantiles,
|
||||
)
|
||||
|
||||
return ms, max_ms, min_ms
|
||||
|
||||
|
||||
def prepare_shapes(args):
|
||||
KN_model_names = []
|
||||
models_tps = list(itertools.product(args.models, args.tp_sizes))
|
||||
for model, tp_size in models_tps:
|
||||
assert model in WEIGHT_SHAPES
|
||||
for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model]):
|
||||
KN[tp_split_dim] = KN[tp_split_dim] // tp_size
|
||||
KN.append(model)
|
||||
KN_model_names.append(KN)
|
||||
return KN_model_names
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--models",
|
||||
nargs="+",
|
||||
type=str,
|
||||
default=["meta-llama/Llama-3.1-8B-Instruct"],
|
||||
help="List of models to benchmark",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tp-sizes",
|
||||
nargs="+",
|
||||
type=int,
|
||||
default=[1],
|
||||
help="List of tensor parallel sizes",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
# Skip in CI environment
|
||||
if IS_CI:
|
||||
print("Skipping QServe W4A8 GEMM benchmark in CI environment")
|
||||
print("QServe operations may have compatibility issues in CI")
|
||||
else:
|
||||
KN_model_names = prepare_shapes(args)
|
||||
|
||||
for K, N, model_name in KN_model_names:
|
||||
print(f"{model_name} N={N} K={K}: ")
|
||||
benchmark.run(
|
||||
print_data=True,
|
||||
N=N,
|
||||
K=K,
|
||||
)
|
||||
|
||||
print("Benchmark finished!")
|
||||
@@ -406,19 +406,6 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
|
||||
m.def("apply_token_bitmask_inplace_cuda(Tensor logits, Tensor bitmask, Tensor? indices=None) -> ()");
|
||||
m.impl("apply_token_bitmask_inplace_cuda", &ApplyTokenBitmaskInplace);
|
||||
|
||||
/*
|
||||
* From csrc/gemm (QServe)
|
||||
*/
|
||||
m.def(
|
||||
"qserve_w4a8_per_chn_gemm(Tensor _in_feats, Tensor _kernel, Tensor _wscales, Tensor _ascales, Tensor _w_szs, "
|
||||
"Tensor _a_ssums, Tensor! _out_feats) -> ()");
|
||||
m.impl("qserve_w4a8_per_chn_gemm", torch::kCUDA, &qserve_w4a8_per_chn_gemm);
|
||||
|
||||
m.def(
|
||||
"qserve_w4a8_per_group_gemm(Tensor _in_feats, Tensor _kernel, Tensor _zeros, Tensor _scales_i8, Tensor _wscales, "
|
||||
"Tensor _ascales, Tensor! _out_feats) -> ()");
|
||||
m.impl("qserve_w4a8_per_group_gemm", torch::kCUDA, &qserve_w4a8_per_group_gemm);
|
||||
|
||||
/*
|
||||
* From csrc/quantization/gguf
|
||||
*/
|
||||
|
||||
@@ -1,710 +0,0 @@
|
||||
// Implemented by Haotian Tang and Shang Yang.
|
||||
// @article{lin2024qserve,
|
||||
// title={QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving},
|
||||
// author={Lin*, Yujun and Tang*, Haotian and Yang*, Shang and Zhang, Zhekai and Xiao, Guangxuan and Gan, Chuang and
|
||||
// Han, Song}, journal={arXiv preprint arXiv:2405.04532}, year={2024}
|
||||
// }
|
||||
// @article{yang2025lserve,
|
||||
// title={LServe: Efficient Long-sequence LLM Serving with Unified Sparse Attention},
|
||||
// author={Yang*, Shang and Guo*, Junxian and Tang, Haotian and Hu, Qinghao and Xiao, Guangxuan and Tang, Jiaming and
|
||||
// Lin, Yujun and Liu, Zhijian and Lu, Yao and Han, Song}, year={2025}
|
||||
// }
|
||||
|
||||
// Adapted from https://github.com/mit-han-lab/omniserve/blob/main/kernels/csrc/qgemm/w4a8_per_chn/gemm_cuda.cu
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_pipeline_primitives.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include "utils.h"
|
||||
|
||||
#define OP_M 16
|
||||
#define OP_N 8
|
||||
#define OP_K 32
|
||||
#define INTRIN_M 16
|
||||
#define INTRIN_N 16
|
||||
#define INTRIN_K 32
|
||||
#define WARP_SIZE 32
|
||||
#define SMEM_PAD_A 0
|
||||
#define SMEM_PAD_B 0
|
||||
#define PACK_SIZE 16
|
||||
#if (__CUDACC_VER_MAJOR__ >= 11) && (__CUDACC_VER_MINOR__ >= 4)
|
||||
#define L2_CACHEHINT(size) ".L2::" #size "B"
|
||||
#else
|
||||
#define L2_CACHEHINT(size)
|
||||
#endif
|
||||
|
||||
#define KERNEL_LAUNCH_CODE \
|
||||
constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N) * (CTA_K / WARP_K); \
|
||||
constexpr int SCALES_SMEM_SIZE = (G >= CTA_K) ? (CTA_N * STAGES * 2) : (CTA_N * (CTA_K / G) * STAGES * 2); \
|
||||
constexpr int kSmemByteSize = \
|
||||
((CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B) / 2) * STAGES + SCALES_SMEM_SIZE) * \
|
||||
sizeof(int8_t); \
|
||||
if (kSmemByteSize >= 99 * 1024) { \
|
||||
printf( \
|
||||
"This kernel requires %d Bytes of shared memory, which exceeds " \
|
||||
"device limit.\n", \
|
||||
kSmemByteSize); \
|
||||
return; \
|
||||
} \
|
||||
int num_blocks_m = (num_out_feats + CTA_M - 1) / CTA_M; \
|
||||
int num_blocks_n = num_out_channels / CTA_N / 1; \
|
||||
const int log_tile = get_log_tile<8>((num_out_feats + CTA_M - 1) / CTA_M); \
|
||||
const int tile_shift = 1 << log_tile; \
|
||||
dim3 num_blocks(num_blocks_n* tile_shift, (num_blocks_m + tile_shift - 1) / tile_shift); \
|
||||
dim3 threads_per_block(WARP_SIZE, NUM_WARPS); \
|
||||
auto kernel_func = dense_kernel0<CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES, G>; \
|
||||
cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize); \
|
||||
kernel_func<<<num_blocks, threads_per_block, kSmemByteSize, stream>>>( \
|
||||
in_feats, kernel, wscales, ascales, w_szs, a_ssums, out_feats, num_in_feats, num_out_channels, num_in_channels);
|
||||
|
||||
template <int N>
|
||||
__inline__ __host__ __device__ int get_log_tile(int n) {
|
||||
if (N >= 8 && n >= 6)
|
||||
return 3;
|
||||
else if (N >= 4 && n >= 3)
|
||||
return 2;
|
||||
else if (N >= 2 && n >= 2)
|
||||
return 1;
|
||||
else
|
||||
return 0;
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
__inline__ __device__ uint2 get_block_idx_mapping(int blockIdx_x, int blockIdx_y, int log_tile) {
|
||||
return make_uint2((blockIdx_x >> log_tile), (blockIdx_y << log_tile) + ((blockIdx_x) & ((1 << (log_tile)) - 1)));
|
||||
}
|
||||
|
||||
__inline__ __device__ uint32_t cast_smem_ptr_to_uint(void const* const ptr) {
|
||||
uint32_t smem_int_ptr;
|
||||
|
||||
asm("{.reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, "
|
||||
"smem_ptr; }\n"
|
||||
: "=r"(smem_int_ptr)
|
||||
: "l"(ptr));
|
||||
|
||||
return smem_int_ptr;
|
||||
}
|
||||
|
||||
__inline__ __device__ void ldmatrix_m8n8_x4_b16(int8_t* shared_warp, int ax0_0, uint32_t addr) {
|
||||
__asm__ __volatile__(
|
||||
"ldmatrix.sync.aligned.m8n8.x4.shared.b16"
|
||||
"{%0, %1, %2, %3}, [%4];"
|
||||
: "=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[0]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[1]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[2]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[3])
|
||||
: "r"(addr));
|
||||
}
|
||||
|
||||
__inline__ __device__ void ldmatrix_m8n8_x4_trans_b16(int8_t* shared_warp, int ax0_0, uint32_t addr) {
|
||||
__asm__ __volatile__(
|
||||
"ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16"
|
||||
"{%0, %1, %2, %3}, [%4];"
|
||||
: "=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[0]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[1]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[2]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[3])
|
||||
: "r"(addr));
|
||||
}
|
||||
|
||||
// function from lmdeploy
|
||||
__inline__ __device__ void cp_async_cg_A(uint32_t smem_int_ptr, const uint4* __restrict__ src, bool mask) {
|
||||
const int cp_size = 16;
|
||||
asm volatile("{"
|
||||
" .reg .pred p;"
|
||||
" setp.ne.b32 p, %0, 0;"
|
||||
" @p cp.async.cg.shared.global" L2_CACHEHINT(128) " [%1], [%2], %3;"
|
||||
"}" ::"r"((int)mask),
|
||||
"r"(smem_int_ptr),
|
||||
"l"(src),
|
||||
"n"(cp_size));
|
||||
}
|
||||
|
||||
__device__ __inline__ void mma_m16n8k32(void* C_warp, void* A_shared_warp, void* B_shared_warp) {
|
||||
__asm__ __volatile__(
|
||||
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32"
|
||||
"{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};"
|
||||
: "=r"(((int*)C_warp)[0]), "=r"(((int*)C_warp)[1]), "=r"(((int*)C_warp)[2]), "=r"(((int*)C_warp)[3])
|
||||
: "r"(((unsigned*)A_shared_warp)[0]),
|
||||
"r"(((unsigned*)A_shared_warp)[1]),
|
||||
"r"(((unsigned*)A_shared_warp)[2]),
|
||||
"r"(((unsigned*)A_shared_warp)[3]),
|
||||
"r"(((unsigned*)B_shared_warp)[0]),
|
||||
"r"(((unsigned*)B_shared_warp)[1]),
|
||||
"r"(((int*)C_warp)[0]),
|
||||
"r"(((int*)C_warp)[1]),
|
||||
"r"(((int*)C_warp)[2]),
|
||||
"r"(((int*)C_warp)[3]));
|
||||
}
|
||||
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int SHARED_K_ITERS, int STAGES>
|
||||
__device__ __inline__ void global_to_share_one_stage_A(
|
||||
int8_t* src,
|
||||
int8_t* dst,
|
||||
int global_ncols,
|
||||
int cta_offset_m,
|
||||
int cta_offset_n,
|
||||
int global_iter_k,
|
||||
int shared_iter_k,
|
||||
bool mask,
|
||||
bool* preds) {
|
||||
constexpr int total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / CTA_SIZE;
|
||||
constexpr int partial_global_iters = total_global_iters / SHARED_K_ITERS;
|
||||
constexpr int cta_step_m_or_n = (CTA_SIZE * PACK_SIZE) / CTA_K;
|
||||
constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K;
|
||||
constexpr int threads_per_row = CTA_K / PACK_SIZE;
|
||||
constexpr int kSmemCol = CTA_K + SMEM_PAD_A;
|
||||
int8_t* dst_hoisted = dst;
|
||||
int8_t* src_hoisted = src + global_iter_k * CTA_K;
|
||||
|
||||
if (mask) {
|
||||
#pragma unroll
|
||||
for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter) {
|
||||
int global_iter = shared_iter_k * partial_global_iters + _global_iter;
|
||||
|
||||
void* dst_ptr = (void*)(dst_hoisted + global_iter * cta_step_m_or_n * kSmemCol);
|
||||
uint4* src_ptr = (uint4*)(src_hoisted + global_iter * cta_step_m_or_n * global_ncols);
|
||||
// *dst_ptr = *src_ptr;
|
||||
if constexpr (STAGES > 1) {
|
||||
uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);
|
||||
cp_async_cg_A(addr, src_ptr, preds[global_iter]);
|
||||
} else {
|
||||
if (preds[global_iter]) *(uint4*)dst_ptr = *src_ptr;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int SHARED_K_ITERS, int STAGES>
|
||||
__device__ __inline__ void global_to_share_one_stage_B(
|
||||
int8_t* src,
|
||||
int8_t* dst,
|
||||
int global_ncols,
|
||||
int cta_offset_m,
|
||||
int cta_offset_n,
|
||||
int global_iter_k,
|
||||
int shared_iter_k,
|
||||
bool mask) {
|
||||
constexpr int total_global_iters = (CTA_N * CTA_K) / 32 / CTA_SIZE;
|
||||
constexpr int NUM_WARPS = CTA_SIZE / WARP_SIZE;
|
||||
constexpr int warps_per_row = CTA_K / 32;
|
||||
constexpr int cta_step_m_or_n = NUM_WARPS / warps_per_row;
|
||||
constexpr int kSmemCol = CTA_K;
|
||||
int8_t* dst_hoisted = dst;
|
||||
int8_t* src_hoisted = src + global_iter_k * CTA_K * PACK_SIZE;
|
||||
|
||||
#pragma unroll
|
||||
for (int global_iter = 0; global_iter < total_global_iters; ++global_iter) {
|
||||
void* dst_ptr = (void*)(dst_hoisted + global_iter * cta_step_m_or_n * kSmemCol * PACK_SIZE);
|
||||
uint4* src_ptr = (uint4*)(src_hoisted + global_iter * cta_step_m_or_n * global_ncols * PACK_SIZE);
|
||||
if constexpr (STAGES > 1) {
|
||||
uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);
|
||||
cp_async_cg_A(addr, src_ptr, mask);
|
||||
} else {
|
||||
if (mask) *(uint4*)dst_ptr = *src_ptr;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int STAGES, int G>
|
||||
__device__ __inline__ void global_to_share_one_stage_zeros(
|
||||
int8_t* src,
|
||||
int8_t* dst,
|
||||
int global_ncols,
|
||||
int cta_offset_m,
|
||||
int cta_offset_n,
|
||||
int global_iter_k,
|
||||
int shared_iter_k,
|
||||
bool mask) {
|
||||
constexpr int threads_needed = CTA_N / PACK_SIZE / 1;
|
||||
constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE;
|
||||
constexpr int total_global_iters = CTA_N / PACK_SIZE / threads_used;
|
||||
constexpr int threads_per_row = CTA_N / PACK_SIZE;
|
||||
constexpr int kSmemCol = CTA_N;
|
||||
bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used);
|
||||
int g_idx = global_iter_k * CTA_K / G;
|
||||
|
||||
void* dst_ptr = (void*)(dst + (threadIdx.x % threads_per_row) * PACK_SIZE);
|
||||
uint4* src_ptr = (uint4*)(src + g_idx * global_ncols + cta_offset_n + (threadIdx.x % threads_per_row) * PACK_SIZE);
|
||||
if (STAGES > 1) {
|
||||
uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);
|
||||
cp_async_cg_A(addr, src_ptr, local_mask);
|
||||
} else {
|
||||
if (local_mask) {
|
||||
*(uint4*)dst_ptr = *src_ptr;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int STAGES>
|
||||
__device__ __inline__ void
|
||||
share_to_reg_one_stage_A(int8_t* src, int8_t* dst, int warp_offset_m, int warp_offset_n, int k_0_1, int shared_iters) {
|
||||
constexpr int kSmemCol = CTA_K + SMEM_PAD_A;
|
||||
int ld_col = (k_0_1 * INTRIN_K + (threadIdx.x / 16) * 16) / PACK_SIZE;
|
||||
|
||||
for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) {
|
||||
int ld_row = warp_offset_m + shared_iter * INTRIN_M + (threadIdx.x % 16);
|
||||
int ld_col_swizzled = ld_col ^ (ld_row / 2) & 3;
|
||||
void* addr_ptr = (void*)(src + ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE);
|
||||
uint32_t addr = cast_smem_ptr_to_uint(addr_ptr);
|
||||
ldmatrix_m8n8_x4_b16(dst, shared_iter, addr);
|
||||
}
|
||||
}
|
||||
|
||||
template <int WARP_K, int CTA_N, int CTA_K, int CTA_SIZE, int STAGES, int G>
|
||||
__device__ __inline__ void share_to_reg_one_stage_B(
|
||||
int8_t* src,
|
||||
int8_t* dst,
|
||||
int8_t* zeros,
|
||||
int8_t* scales_i8,
|
||||
int warp_offset_m,
|
||||
int warp_offset_n,
|
||||
int k_0_0,
|
||||
int k_0_1,
|
||||
int shared_iters) {
|
||||
constexpr int kSmemCol = CTA_K + SMEM_PAD_B;
|
||||
#pragma unroll
|
||||
for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) {
|
||||
uint4 loaded =
|
||||
*((uint4*)(src) + warp_offset_n / 32 * kSmemCol + shared_iter * 32 / 32 * kSmemCol + k_0_1 * INTRIN_K +
|
||||
threadIdx.x);
|
||||
|
||||
auto ptr = (uint32_t*)dst + shared_iter * 8;
|
||||
ptr[0] = loaded.x & 0x0F0F0F0F;
|
||||
ptr[4] = (loaded.x & 0xF0F0F0F0) >> 4;
|
||||
ptr[2] = loaded.y & 0x0F0F0F0F;
|
||||
ptr[6] = (loaded.y & 0xF0F0F0F0) >> 4;
|
||||
ptr[1] = loaded.z & 0x0F0F0F0F;
|
||||
ptr[5] = (loaded.z & 0xF0F0F0F0) >> 4;
|
||||
ptr[3] = loaded.w & 0x0F0F0F0F;
|
||||
ptr[7] = (loaded.w & 0xF0F0F0F0) >> 4;
|
||||
}
|
||||
}
|
||||
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int WARP_M, int WARP_N, int WARP_K, int STAGES, int G>
|
||||
__global__ void dense_kernel0(
|
||||
int8_t* __restrict__ A,
|
||||
int8_t* __restrict__ B,
|
||||
half2* __restrict__ wscales,
|
||||
half* __restrict__ ascales,
|
||||
half2* __restrict__ w_szs,
|
||||
half* __restrict__ a_ssums,
|
||||
half* __restrict__ C,
|
||||
int M,
|
||||
int64_t N,
|
||||
int64_t K) {
|
||||
constexpr int SPLITK = 1;
|
||||
constexpr int NUM_WARPS_MN = CTA_M / WARP_M * CTA_N / WARP_N;
|
||||
constexpr int NUM_WARPS = NUM_WARPS_MN * CTA_K / WARP_K;
|
||||
constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE;
|
||||
constexpr int CTA_SIZE_MN = NUM_WARPS_MN * WARP_SIZE;
|
||||
constexpr int SLICES = CTA_K / WARP_K;
|
||||
int num_blocks_n = (N + CTA_N - 1) / CTA_N;
|
||||
int num_blocks_m = (M + CTA_M - 1) / CTA_M;
|
||||
|
||||
int blockIdx_n = blockIdx.x;
|
||||
int blockIdx_m = blockIdx.y;
|
||||
const int log_tile = get_log_tile<8>((M + CTA_M - 1) / CTA_M);
|
||||
const uint2 block_idx_mapping = get_block_idx_mapping(blockIdx_n, blockIdx_m, log_tile);
|
||||
blockIdx_n = block_idx_mapping.x;
|
||||
blockIdx_m = block_idx_mapping.y;
|
||||
|
||||
int C_warp[CTA_M * CTA_N / CTA_SIZE_MN];
|
||||
constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A;
|
||||
constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B;
|
||||
constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA;
|
||||
constexpr int kSmemSizeBPerStage = CTA_N * kSmemPadKB / 2;
|
||||
constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES;
|
||||
constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES;
|
||||
|
||||
constexpr int scales_load_interval = G >= CTA_K ? G / CTA_K : 1;
|
||||
constexpr int scales_per_load = G < CTA_K ? CTA_K / G : 1;
|
||||
constexpr int kSmemSizeScales = CTA_N * STAGES;
|
||||
|
||||
extern __shared__ int8_t mem_shared[];
|
||||
int8_t* A_shared = mem_shared;
|
||||
|
||||
int8_t* B_shared = mem_shared + kSmemSizeA;
|
||||
int8_t* zeros_shared = mem_shared + kSmemSizeA + kSmemSizeB;
|
||||
int8_t* scales_i8_shared = mem_shared + kSmemSizeA + kSmemSizeB + kSmemSizeScales;
|
||||
|
||||
int8_t A_shared_warp_[2][WARP_M * WARP_K / WARP_SIZE];
|
||||
int8_t B_shared_warp_[2][WARP_N * WARP_K / WARP_SIZE];
|
||||
constexpr int A_total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / CTA_SIZE;
|
||||
constexpr int B_total_global_iters = (CTA_N * CTA_K) / PACK_SIZE / CTA_SIZE;
|
||||
constexpr int A_src_step_m = (CTA_SIZE * PACK_SIZE) / CTA_K;
|
||||
constexpr int A_warp_step_m = (WARP_SIZE * PACK_SIZE) / CTA_K;
|
||||
constexpr int A_threads_per_row = CTA_K / PACK_SIZE;
|
||||
|
||||
constexpr int B_warps_per_row = CTA_K / 32;
|
||||
constexpr int B_src_step_n = NUM_WARPS / B_warps_per_row;
|
||||
|
||||
int cta_offset_m = blockIdx_m * CTA_M;
|
||||
int cta_offset_n = blockIdx_n * CTA_N;
|
||||
int warp_mn = threadIdx.y % NUM_WARPS_MN;
|
||||
int slice_id = threadIdx.y / NUM_WARPS_MN;
|
||||
int warp_offset_m = (warp_mn % (CTA_M / WARP_M)) * WARP_M;
|
||||
int warp_offset_n = (warp_mn / (CTA_M / WARP_M)) * WARP_N;
|
||||
int warp_offset_k = slice_id * WARP_K;
|
||||
|
||||
for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE_MN; i++)
|
||||
C_warp[i] = 0;
|
||||
|
||||
int gemm_iters = (K + CTA_K - 1) / CTA_K;
|
||||
int k_0_0_ld = 0;
|
||||
int k_0_0 = 0;
|
||||
constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1;
|
||||
int A_hoisted_row = threadIdx.y * A_warp_step_m + (threadIdx.x / A_threads_per_row);
|
||||
int A_hoisted_col = (threadIdx.x % A_threads_per_row);
|
||||
int A_hoisted_col_swizzled = A_hoisted_col ^ (A_hoisted_row / 2) & 3;
|
||||
|
||||
int8_t* A_shared_hoisted = A_shared + A_hoisted_row * kSmemPadKA + A_hoisted_col_swizzled * PACK_SIZE;
|
||||
int8_t* B_shared_hoisted = B_shared + (threadIdx.y % B_warps_per_row) * 32 * PACK_SIZE +
|
||||
(threadIdx.y / B_warps_per_row) * kSmemPadKB * PACK_SIZE + threadIdx.x * PACK_SIZE;
|
||||
int8_t* A_hoisted = A + cta_offset_m * K + A_hoisted_row * K + A_hoisted_col * PACK_SIZE;
|
||||
int8_t* B_hoisted = B + cta_offset_n / 32 * K * PACK_SIZE + (threadIdx.y % B_warps_per_row) * 32 * PACK_SIZE +
|
||||
(threadIdx.y / B_warps_per_row) * K * PACK_SIZE + threadIdx.x * PACK_SIZE;
|
||||
|
||||
bool A_g2s_preds[A_total_global_iters];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < A_total_global_iters; i++) {
|
||||
A_g2s_preds[i] = (cta_offset_m + A_hoisted_row + i * A_src_step_m) < M;
|
||||
}
|
||||
|
||||
int* C_shared = reinterpret_cast<int*>(mem_shared);
|
||||
|
||||
#pragma unroll
|
||||
for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld) {
|
||||
global_to_share_one_stage_A<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(
|
||||
A_hoisted,
|
||||
A_shared_hoisted + k_0_0_ld * kSmemSizeAPerStage,
|
||||
K,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
0,
|
||||
true,
|
||||
A_g2s_preds);
|
||||
global_to_share_one_stage_B<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(
|
||||
B_hoisted, B_shared_hoisted + k_0_0_ld * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, k_0_0_ld, 0, true);
|
||||
|
||||
if constexpr (STAGES > 1) __pipeline_commit();
|
||||
}
|
||||
if constexpr (STAGES > 1) __pipeline_wait_prior(STAGES - 2);
|
||||
__syncthreads();
|
||||
|
||||
share_to_reg_one_stage_A<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
|
||||
A_shared + warp_offset_k, A_shared_warp_[0], warp_offset_m, warp_offset_n, 0, WARP_M / INTRIN_M);
|
||||
share_to_reg_one_stage_B<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
|
||||
B_shared + warp_offset_k * PACK_SIZE,
|
||||
B_shared_warp_[0],
|
||||
zeros_shared,
|
||||
scales_i8_shared,
|
||||
warp_offset_m,
|
||||
warp_offset_n,
|
||||
0,
|
||||
0,
|
||||
WARP_N / 32);
|
||||
constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K;
|
||||
|
||||
for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld) {
|
||||
int ld_stage = k_0_0_ld % STAGES;
|
||||
int compute_stage = k_0_0 % STAGES;
|
||||
int8_t* A_shared_this_compute_stage;
|
||||
int8_t* B_shared_this_compute_stage;
|
||||
int8_t* zeros_shared_this_compute_stage;
|
||||
int8_t* scales_i8_shared_this_compute_stage;
|
||||
|
||||
for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k) {
|
||||
A_shared_this_compute_stage = A_shared + compute_stage * kSmemSizeAPerStage + warp_offset_k;
|
||||
B_shared_this_compute_stage = B_shared + compute_stage * kSmemSizeBPerStage + warp_offset_k * PACK_SIZE;
|
||||
zeros_shared_this_compute_stage = zeros_shared + (compute_stage)*CTA_N;
|
||||
scales_i8_shared_this_compute_stage = scales_i8_shared + (compute_stage)*CTA_N;
|
||||
|
||||
share_to_reg_one_stage_A<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
|
||||
A_shared_this_compute_stage,
|
||||
A_shared_warp_[(iter_k + 1) % 2],
|
||||
warp_offset_m,
|
||||
warp_offset_n,
|
||||
(iter_k + 1) % SHARED_K_ITERS,
|
||||
WARP_M / INTRIN_M);
|
||||
share_to_reg_one_stage_B<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
|
||||
B_shared_this_compute_stage,
|
||||
B_shared_warp_[(iter_k + 1) % 2],
|
||||
zeros_shared_this_compute_stage,
|
||||
scales_i8_shared_this_compute_stage,
|
||||
warp_offset_m,
|
||||
warp_offset_n,
|
||||
k_0_0 + (iter_k == SHARED_K_ITERS - 1),
|
||||
(iter_k + 1) % SHARED_K_ITERS,
|
||||
WARP_N / 32);
|
||||
int8_t* A_shared_warp = A_shared_warp_[iter_k % 2];
|
||||
int8_t* B_shared_warp = B_shared_warp_[iter_k % 2];
|
||||
|
||||
for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4) {
|
||||
for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3) {
|
||||
mma_m16n8k32(
|
||||
(void*)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8),
|
||||
(void*)(A_shared_warp + i_0_3 * 16),
|
||||
(void*)(B_shared_warp + j_0_4 * 16));
|
||||
mma_m16n8k32(
|
||||
(void*)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4),
|
||||
(void*)(A_shared_warp + i_0_3 * 16),
|
||||
(void*)(B_shared_warp + j_0_4 * 16 + 8));
|
||||
}
|
||||
}
|
||||
|
||||
if (iter_k < SHARED_K_ITERS - 1) {
|
||||
if constexpr (STAGES == 1) __syncthreads();
|
||||
global_to_share_one_stage_A<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(
|
||||
A_hoisted,
|
||||
A_shared_hoisted + ld_stage * kSmemSizeAPerStage,
|
||||
K,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
iter_k,
|
||||
k_0_0_ld < gemm_iters,
|
||||
A_g2s_preds);
|
||||
global_to_share_one_stage_B<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(
|
||||
B_hoisted,
|
||||
B_shared_hoisted + ld_stage * kSmemSizeBPerStage,
|
||||
K,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
iter_k,
|
||||
k_0_0_ld < gemm_iters);
|
||||
}
|
||||
|
||||
if (iter_k == SHARED_K_ITERS - 2) {
|
||||
if constexpr (STAGES == 1 && SHARED_K_ITERS > 2) {
|
||||
__syncthreads();
|
||||
}
|
||||
global_to_share_one_stage_A<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(
|
||||
A_hoisted,
|
||||
A_shared_hoisted + ld_stage * kSmemSizeAPerStage,
|
||||
K,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
iter_k + 1,
|
||||
k_0_0_ld < gemm_iters,
|
||||
A_g2s_preds);
|
||||
global_to_share_one_stage_B<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(
|
||||
B_hoisted,
|
||||
B_shared_hoisted + ld_stage * kSmemSizeBPerStage,
|
||||
K,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
iter_k + 1,
|
||||
k_0_0_ld < gemm_iters);
|
||||
if constexpr (STAGES > 1) {
|
||||
__pipeline_commit();
|
||||
__pipeline_wait_prior(STAGES - 2);
|
||||
}
|
||||
compute_stage = (k_0_0 + 1) % STAGES;
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
__pipeline_commit();
|
||||
__pipeline_wait_prior(0);
|
||||
__syncthreads();
|
||||
|
||||
if constexpr (SLICES > 1) {
|
||||
#pragma unroll
|
||||
for (int z = 0; z < SLICES; ++z) {
|
||||
if (slice_id == z) {
|
||||
#pragma unroll
|
||||
for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) {
|
||||
#pragma unroll
|
||||
for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) {
|
||||
#pragma unroll
|
||||
for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) {
|
||||
if (z > 0) {
|
||||
C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] += C_shared
|
||||
[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 +
|
||||
((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) +
|
||||
(threadIdx.x % 4) * 2];
|
||||
}
|
||||
C_shared
|
||||
[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 +
|
||||
((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) +
|
||||
(threadIdx.x % 4) * 2] = C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id];
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (slice_id == 0) {
|
||||
#pragma unroll
|
||||
for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) {
|
||||
#pragma unroll
|
||||
for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) {
|
||||
#pragma unroll
|
||||
for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) {
|
||||
C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] = C_shared
|
||||
[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 +
|
||||
((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) +
|
||||
(threadIdx.x % 4) * 2];
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int row_wb_thd = cta_offset_m + warp_offset_m + (threadIdx.x / 4);
|
||||
int col_wb_thd = cta_offset_n + warp_offset_n + (threadIdx.x % 4) * 2;
|
||||
if (slice_id == 0) {
|
||||
for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) {
|
||||
int row_wb_1 = row_wb_thd + ax0_0_1 * OP_M;
|
||||
for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) {
|
||||
int col_wb_1 = col_wb_thd + ax1_0_1 * 16;
|
||||
int* C_warp_local = C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8;
|
||||
for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) {
|
||||
int row_wb = row_wb_1 + (local_id % 4) / 2 * 8;
|
||||
if (row_wb < M) {
|
||||
int col_wb = col_wb_1 + (local_id / 4) * 8 + (local_id % 2);
|
||||
float2 wscale = __half22float2(*(wscales + col_wb / 2));
|
||||
float2 w_sz = __half22float2(*(w_szs + col_wb / 2));
|
||||
float ascale = __half2float(ascales[row_wb]);
|
||||
float a_ssum = __half2float(a_ssums[row_wb]);
|
||||
float2 psums =
|
||||
make_float2(__int2float_rn(C_warp_local[local_id]), __int2float_rn(C_warp_local[local_id + 1]));
|
||||
psums.x = psums.x * wscale.x * ascale - w_sz.x * a_ssum;
|
||||
psums.y = psums.y * wscale.y * ascale - w_sz.y * a_ssum;
|
||||
*reinterpret_cast<half2*>(C + row_wb * N + col_wb) = __float22half2_rn(psums);
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#else
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int WARP_M, int WARP_N, int WARP_K, int STAGES, int G>
|
||||
__global__ void dense_kernel0(
|
||||
int8_t* __restrict__ A,
|
||||
int8_t* __restrict__ B,
|
||||
half2* __restrict__ wscales,
|
||||
half* __restrict__ ascales,
|
||||
half2* __restrict__ w_szs,
|
||||
half* __restrict__ a_ssums,
|
||||
half* __restrict__ C,
|
||||
int M,
|
||||
int64_t N,
|
||||
int64_t K) {
|
||||
// Not implemented for SM < 800
|
||||
assert(false);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
|
||||
void qserve_w4a8_per_chn_gemm(
|
||||
const torch::Tensor& _in_feats,
|
||||
const torch::Tensor& _kernel,
|
||||
const torch::Tensor& _wscales,
|
||||
const torch::Tensor& _ascales,
|
||||
const torch::Tensor& _w_szs,
|
||||
const torch::Tensor& _a_ssums,
|
||||
torch::Tensor& _out_feats) {
|
||||
// Check input tensor
|
||||
TORCH_CHECK(_in_feats.is_cuda(), "_in_feats must be a CUDA tensor");
|
||||
TORCH_CHECK(_in_feats.dim() == 2, "_in_feats must be a 2D tensor");
|
||||
TORCH_CHECK(_in_feats.is_contiguous(), "_in_feats must be contiguous");
|
||||
TORCH_CHECK(_in_feats.scalar_type() == torch::kInt8, "_in_feats must be int8");
|
||||
// Check kernel tensor
|
||||
TORCH_CHECK(_kernel.is_cuda(), "_kernel must be a CUDA tensor");
|
||||
TORCH_CHECK(_kernel.dim() == 2, "_kernel must be a 2D tensor");
|
||||
TORCH_CHECK(_kernel.is_contiguous(), "_kernel must be contiguous");
|
||||
TORCH_CHECK(_kernel.scalar_type() == torch::kInt8, "_kernel must be int8");
|
||||
// Check output tensor
|
||||
TORCH_CHECK(_out_feats.is_cuda(), "_out_feats must be a CUDA tensor");
|
||||
TORCH_CHECK(_out_feats.is_contiguous(), "_out_feats must be contiguous");
|
||||
TORCH_CHECK(_out_feats.scalar_type() == torch::kHalf, "_out_feats must be half");
|
||||
|
||||
int num_in_feats = _in_feats.size(0);
|
||||
int num_in_channels = _in_feats.size(1);
|
||||
int num_out_feats = _out_feats.size(-2);
|
||||
int num_out_channels = _out_feats.size(-1);
|
||||
|
||||
// Check matmul shape
|
||||
TORCH_CHECK(num_out_channels == _kernel.size(0), "num_out_channels must be equal to _kernel.size(0)");
|
||||
TORCH_CHECK(num_in_feats == num_out_feats, "num_in_feats must be equal to num_out_feats");
|
||||
|
||||
// Check _ascales
|
||||
TORCH_CHECK(_ascales.is_cuda(), "_ascales must be a CUDA tensor");
|
||||
TORCH_CHECK(_ascales.is_contiguous(), "_ascales must be contiguous");
|
||||
TORCH_CHECK(_ascales.scalar_type() == torch::kHalf, "_ascales must be half");
|
||||
TORCH_CHECK(_ascales.numel() == num_in_feats, "_ascales must have num_in_feats elements");
|
||||
|
||||
// Check _wscales
|
||||
TORCH_CHECK(_wscales.is_cuda(), "_wscales must be a CUDA tensor");
|
||||
TORCH_CHECK(_wscales.is_contiguous(), "_wscales must be contiguous");
|
||||
TORCH_CHECK(_wscales.scalar_type() == torch::kHalf, "_wscales must be half");
|
||||
TORCH_CHECK(_wscales.numel() == num_out_channels, "_wscales must have num_out_channels elements");
|
||||
|
||||
// Check _w_szs
|
||||
TORCH_CHECK(_w_szs.is_cuda(), "_w_szs must be a CUDA tensor");
|
||||
TORCH_CHECK(_w_szs.is_contiguous(), "_w_szs must be contiguous");
|
||||
TORCH_CHECK(_w_szs.scalar_type() == torch::kHalf, "_w_szs must be half");
|
||||
TORCH_CHECK(_w_szs.numel() == num_out_channels, "_w_szs must have num_out_channels elements");
|
||||
|
||||
// Check _a_ssums
|
||||
TORCH_CHECK(_a_ssums.is_cuda(), "_a_ssums must be a CUDA tensor");
|
||||
TORCH_CHECK(_a_ssums.is_contiguous(), "_a_ssums must be contiguous");
|
||||
TORCH_CHECK(_a_ssums.scalar_type() == torch::kHalf, "_a_ssums must be half");
|
||||
TORCH_CHECK(_a_ssums.numel() == num_in_feats, "_a_ssums must have num_in_feats elements");
|
||||
|
||||
auto in_feats = reinterpret_cast<int8_t*>(_in_feats.data_ptr<int8_t>());
|
||||
auto kernel = reinterpret_cast<int8_t*>(_kernel.data_ptr<int8_t>());
|
||||
auto w_szs = reinterpret_cast<half2*>(_w_szs.data_ptr());
|
||||
auto a_ssums = reinterpret_cast<half*>(_a_ssums.data_ptr());
|
||||
auto wscales = reinterpret_cast<half2*>(_wscales.data_ptr());
|
||||
auto ascales = reinterpret_cast<half*>(_ascales.data_ptr());
|
||||
auto out_feats = reinterpret_cast<half*>(_out_feats.data_ptr<at::Half>());
|
||||
auto stream = at::cuda::getCurrentCUDAStream(_in_feats.get_device());
|
||||
|
||||
auto sm_version = getSMVersion();
|
||||
if (sm_version >= 80) {
|
||||
constexpr int G = 128;
|
||||
|
||||
if (num_out_feats > 256) {
|
||||
constexpr int CTA_M = 128;
|
||||
constexpr int CTA_N = 128;
|
||||
constexpr int CTA_K = 64;
|
||||
constexpr int WARP_M = 64;
|
||||
constexpr int WARP_N = 32;
|
||||
constexpr int WARP_K = 64;
|
||||
constexpr int STAGES = 3;
|
||||
KERNEL_LAUNCH_CODE
|
||||
} else if (num_out_feats >= 128) {
|
||||
constexpr int CTA_M = 64;
|
||||
constexpr int CTA_N = 64;
|
||||
constexpr int CTA_K = 64;
|
||||
constexpr int WARP_M = 32;
|
||||
constexpr int WARP_N = 32;
|
||||
constexpr int WARP_K = 64;
|
||||
constexpr int STAGES = 4;
|
||||
KERNEL_LAUNCH_CODE
|
||||
} else {
|
||||
constexpr int CTA_M = 32;
|
||||
constexpr int CTA_N = 64;
|
||||
constexpr int CTA_K = 128;
|
||||
constexpr int WARP_M = 32;
|
||||
constexpr int WARP_N = 32;
|
||||
constexpr int WARP_K = 64;
|
||||
constexpr int STAGES = 3;
|
||||
KERNEL_LAUNCH_CODE
|
||||
}
|
||||
} else {
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false, "No implemented qserve_w4a8_per_chn_gemm for current compute capability: ", sm_version);
|
||||
}
|
||||
return;
|
||||
}
|
||||
@@ -1,795 +0,0 @@
|
||||
// Implemented by Haotian Tang and Shang Yang.
|
||||
// @article{lin2024qserve,
|
||||
// title={QServe: W4A8KV4 Quantization and System Co-design for Efficient LLM Serving},
|
||||
// author={Lin*, Yujun and Tang*, Haotian and Yang*, Shang and Zhang, Zhekai and Xiao, Guangxuan and Gan, Chuang and
|
||||
// Han, Song}, journal={arXiv preprint arXiv:2405.04532}, year={2024}
|
||||
// }
|
||||
// @article{yang2025lserve,
|
||||
// title={LServe: Efficient Long-sequence LLM Serving with Unified Sparse Attention},
|
||||
// author={Yang*, Shang and Guo*, Junxian and Tang, Haotian and Hu, Qinghao and Xiao, Guangxuan and Tang, Jiaming and
|
||||
// Lin, Yujun and Liu, Zhijian and Lu, Yao and Han, Song}, year={2025}
|
||||
// }
|
||||
|
||||
// Adapted from https://github.com/mit-han-lab/omniserve/blob/main/kernels/csrc/qgemm/w4a8_per_group/gemm_cuda.cu
|
||||
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <cuda_fp16.h>
|
||||
#include <cuda_pipeline_primitives.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include "utils.h"
|
||||
|
||||
#define OP_M 16
|
||||
#define OP_N 8
|
||||
#define OP_K 32
|
||||
#define INTRIN_M 16
|
||||
#define INTRIN_N 16
|
||||
#define INTRIN_K 32
|
||||
#define WARP_SIZE 32
|
||||
#define SMEM_PAD_A 0
|
||||
#define SMEM_PAD_B 0
|
||||
#define PACK_SIZE 16
|
||||
#if (__CUDACC_VER_MAJOR__ >= 11) && (__CUDACC_VER_MINOR__ >= 4)
|
||||
#define L2_CACHEHINT(size) ".L2::" #size "B"
|
||||
#else
|
||||
#define L2_CACHEHINT(size)
|
||||
#endif
|
||||
|
||||
#define KERNEL_LAUNCH_CODE \
|
||||
constexpr int NUM_WARPS = (CTA_M / WARP_M) * (CTA_N / WARP_N) * (CTA_K / WARP_K); \
|
||||
constexpr int SCALES_SMEM_SIZE = (G >= CTA_K) ? (CTA_N * STAGES * 2) : (CTA_N * (CTA_K / G) * STAGES * 2); \
|
||||
constexpr int kSmemByteSize = \
|
||||
((CTA_M * (CTA_K + SMEM_PAD_A) + CTA_N * (CTA_K + SMEM_PAD_B) / 2) * STAGES + SCALES_SMEM_SIZE) * \
|
||||
sizeof(int8_t); \
|
||||
if (kSmemByteSize >= 99 * 1024) { \
|
||||
printf( \
|
||||
"This kernel requires %d Bytes of shared memory, which exceeds " \
|
||||
"device limit.\n", \
|
||||
kSmemByteSize); \
|
||||
return; \
|
||||
} \
|
||||
int num_blocks_m = (num_out_feats + CTA_M - 1) / CTA_M; \
|
||||
int num_blocks_n = num_out_channels / CTA_N / 1; \
|
||||
const int log_tile = get_log_tile<8>((num_out_feats + CTA_M - 1) / CTA_M); \
|
||||
const int tile_shift = 1 << log_tile; \
|
||||
dim3 num_blocks(num_blocks_n* tile_shift, (num_blocks_m + tile_shift - 1) / tile_shift); \
|
||||
dim3 threads_per_block(WARP_SIZE, NUM_WARPS); \
|
||||
auto kernel_func = dense_kernel0<CTA_M, CTA_N, CTA_K, WARP_M, WARP_N, WARP_K, STAGES, G>; \
|
||||
cudaFuncSetAttribute(kernel_func, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemByteSize); \
|
||||
kernel_func<<<num_blocks, threads_per_block, kSmemByteSize, stream>>>( \
|
||||
in_feats, \
|
||||
kernel, \
|
||||
zeros, \
|
||||
scales_i8, \
|
||||
wscales, \
|
||||
ascales, \
|
||||
out_feats, \
|
||||
num_in_feats, \
|
||||
num_out_channels, \
|
||||
num_in_channels);
|
||||
|
||||
template <int N>
|
||||
__inline__ __host__ __device__ int get_log_tile(int n) {
|
||||
if (N >= 8 && n >= 6)
|
||||
return 3;
|
||||
else if (N >= 4 && n >= 3)
|
||||
return 2;
|
||||
else if (N >= 2 && n >= 2)
|
||||
return 1;
|
||||
else
|
||||
return 0;
|
||||
}
|
||||
|
||||
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
|
||||
__inline__ __device__ uint2 get_block_idx_mapping(int blockIdx_x, int blockIdx_y, int log_tile) {
|
||||
return make_uint2((blockIdx_x >> log_tile), (blockIdx_y << log_tile) + ((blockIdx_x) & ((1 << (log_tile)) - 1)));
|
||||
}
|
||||
|
||||
__inline__ __device__ uint32_t cast_smem_ptr_to_uint(void const* const ptr) {
|
||||
uint32_t smem_int_ptr;
|
||||
|
||||
asm("{.reg .u64 smem_ptr; cvta.to.shared.u64 smem_ptr, %1; cvt.u32.u64 %0, "
|
||||
"smem_ptr; }\n"
|
||||
: "=r"(smem_int_ptr)
|
||||
: "l"(ptr));
|
||||
|
||||
return smem_int_ptr;
|
||||
}
|
||||
|
||||
__inline__ __device__ void ldmatrix_m8n8_x4_b16(int8_t* shared_warp, int ax0_0, uint32_t addr) {
|
||||
__asm__ __volatile__(
|
||||
"ldmatrix.sync.aligned.m8n8.x4.shared.b16"
|
||||
"{%0, %1, %2, %3}, [%4];"
|
||||
: "=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[0]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[1]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[2]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[3])
|
||||
: "r"(addr));
|
||||
}
|
||||
|
||||
__inline__ __device__ void ldmatrix_m8n8_x4_trans_b16(int8_t* shared_warp, int ax0_0, uint32_t addr) {
|
||||
__asm__ __volatile__(
|
||||
"ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16"
|
||||
"{%0, %1, %2, %3}, [%4];"
|
||||
: "=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[0]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[1]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[2]),
|
||||
"=r"(((unsigned*)(shared_warp + (ax0_0 * 16)))[3])
|
||||
: "r"(addr));
|
||||
}
|
||||
|
||||
// function from lmdeploy
|
||||
__inline__ __device__ void cp_async_cg_A(uint32_t smem_int_ptr, const uint4* __restrict__ src, bool mask) {
|
||||
const int cp_size = 16;
|
||||
asm volatile("{"
|
||||
" .reg .pred p;"
|
||||
" setp.ne.b32 p, %0, 0;"
|
||||
" @p cp.async.cg.shared.global" L2_CACHEHINT(128) " [%1], [%2], %3;"
|
||||
"}" ::"r"((int)mask),
|
||||
"r"(smem_int_ptr),
|
||||
"l"(src),
|
||||
"n"(cp_size));
|
||||
}
|
||||
|
||||
__device__ __inline__ void mma_m16n8k32(void* C_warp, void* A_shared_warp, void* B_shared_warp) {
|
||||
__asm__ __volatile__(
|
||||
"mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32"
|
||||
"{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%10, %11, %12, %13};"
|
||||
: "=r"(((int*)C_warp)[0]), "=r"(((int*)C_warp)[1]), "=r"(((int*)C_warp)[2]), "=r"(((int*)C_warp)[3])
|
||||
: "r"(((unsigned*)A_shared_warp)[0]),
|
||||
"r"(((unsigned*)A_shared_warp)[1]),
|
||||
"r"(((unsigned*)A_shared_warp)[2]),
|
||||
"r"(((unsigned*)A_shared_warp)[3]),
|
||||
"r"(((unsigned*)B_shared_warp)[0]),
|
||||
"r"(((unsigned*)B_shared_warp)[1]),
|
||||
"r"(((int*)C_warp)[0]),
|
||||
"r"(((int*)C_warp)[1]),
|
||||
"r"(((int*)C_warp)[2]),
|
||||
"r"(((int*)C_warp)[3]));
|
||||
}
|
||||
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int SHARED_K_ITERS, int STAGES>
|
||||
__device__ __inline__ void global_to_share_one_stage_A(
|
||||
int8_t* src,
|
||||
int8_t* dst,
|
||||
int global_ncols,
|
||||
int cta_offset_m,
|
||||
int cta_offset_n,
|
||||
int global_iter_k,
|
||||
int shared_iter_k,
|
||||
bool mask,
|
||||
bool* preds) {
|
||||
constexpr int total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / CTA_SIZE;
|
||||
constexpr int partial_global_iters = total_global_iters / SHARED_K_ITERS;
|
||||
constexpr int cta_step_m_or_n = (CTA_SIZE * PACK_SIZE) / CTA_K;
|
||||
constexpr int warp_step_m_or_n = (WARP_SIZE * PACK_SIZE) / CTA_K;
|
||||
constexpr int threads_per_row = CTA_K / PACK_SIZE;
|
||||
constexpr int kSmemCol = CTA_K + SMEM_PAD_A;
|
||||
int8_t* dst_hoisted = dst;
|
||||
int8_t* src_hoisted = src + global_iter_k * CTA_K;
|
||||
|
||||
if (mask) {
|
||||
#pragma unroll
|
||||
for (int _global_iter = 0; _global_iter < partial_global_iters; ++_global_iter) {
|
||||
int global_iter = shared_iter_k * partial_global_iters + _global_iter;
|
||||
void* dst_ptr = (void*)(dst_hoisted + global_iter * cta_step_m_or_n * kSmemCol);
|
||||
uint4* src_ptr = (uint4*)(src_hoisted + global_iter * cta_step_m_or_n * global_ncols);
|
||||
if constexpr (STAGES > 1) {
|
||||
uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);
|
||||
cp_async_cg_A(addr, src_ptr, preds[global_iter]);
|
||||
} else {
|
||||
if (preds[global_iter]) *(uint4*)dst_ptr = *src_ptr;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int SHARED_K_ITERS, int STAGES>
|
||||
__device__ __inline__ void global_to_share_one_stage_B(
|
||||
int8_t* src,
|
||||
int8_t* dst,
|
||||
int global_ncols,
|
||||
int cta_offset_m,
|
||||
int cta_offset_n,
|
||||
int global_iter_k,
|
||||
int shared_iter_k,
|
||||
bool mask) {
|
||||
constexpr int total_global_iters = (CTA_N * CTA_K) / 32 / CTA_SIZE;
|
||||
constexpr int NUM_WARPS = CTA_SIZE / WARP_SIZE;
|
||||
constexpr int warps_per_row = CTA_K / 32;
|
||||
constexpr int cta_step_m_or_n = NUM_WARPS / warps_per_row;
|
||||
constexpr int kSmemCol = CTA_K;
|
||||
int8_t* dst_hoisted = dst;
|
||||
int8_t* src_hoisted = src + global_iter_k * CTA_K * PACK_SIZE;
|
||||
|
||||
#pragma unroll
|
||||
for (int global_iter = 0; global_iter < total_global_iters; ++global_iter) {
|
||||
void* dst_ptr = (void*)(dst_hoisted + global_iter * cta_step_m_or_n * kSmemCol * PACK_SIZE);
|
||||
uint4* src_ptr = (uint4*)(src_hoisted + global_iter * cta_step_m_or_n * global_ncols * PACK_SIZE);
|
||||
if constexpr (STAGES > 1) {
|
||||
uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);
|
||||
cp_async_cg_A(addr, src_ptr, mask);
|
||||
} else {
|
||||
if (mask) *(uint4*)dst_ptr = *src_ptr;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int STAGES, int G>
|
||||
__device__ __inline__ void global_to_share_one_stage_zeros(
|
||||
int8_t* src,
|
||||
int8_t* dst,
|
||||
int global_ncols,
|
||||
int cta_offset_m,
|
||||
int cta_offset_n,
|
||||
int global_iter_k,
|
||||
int shared_iter_k,
|
||||
bool mask) {
|
||||
constexpr int threads_needed = CTA_N / PACK_SIZE / 1;
|
||||
constexpr int threads_used = threads_needed < CTA_SIZE ? threads_needed : CTA_SIZE;
|
||||
constexpr int total_global_iters = CTA_N / PACK_SIZE / threads_used;
|
||||
constexpr int threads_per_row = CTA_N / PACK_SIZE;
|
||||
constexpr int kSmemCol = CTA_N;
|
||||
bool local_mask = mask & (threadIdx.y * WARP_SIZE + threadIdx.x < threads_used);
|
||||
int g_idx = global_iter_k * CTA_K / G;
|
||||
|
||||
void* dst_ptr = (void*)(dst + (threadIdx.x % threads_per_row) * PACK_SIZE);
|
||||
uint4* src_ptr = (uint4*)(src + g_idx * global_ncols + cta_offset_n + (threadIdx.x % threads_per_row) * PACK_SIZE);
|
||||
if (STAGES > 1) {
|
||||
uint32_t addr = cast_smem_ptr_to_uint(dst_ptr);
|
||||
cp_async_cg_A(addr, src_ptr, local_mask);
|
||||
} else {
|
||||
if (local_mask) {
|
||||
*(uint4*)dst_ptr = *src_ptr;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int CTA_SIZE, int STAGES>
|
||||
__device__ __inline__ void
|
||||
share_to_reg_one_stage_A(int8_t* src, int8_t* dst, int warp_offset_m, int warp_offset_n, int k_0_1, int shared_iters) {
|
||||
constexpr int kSmemCol = CTA_K + SMEM_PAD_A;
|
||||
int ld_col = (k_0_1 * INTRIN_K + (threadIdx.x / 16) * 16) / PACK_SIZE;
|
||||
|
||||
for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) {
|
||||
int ld_row = warp_offset_m + shared_iter * INTRIN_M + (threadIdx.x % 16);
|
||||
int ld_col_swizzled = ld_col ^ (ld_row / 2) & 3;
|
||||
void* addr_ptr = (void*)(src + ld_row * kSmemCol + ld_col_swizzled * PACK_SIZE);
|
||||
uint32_t addr = cast_smem_ptr_to_uint(addr_ptr);
|
||||
ldmatrix_m8n8_x4_b16(dst, shared_iter, addr);
|
||||
}
|
||||
}
|
||||
|
||||
template <int WARP_K, int CTA_N, int CTA_K, int CTA_SIZE, int STAGES, int G>
|
||||
__device__ __inline__ void share_to_reg_one_stage_B(
|
||||
int8_t* src,
|
||||
int8_t* dst,
|
||||
int8_t* zeros,
|
||||
int8_t* scales_i8,
|
||||
int warp_offset_m,
|
||||
int warp_offset_n,
|
||||
int k_0_0,
|
||||
int k_0_1,
|
||||
int shared_iters) {
|
||||
constexpr int kSmemCol = CTA_K + SMEM_PAD_B;
|
||||
#pragma unroll
|
||||
for (int shared_iter = 0; shared_iter < shared_iters; ++shared_iter) {
|
||||
uint4 loaded =
|
||||
*((uint4*)(src) + warp_offset_n / 32 * kSmemCol + shared_iter * 32 / 32 * kSmemCol + k_0_1 * INTRIN_K +
|
||||
threadIdx.x);
|
||||
uint32_t loaded_0 = loaded.x & 0x0F0F0F0F;
|
||||
uint32_t loaded_4 = (loaded.x & 0xF0F0F0F0) >> 4;
|
||||
uint32_t loaded_2 = loaded.y & 0x0F0F0F0F;
|
||||
uint32_t loaded_6 = (loaded.y & 0xF0F0F0F0) >> 4;
|
||||
uint32_t loaded_1 = loaded.z & 0x0F0F0F0F;
|
||||
uint32_t loaded_5 = (loaded.z & 0xF0F0F0F0) >> 4;
|
||||
uint32_t loaded_3 = loaded.w & 0x0F0F0F0F;
|
||||
uint32_t loaded_7 = (loaded.w & 0xF0F0F0F0) >> 4;
|
||||
|
||||
auto ptr = (uint32_t*)dst + shared_iter * 8;
|
||||
int scales_zeros_offset = warp_offset_n + (threadIdx.x / 4) * 4 + shared_iter * 32;
|
||||
uint32_t packed_scales = *reinterpret_cast<uint32_t*>(scales_i8 + scales_zeros_offset);
|
||||
uint32_t packed_zeros = *reinterpret_cast<uint32_t*>(zeros + scales_zeros_offset);
|
||||
|
||||
uint32_t scale_0 = packed_scales & 0xFF;
|
||||
uint32_t zero_point_0 = __byte_perm(packed_zeros, 0, 0x00000000);
|
||||
uint32_t ptr_0 = loaded_0 * scale_0;
|
||||
uint32_t ptr_1 = loaded_1 * scale_0;
|
||||
ptr[0] = __vadd4(ptr_0, zero_point_0);
|
||||
ptr[1] = __vadd4(ptr_1, zero_point_0);
|
||||
|
||||
uint32_t scale_1 = (packed_scales & 0xFF00) >> 8;
|
||||
uint32_t zero_point_1 = __byte_perm(packed_zeros, 0, 0x00001111);
|
||||
uint32_t ptr_2 = loaded_2 * scale_1;
|
||||
uint32_t ptr_3 = loaded_3 * scale_1;
|
||||
ptr[2] = __vadd4(ptr_2, zero_point_1);
|
||||
ptr[3] = __vadd4(ptr_3, zero_point_1);
|
||||
|
||||
uint32_t scale_2 = (packed_scales & 0xFF0000) >> 16;
|
||||
uint32_t zero_point_2 = __byte_perm(packed_zeros, 0, 0x00002222);
|
||||
uint32_t ptr_4 = loaded_4 * scale_2;
|
||||
uint32_t ptr_5 = loaded_5 * scale_2;
|
||||
ptr[4] = __vadd4(ptr_4, zero_point_2);
|
||||
ptr[5] = __vadd4(ptr_5, zero_point_2);
|
||||
|
||||
uint32_t scale_3 = (packed_scales & 0xFF000000) >> 24;
|
||||
uint32_t zero_point_3 = __byte_perm(packed_zeros, 0, 0x00003333);
|
||||
uint32_t ptr_6 = loaded_6 * scale_3;
|
||||
uint32_t ptr_7 = loaded_7 * scale_3;
|
||||
ptr[6] = __vadd4(ptr_6, zero_point_3);
|
||||
ptr[7] = __vadd4(ptr_7, zero_point_3);
|
||||
}
|
||||
}
|
||||
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int WARP_M, int WARP_N, int WARP_K, int STAGES, int G>
|
||||
__global__ void dense_kernel0(
|
||||
int8_t* __restrict__ A,
|
||||
int8_t* __restrict__ B,
|
||||
int8_t* __restrict__ zeros,
|
||||
int8_t* __restrict__ scales_i8,
|
||||
half2* __restrict__ wscales,
|
||||
half* __restrict__ ascales,
|
||||
half* __restrict__ C,
|
||||
int M,
|
||||
int64_t N,
|
||||
int64_t K) {
|
||||
constexpr int SPLITK = 1;
|
||||
constexpr int NUM_WARPS_MN = CTA_M / WARP_M * CTA_N / WARP_N;
|
||||
constexpr int NUM_WARPS = NUM_WARPS_MN * CTA_K / WARP_K;
|
||||
constexpr int CTA_SIZE = NUM_WARPS * WARP_SIZE;
|
||||
constexpr int CTA_SIZE_MN = NUM_WARPS_MN * WARP_SIZE;
|
||||
constexpr int SLICES = CTA_K / WARP_K;
|
||||
int num_blocks_n = (N + CTA_N - 1) / CTA_N;
|
||||
int num_blocks_m = (M + CTA_M - 1) / CTA_M;
|
||||
|
||||
int blockIdx_n = blockIdx.x;
|
||||
int blockIdx_m = blockIdx.y;
|
||||
const int log_tile = get_log_tile<8>((M + CTA_M - 1) / CTA_M);
|
||||
const uint2 block_idx_mapping = get_block_idx_mapping(blockIdx_n, blockIdx_m, log_tile);
|
||||
blockIdx_n = block_idx_mapping.x;
|
||||
blockIdx_m = block_idx_mapping.y;
|
||||
|
||||
int C_warp[CTA_M * CTA_N / CTA_SIZE_MN];
|
||||
constexpr int kSmemPadKA = CTA_K + SMEM_PAD_A;
|
||||
constexpr int kSmemPadKB = CTA_K + SMEM_PAD_B;
|
||||
constexpr int kSmemSizeAPerStage = CTA_M * kSmemPadKA;
|
||||
constexpr int kSmemSizeBPerStage = CTA_N * kSmemPadKB / 2;
|
||||
constexpr int kSmemSizeA = kSmemSizeAPerStage * STAGES;
|
||||
constexpr int kSmemSizeB = kSmemSizeBPerStage * STAGES;
|
||||
|
||||
constexpr int scales_load_interval = G >= CTA_K ? G / CTA_K : 1;
|
||||
constexpr int scales_per_load = G < CTA_K ? CTA_K / G : 1;
|
||||
constexpr int kSmemSizeScales = CTA_N * STAGES;
|
||||
|
||||
extern __shared__ int8_t mem_shared[];
|
||||
int8_t* A_shared = mem_shared;
|
||||
|
||||
int8_t* B_shared = mem_shared + kSmemSizeA;
|
||||
int8_t* zeros_shared = mem_shared + kSmemSizeA + kSmemSizeB;
|
||||
int8_t* scales_i8_shared = mem_shared + kSmemSizeA + kSmemSizeB + kSmemSizeScales;
|
||||
|
||||
int8_t A_shared_warp_[2][WARP_M * WARP_K / WARP_SIZE];
|
||||
int8_t B_shared_warp_[2][WARP_N * WARP_K / WARP_SIZE];
|
||||
constexpr int A_total_global_iters = (CTA_M * CTA_K) / PACK_SIZE / CTA_SIZE;
|
||||
constexpr int B_total_global_iters = (CTA_N * CTA_K) / PACK_SIZE / CTA_SIZE;
|
||||
constexpr int A_src_step_m = (CTA_SIZE * PACK_SIZE) / CTA_K;
|
||||
constexpr int A_warp_step_m = (WARP_SIZE * PACK_SIZE) / CTA_K;
|
||||
constexpr int A_threads_per_row = CTA_K / PACK_SIZE;
|
||||
|
||||
constexpr int B_warps_per_row = CTA_K / 32;
|
||||
constexpr int B_src_step_n = NUM_WARPS / B_warps_per_row;
|
||||
|
||||
int cta_offset_m = blockIdx_m * CTA_M;
|
||||
int cta_offset_n = blockIdx_n * CTA_N;
|
||||
int warp_mn = threadIdx.y % NUM_WARPS_MN;
|
||||
int slice_id = threadIdx.y / NUM_WARPS_MN;
|
||||
int warp_offset_m = (warp_mn % (CTA_M / WARP_M)) * WARP_M;
|
||||
int warp_offset_n = (warp_mn / (CTA_M / WARP_M)) * WARP_N;
|
||||
int warp_offset_k = slice_id * WARP_K;
|
||||
|
||||
for (int i = 0; i < CTA_M * CTA_N / CTA_SIZE_MN; i++)
|
||||
C_warp[i] = 0;
|
||||
|
||||
int gemm_iters = (K + CTA_K - 1) / CTA_K;
|
||||
|
||||
int k_0_0_ld = 0;
|
||||
int k_0_0 = 0;
|
||||
constexpr int prologue_stages = STAGES == 1 ? 1 : STAGES - 1;
|
||||
int A_hoisted_row = threadIdx.y * A_warp_step_m + (threadIdx.x / A_threads_per_row);
|
||||
int A_hoisted_col = (threadIdx.x % A_threads_per_row);
|
||||
int A_hoisted_col_swizzled = A_hoisted_col ^ (A_hoisted_row / 2) & 3;
|
||||
|
||||
int8_t* A_shared_hoisted = A_shared + A_hoisted_row * kSmemPadKA + A_hoisted_col_swizzled * PACK_SIZE;
|
||||
int8_t* B_shared_hoisted = B_shared + (threadIdx.y % B_warps_per_row) * 32 * PACK_SIZE +
|
||||
(threadIdx.y / B_warps_per_row) * kSmemPadKB * PACK_SIZE + threadIdx.x * PACK_SIZE;
|
||||
int8_t* A_hoisted = A + cta_offset_m * K + A_hoisted_row * K + A_hoisted_col * PACK_SIZE;
|
||||
int8_t* B_hoisted = B + cta_offset_n / 32 * K * PACK_SIZE + (threadIdx.y % B_warps_per_row) * 32 * PACK_SIZE +
|
||||
(threadIdx.y / B_warps_per_row) * K * PACK_SIZE + threadIdx.x * PACK_SIZE;
|
||||
|
||||
bool A_g2s_preds[A_total_global_iters];
|
||||
#pragma unroll
|
||||
for (int i = 0; i < A_total_global_iters; i++) {
|
||||
A_g2s_preds[i] = (cta_offset_m + A_hoisted_row + i * A_src_step_m) < M;
|
||||
}
|
||||
|
||||
int* C_shared = reinterpret_cast<int*>(mem_shared);
|
||||
|
||||
#pragma unroll
|
||||
for (k_0_0_ld = 0; k_0_0_ld < prologue_stages; ++k_0_0_ld) {
|
||||
global_to_share_one_stage_A<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(
|
||||
A_hoisted,
|
||||
A_shared_hoisted + k_0_0_ld * kSmemSizeAPerStage,
|
||||
K,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
0,
|
||||
true,
|
||||
A_g2s_preds);
|
||||
global_to_share_one_stage_B<CTA_M, CTA_N, CTA_K, CTA_SIZE, 1, STAGES>(
|
||||
B_hoisted, B_shared_hoisted + k_0_0_ld * kSmemSizeBPerStage, K, cta_offset_m, cta_offset_n, k_0_0_ld, 0, true);
|
||||
global_to_share_one_stage_zeros<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
|
||||
zeros, zeros_shared + (k_0_0_ld)*CTA_N, N, cta_offset_m, cta_offset_n, k_0_0_ld, 0, k_0_0_ld < gemm_iters);
|
||||
global_to_share_one_stage_zeros<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
|
||||
scales_i8,
|
||||
scales_i8_shared + (k_0_0_ld)*CTA_N,
|
||||
N,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
0,
|
||||
k_0_0_ld < gemm_iters);
|
||||
|
||||
if constexpr (STAGES > 1) __pipeline_commit();
|
||||
}
|
||||
if constexpr (STAGES > 1) __pipeline_wait_prior(STAGES - 2);
|
||||
__syncthreads();
|
||||
|
||||
share_to_reg_one_stage_A<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
|
||||
A_shared + warp_offset_k, A_shared_warp_[0], warp_offset_m, warp_offset_n, 0, WARP_M / INTRIN_M);
|
||||
share_to_reg_one_stage_B<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
|
||||
B_shared + warp_offset_k * PACK_SIZE,
|
||||
B_shared_warp_[0],
|
||||
zeros_shared,
|
||||
scales_i8_shared,
|
||||
warp_offset_m,
|
||||
warp_offset_n,
|
||||
0,
|
||||
0,
|
||||
WARP_N / 32);
|
||||
constexpr int SHARED_K_ITERS = WARP_K / INTRIN_K;
|
||||
|
||||
for (; k_0_0 < gemm_iters; ++k_0_0, ++k_0_0_ld) {
|
||||
int ld_stage = k_0_0_ld % STAGES;
|
||||
int compute_stage = k_0_0 % STAGES;
|
||||
int8_t* A_shared_this_compute_stage;
|
||||
int8_t* B_shared_this_compute_stage;
|
||||
int8_t* zeros_shared_this_compute_stage;
|
||||
int8_t* scales_i8_shared_this_compute_stage;
|
||||
|
||||
for (int iter_k = 0; iter_k < SHARED_K_ITERS; ++iter_k) {
|
||||
A_shared_this_compute_stage = A_shared + compute_stage * kSmemSizeAPerStage + warp_offset_k;
|
||||
B_shared_this_compute_stage = B_shared + compute_stage * kSmemSizeBPerStage + warp_offset_k * PACK_SIZE;
|
||||
zeros_shared_this_compute_stage = zeros_shared + (compute_stage)*CTA_N;
|
||||
scales_i8_shared_this_compute_stage = scales_i8_shared + (compute_stage)*CTA_N;
|
||||
|
||||
share_to_reg_one_stage_A<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES>(
|
||||
A_shared_this_compute_stage,
|
||||
A_shared_warp_[(iter_k + 1) % 2],
|
||||
warp_offset_m,
|
||||
warp_offset_n,
|
||||
(iter_k + 1) % SHARED_K_ITERS,
|
||||
WARP_M / INTRIN_M);
|
||||
share_to_reg_one_stage_B<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
|
||||
B_shared_this_compute_stage,
|
||||
B_shared_warp_[(iter_k + 1) % 2],
|
||||
zeros_shared_this_compute_stage,
|
||||
scales_i8_shared_this_compute_stage,
|
||||
warp_offset_m,
|
||||
warp_offset_n,
|
||||
k_0_0 + (iter_k == SHARED_K_ITERS - 1),
|
||||
(iter_k + 1) % SHARED_K_ITERS,
|
||||
WARP_N / 32);
|
||||
int8_t* A_shared_warp = A_shared_warp_[iter_k % 2];
|
||||
int8_t* B_shared_warp = B_shared_warp_[iter_k % 2];
|
||||
|
||||
for (int j_0_4 = 0; j_0_4 < WARP_N / INTRIN_N; ++j_0_4) {
|
||||
for (int i_0_3 = 0; i_0_3 < WARP_M / INTRIN_M; ++i_0_3) {
|
||||
mma_m16n8k32(
|
||||
(void*)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8),
|
||||
(void*)(A_shared_warp + i_0_3 * 16),
|
||||
(void*)(B_shared_warp + j_0_4 * 16));
|
||||
mma_m16n8k32(
|
||||
(void*)(C_warp + i_0_3 * WARP_N / INTRIN_N * 8 + j_0_4 * 8 + 4),
|
||||
(void*)(A_shared_warp + i_0_3 * 16),
|
||||
(void*)(B_shared_warp + j_0_4 * 16 + 8));
|
||||
}
|
||||
}
|
||||
|
||||
if (iter_k < SHARED_K_ITERS - 1) {
|
||||
if constexpr (STAGES == 1) __syncthreads();
|
||||
global_to_share_one_stage_A<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(
|
||||
A_hoisted,
|
||||
A_shared_hoisted + ld_stage * kSmemSizeAPerStage,
|
||||
K,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
iter_k,
|
||||
k_0_0_ld < gemm_iters,
|
||||
A_g2s_preds);
|
||||
global_to_share_one_stage_B<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(
|
||||
B_hoisted,
|
||||
B_shared_hoisted + ld_stage * kSmemSizeBPerStage,
|
||||
K,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
iter_k,
|
||||
k_0_0_ld < gemm_iters);
|
||||
}
|
||||
|
||||
if (iter_k == SHARED_K_ITERS - 2) {
|
||||
if constexpr (STAGES == 1 && SHARED_K_ITERS > 2) {
|
||||
__syncthreads();
|
||||
}
|
||||
global_to_share_one_stage_A<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(
|
||||
A_hoisted,
|
||||
A_shared_hoisted + ld_stage * kSmemSizeAPerStage,
|
||||
K,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
iter_k + 1,
|
||||
k_0_0_ld < gemm_iters,
|
||||
A_g2s_preds);
|
||||
global_to_share_one_stage_B<CTA_M, CTA_N, CTA_K, CTA_SIZE, WARP_K / INTRIN_K, STAGES>(
|
||||
B_hoisted,
|
||||
B_shared_hoisted + ld_stage * kSmemSizeBPerStage,
|
||||
K,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
iter_k + 1,
|
||||
k_0_0_ld < gemm_iters);
|
||||
global_to_share_one_stage_zeros<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
|
||||
zeros,
|
||||
zeros_shared + (ld_stage)*CTA_N,
|
||||
N,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
iter_k,
|
||||
k_0_0_ld < gemm_iters);
|
||||
global_to_share_one_stage_zeros<CTA_M, CTA_N, CTA_K, CTA_SIZE, STAGES, G>(
|
||||
scales_i8,
|
||||
scales_i8_shared + (ld_stage)*CTA_N,
|
||||
N,
|
||||
cta_offset_m,
|
||||
cta_offset_n,
|
||||
k_0_0_ld,
|
||||
iter_k,
|
||||
k_0_0_ld < gemm_iters);
|
||||
if constexpr (STAGES > 1) {
|
||||
__pipeline_commit();
|
||||
__pipeline_wait_prior(STAGES - 2);
|
||||
}
|
||||
compute_stage = (k_0_0 + 1) % STAGES;
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
}
|
||||
__pipeline_commit();
|
||||
__pipeline_wait_prior(0);
|
||||
__syncthreads();
|
||||
|
||||
if constexpr (SLICES > 1) {
|
||||
#pragma unroll
|
||||
for (int z = 0; z < SLICES; ++z) {
|
||||
if (slice_id == z) {
|
||||
#pragma unroll
|
||||
for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) {
|
||||
#pragma unroll
|
||||
for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) {
|
||||
#pragma unroll
|
||||
for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) {
|
||||
if (z > 0) {
|
||||
C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] += C_shared
|
||||
[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 +
|
||||
((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) +
|
||||
(threadIdx.x % 4) * 2];
|
||||
}
|
||||
C_shared
|
||||
[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 +
|
||||
((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) +
|
||||
(threadIdx.x % 4) * 2] = C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id];
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
if (slice_id == 0) {
|
||||
#pragma unroll
|
||||
for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) {
|
||||
#pragma unroll
|
||||
for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) {
|
||||
#pragma unroll
|
||||
for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; ++local_id) {
|
||||
C_warp[ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8 + local_id] = C_shared
|
||||
[warp_offset_m * CTA_N + ax0_0_1 * OP_M * CTA_N + warp_offset_n + ax1_0_1 * 16 +
|
||||
((local_id % 4) / 2 * 8 + (threadIdx.x / 4)) * CTA_N + (local_id / 4) * 8 + (local_id % 2) +
|
||||
(threadIdx.x % 4) * 2];
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
int row_wb_thd = cta_offset_m + warp_offset_m + (threadIdx.x / 4);
|
||||
int col_wb_thd = cta_offset_n + warp_offset_n + (threadIdx.x % 4) * 2;
|
||||
if (slice_id == 0) {
|
||||
for (int ax0_0_1 = 0; ax0_0_1 < WARP_M / INTRIN_M; ++ax0_0_1) {
|
||||
int row_wb_1 = row_wb_thd + ax0_0_1 * OP_M;
|
||||
for (int ax1_0_1 = 0; ax1_0_1 < WARP_N / INTRIN_N; ++ax1_0_1) {
|
||||
int col_wb_1 = col_wb_thd + ax1_0_1 * 16;
|
||||
int* C_warp_local = C_warp + ax0_0_1 * WARP_N / INTRIN_N * 8 + ax1_0_1 * 8;
|
||||
for (int local_id = 0; local_id < OP_M * 16 / WARP_SIZE; local_id += 2) {
|
||||
int row_wb = row_wb_1 + (local_id % 4) / 2 * 8;
|
||||
if (row_wb < M) {
|
||||
int col_wb = col_wb_1 + (local_id / 4) * 8 + (local_id % 2);
|
||||
float2 wscale = __half22float2(*(wscales + col_wb / 2));
|
||||
float ascale = __half2float(ascales[row_wb]);
|
||||
float2 psums =
|
||||
make_float2(__int2float_rn(C_warp_local[local_id]), __int2float_rn(C_warp_local[local_id + 1]));
|
||||
psums.x *= wscale.x * ascale;
|
||||
psums.y *= wscale.y * ascale;
|
||||
*reinterpret_cast<half2*>(C + row_wb * N + col_wb) = __float22half2_rn(psums);
|
||||
}
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#else
|
||||
template <int CTA_M, int CTA_N, int CTA_K, int WARP_M, int WARP_N, int WARP_K, int STAGES, int G>
|
||||
__global__ void dense_kernel0(
|
||||
int8_t* __restrict__ A,
|
||||
int8_t* __restrict__ B,
|
||||
int8_t* __restrict__ zeros,
|
||||
int8_t* __restrict__ scales_i8,
|
||||
half2* __restrict__ wscales,
|
||||
half* __restrict__ ascales,
|
||||
half* __restrict__ C,
|
||||
int M,
|
||||
int64_t N,
|
||||
int64_t K) {
|
||||
// Not implemented for SM < 800
|
||||
assert(false);
|
||||
return;
|
||||
}
|
||||
#endif
|
||||
|
||||
void qserve_w4a8_per_group_gemm(
|
||||
const torch::Tensor& _in_feats,
|
||||
const torch::Tensor& _kernel,
|
||||
const torch::Tensor& _zeros,
|
||||
const torch::Tensor& _scales_i8,
|
||||
const torch::Tensor& _wscales,
|
||||
const torch::Tensor& _ascales,
|
||||
torch::Tensor& _out_feats) {
|
||||
// Check input tensor
|
||||
TORCH_CHECK(_in_feats.is_cuda(), "_in_feats must be a CUDA tensor");
|
||||
TORCH_CHECK(_in_feats.dim() == 2, "_in_feats must be a 2D tensor");
|
||||
TORCH_CHECK(_in_feats.is_contiguous(), "_in_feats must be contiguous");
|
||||
TORCH_CHECK(_in_feats.scalar_type() == torch::kInt8, "_in_feats must be int8");
|
||||
// Check kernel tensor
|
||||
TORCH_CHECK(_kernel.is_cuda(), "_kernel must be a CUDA tensor");
|
||||
TORCH_CHECK(_kernel.dim() == 2, "_kernel must be a 2D tensor");
|
||||
TORCH_CHECK(_kernel.is_contiguous(), "_kernel must be contiguous");
|
||||
TORCH_CHECK(_kernel.scalar_type() == torch::kInt8, "_kernel must be int8");
|
||||
// Check output tensor
|
||||
TORCH_CHECK(_out_feats.is_cuda(), "_out_feats must be a CUDA tensor");
|
||||
TORCH_CHECK(_out_feats.is_contiguous(), "_out_feats must be contiguous");
|
||||
TORCH_CHECK(_out_feats.scalar_type() == torch::kHalf, "_out_feats must be half");
|
||||
|
||||
int num_in_feats = _in_feats.size(0);
|
||||
int num_in_channels = _in_feats.size(1);
|
||||
int num_out_feats = _out_feats.size(-2);
|
||||
int num_out_channels = _out_feats.size(-1);
|
||||
|
||||
// Check matmul shape
|
||||
TORCH_CHECK(num_out_channels == _kernel.size(0), "num_out_channels must be equal to _kernel.size(0)");
|
||||
TORCH_CHECK(num_in_feats == num_out_feats, "num_in_feats must be equal to num_out_feats");
|
||||
|
||||
// Check _ascales
|
||||
TORCH_CHECK(_ascales.is_cuda(), "_ascales must be a CUDA tensor");
|
||||
TORCH_CHECK(_ascales.is_contiguous(), "_ascales must be contiguous");
|
||||
TORCH_CHECK(_ascales.scalar_type() == torch::kHalf, "_ascales must be half");
|
||||
TORCH_CHECK(_ascales.numel() == num_in_feats, "_ascales must have num_in_feats elements");
|
||||
|
||||
// Check _wscales
|
||||
TORCH_CHECK(_wscales.is_cuda(), "_wscales must be a CUDA tensor");
|
||||
TORCH_CHECK(_wscales.is_contiguous(), "_wscales must be contiguous");
|
||||
TORCH_CHECK(_wscales.scalar_type() == torch::kHalf, "_wscales must be half");
|
||||
TORCH_CHECK(_wscales.numel() == num_out_channels, "_wscales must have num_out_channels elements");
|
||||
|
||||
// Check _scales_i8
|
||||
TORCH_CHECK(_scales_i8.is_cuda(), "_scales_i8 must be a CUDA tensor");
|
||||
TORCH_CHECK(_scales_i8.dim() == 2, "_scales_i8 must be a 2D tensor");
|
||||
TORCH_CHECK(_scales_i8.is_contiguous(), "_scales_i8 must be contiguous");
|
||||
TORCH_CHECK(_scales_i8.scalar_type() == torch::kInt8, "_scales_i8 must be int8");
|
||||
TORCH_CHECK(num_in_channels % _scales_i8.size(0) == 0, "num_in_channels must be divisible by _scales_i8.size(0)");
|
||||
TORCH_CHECK(num_out_channels == _scales_i8.size(1), "num_out_channels must be equal to _scales_i8.size(1)");
|
||||
|
||||
// Check _zeros
|
||||
TORCH_CHECK(_zeros.is_cuda(), "_zeros must be a CUDA tensor");
|
||||
TORCH_CHECK(_zeros.dim() == 2, "_zeros must be a 2D tensor");
|
||||
TORCH_CHECK(_zeros.is_contiguous(), "_zeros must be contiguous");
|
||||
TORCH_CHECK(_zeros.scalar_type() == torch::kInt8, "_zeros must be int8");
|
||||
TORCH_CHECK(num_in_channels % _zeros.size(0) == 0, "num_in_channels must be divisible by _zeros.size(0)");
|
||||
TORCH_CHECK(num_out_channels == _zeros.size(1), "num_out_channels must be equal to _zeros.size(1)");
|
||||
|
||||
// Check group size
|
||||
auto group_size = num_in_channels / _scales_i8.size(0);
|
||||
TORCH_CHECK(group_size == 128, "group_size must be 128");
|
||||
|
||||
auto in_feats = reinterpret_cast<int8_t*>(_in_feats.data_ptr<int8_t>());
|
||||
auto kernel = reinterpret_cast<int8_t*>(_kernel.data_ptr<int8_t>());
|
||||
auto zeros = reinterpret_cast<int8_t*>(_zeros.data_ptr<int8_t>());
|
||||
auto scales_i8 = reinterpret_cast<int8_t*>(_scales_i8.data_ptr<int8_t>());
|
||||
auto wscales = reinterpret_cast<half2*>(_wscales.data_ptr());
|
||||
auto ascales = reinterpret_cast<half*>(_ascales.data_ptr());
|
||||
// auto options =
|
||||
// torch::TensorOptions().dtype(torch::kHalf).device(_in_feats.device());
|
||||
auto out_feats = reinterpret_cast<half*>(_out_feats.data_ptr<at::Half>());
|
||||
auto stream = at::cuda::getCurrentCUDAStream(_in_feats.get_device());
|
||||
auto sm_version = getSMVersion();
|
||||
if (sm_version >= 80) {
|
||||
constexpr int G = 128;
|
||||
|
||||
if (num_out_feats > 128) {
|
||||
constexpr int CTA_M = 128;
|
||||
constexpr int CTA_N = 64;
|
||||
constexpr int CTA_K = 64;
|
||||
constexpr int WARP_M = 64;
|
||||
constexpr int WARP_N = 32;
|
||||
constexpr int WARP_K = 64;
|
||||
constexpr int STAGES = 4;
|
||||
KERNEL_LAUNCH_CODE
|
||||
} else if (num_out_feats >= 128) {
|
||||
if (num_in_channels <= 4096) {
|
||||
constexpr int CTA_M = 64;
|
||||
constexpr int CTA_N = 64;
|
||||
constexpr int CTA_K = 64;
|
||||
constexpr int WARP_M = 32;
|
||||
constexpr int WARP_N = 32;
|
||||
constexpr int WARP_K = 64;
|
||||
constexpr int STAGES = 4;
|
||||
KERNEL_LAUNCH_CODE
|
||||
} else {
|
||||
constexpr int CTA_M = 64;
|
||||
constexpr int CTA_N = 64;
|
||||
constexpr int CTA_K = 128;
|
||||
constexpr int WARP_M = 32;
|
||||
constexpr int WARP_N = 32;
|
||||
constexpr int WARP_K = 64;
|
||||
constexpr int STAGES = 3;
|
||||
KERNEL_LAUNCH_CODE
|
||||
}
|
||||
} else {
|
||||
constexpr int CTA_M = 32;
|
||||
constexpr int CTA_N = 64;
|
||||
constexpr int CTA_K = 128;
|
||||
constexpr int WARP_M = 32;
|
||||
constexpr int WARP_N = 32;
|
||||
constexpr int WARP_K = 64;
|
||||
constexpr int STAGES = 3;
|
||||
KERNEL_LAUNCH_CODE
|
||||
}
|
||||
} else {
|
||||
TORCH_CHECK_NOT_IMPLEMENTED(
|
||||
false, "No implemented qserve_w4a8_per_group_gemm for current compute capability: ", sm_version);
|
||||
}
|
||||
return;
|
||||
}
|
||||
@@ -726,27 +726,6 @@ void convert_vertical_slash_indexes_mergehead(
|
||||
*/
|
||||
void ApplyTokenBitmaskInplace(at::Tensor logits, at::Tensor bitmask, at::optional<at::Tensor> indices = at::nullopt);
|
||||
|
||||
/*
|
||||
* From csrc/gemm (QServe)
|
||||
*/
|
||||
void qserve_w4a8_per_chn_gemm(
|
||||
const torch::Tensor& _in_feats,
|
||||
const torch::Tensor& _kernel,
|
||||
const torch::Tensor& _wscales,
|
||||
const torch::Tensor& _ascales,
|
||||
const torch::Tensor& _w_szs,
|
||||
const torch::Tensor& _a_ssums,
|
||||
torch::Tensor& _out_feats);
|
||||
|
||||
void qserve_w4a8_per_group_gemm(
|
||||
const torch::Tensor& _in_feats,
|
||||
const torch::Tensor& _kernel,
|
||||
const torch::Tensor& _zeros,
|
||||
const torch::Tensor& _scales_i8,
|
||||
const torch::Tensor& _wscales,
|
||||
const torch::Tensor& _ascales,
|
||||
torch::Tensor& _out_feats);
|
||||
|
||||
/*
|
||||
* From csrc/quantization/gguf
|
||||
*/
|
||||
|
||||
@@ -61,8 +61,6 @@ else:
|
||||
gptq_gemm,
|
||||
gptq_shuffle,
|
||||
int8_scaled_mm,
|
||||
qserve_w4a8_per_chn_gemm,
|
||||
qserve_w4a8_per_group_gemm,
|
||||
sgl_per_token_group_quant_8bit,
|
||||
sgl_per_token_group_quant_fp8,
|
||||
sgl_per_token_group_quant_int8,
|
||||
@@ -192,8 +190,6 @@ else:
|
||||
"moe_sum",
|
||||
"moe_sum_reduce",
|
||||
"prepare_moe_input",
|
||||
"qserve_w4a8_per_chn_gemm",
|
||||
"qserve_w4a8_per_group_gemm",
|
||||
"reconstruct_indices_from_tree_mask",
|
||||
"rmsnorm",
|
||||
"rotary_embedding",
|
||||
|
||||
@@ -137,50 +137,6 @@ def sgl_per_token_quant_fp8(
|
||||
torch.ops.sgl_kernel.sgl_per_token_quant_fp8.default(input, output_q, output_s)
|
||||
|
||||
|
||||
def qserve_w4a8_per_chn_gemm(
|
||||
in_feats: torch.Tensor,
|
||||
kernel: torch.Tensor,
|
||||
wscales: torch.Tensor,
|
||||
ascales: torch.Tensor,
|
||||
w_szs: torch.Tensor,
|
||||
a_ssums: torch.Tensor,
|
||||
out_feats: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
if out_feats is None:
|
||||
# NOTE(HandH1998): qserve_w4a8_per_chn_gemm only supports out dtype=torch.float16 now
|
||||
out_feats = torch.empty(
|
||||
(in_feats.shape[0], kernel.shape[0]),
|
||||
device=in_feats.device,
|
||||
dtype=torch.float16,
|
||||
)
|
||||
torch.ops.sgl_kernel.qserve_w4a8_per_chn_gemm.default(
|
||||
in_feats, kernel, wscales, ascales, w_szs, a_ssums, out_feats
|
||||
)
|
||||
return out_feats
|
||||
|
||||
|
||||
def qserve_w4a8_per_group_gemm(
|
||||
in_feats: torch.Tensor,
|
||||
kernel: torch.Tensor,
|
||||
zeros: torch.Tensor,
|
||||
scales_i8: torch.Tensor,
|
||||
wscales: torch.Tensor,
|
||||
ascales: torch.Tensor,
|
||||
out_feats: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
if out_feats is None:
|
||||
# NOTE(HandH1998): qserve_w4a8_per_group_gemm only supports out dtype=torch.float16 now
|
||||
out_feats = torch.empty(
|
||||
(in_feats.shape[0], kernel.shape[0]),
|
||||
device=in_feats.device,
|
||||
dtype=torch.float16,
|
||||
)
|
||||
torch.ops.sgl_kernel.qserve_w4a8_per_group_gemm.default(
|
||||
in_feats, kernel, zeros, scales_i8, wscales, ascales, out_feats
|
||||
)
|
||||
return out_feats
|
||||
|
||||
|
||||
def shuffle_rows(input_tensor, dst2src_map, output_tensor_shape):
|
||||
output_tensor = torch.empty(
|
||||
output_tensor_shape,
|
||||
|
||||
@@ -1,120 +0,0 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel import qserve_w4a8_per_chn_gemm
|
||||
|
||||
|
||||
# Adapted from https://github.com/mit-han-lab/omniserve/blob/main/omniserve/modeling/layers/quantized_linear/w4a8_linear.py
|
||||
def convert_to_qserve_format(qweight, scale, zero):
|
||||
assert qweight.min() >= 0 and qweight.max() <= 15, "Quantized weight out of range"
|
||||
in_features = qweight.shape[1]
|
||||
out_features = qweight.shape[0]
|
||||
assert in_features % 32 == 0, "Input features must be divisible by 32"
|
||||
assert out_features % 32 == 0, "Output features must be divisible by 32"
|
||||
|
||||
# ---- Repack the weight ---- #
|
||||
# pack to M // 32, K // 32, (8, 4), ([2], 2, 2, 4)
|
||||
qweight_unpack_reorder = (
|
||||
qweight.reshape(
|
||||
out_features // 32,
|
||||
2,
|
||||
2,
|
||||
8,
|
||||
in_features // 32,
|
||||
2,
|
||||
4,
|
||||
4,
|
||||
)
|
||||
.permute(0, 4, 3, 6, 1, 5, 2, 7)
|
||||
.contiguous()
|
||||
)
|
||||
qweight_unpack_reorder = (
|
||||
qweight_unpack_reorder.permute(0, 1, 2, 3, 5, 6, 7, 4)
|
||||
.contiguous()
|
||||
.to(torch.int8)
|
||||
)
|
||||
# B_fp16_reorder = B_fp16_reorder[:, :, :, :, :, :, [3, 2, 1, 0]].contiguous()
|
||||
# [16, 0, 17, 1, ...]
|
||||
qweight_unpack_repacked = (
|
||||
qweight_unpack_reorder[..., 1] << 4
|
||||
) + qweight_unpack_reorder[..., 0]
|
||||
qweight_unpack_repacked = qweight_unpack_repacked.reshape(
|
||||
out_features // 32, in_features // 32, 32, 16
|
||||
)
|
||||
qweight_unpack_repacked = qweight_unpack_repacked.reshape(
|
||||
out_features, in_features // 2
|
||||
).contiguous()
|
||||
|
||||
# ---- Pack the scales ---- #
|
||||
scale = scale.reshape(out_features).to(torch.float16).contiguous()
|
||||
szero = zero.reshape(out_features).to(torch.float16).contiguous() * scale
|
||||
|
||||
return qweight_unpack_repacked, scale, szero
|
||||
|
||||
|
||||
# INT4 Quantization
|
||||
def asym_quantize_tensor(tensor):
|
||||
tensor_min = tensor.min(dim=-1, keepdim=True)[0]
|
||||
tensor_max = tensor.max(dim=-1, keepdim=True)[0]
|
||||
q_min = 0
|
||||
q_max = 15
|
||||
tensor_scale = (tensor_max - tensor_min) / (q_max - q_min)
|
||||
tensor_zero = q_min - torch.round(tensor_min / tensor_scale)
|
||||
tensor_q = torch.clamp(
|
||||
torch.round(tensor / tensor_scale) + tensor_zero, q_min, q_max
|
||||
).to(torch.int8)
|
||||
return tensor_q, tensor_scale.to(torch.float16), tensor_zero.to(torch.int8)
|
||||
|
||||
|
||||
# INT8 Quantization
|
||||
def sym_quantize_tensor(tensor):
|
||||
tensor_scale = tensor.abs().max(dim=-1, keepdim=True)[0] / 127
|
||||
tensor_q = torch.clamp(torch.round(tensor / tensor_scale), -128, 127).to(torch.int8)
|
||||
return tensor_q, tensor_scale.to(torch.float16)
|
||||
|
||||
|
||||
def torch_w4a8_per_chn_gemm(a, b, a_scale, b_scale, b_zero, out_dtype):
|
||||
print(a.shape)
|
||||
print(b.shape)
|
||||
print(b_zero.shape)
|
||||
o = torch.matmul(
|
||||
a.to(torch.float16), (b.to(torch.float16) - b_zero.to(torch.float16)).t()
|
||||
)
|
||||
o = o * a_scale.view(-1, 1) * b_scale.view(1, -1)
|
||||
return o.to(out_dtype)
|
||||
|
||||
|
||||
def _test_accuracy_once(M, N, K, out_dtype, device):
|
||||
# to avoid overflow, multiply 0.01
|
||||
a = torch.randn((M, K), device=device, dtype=torch.float32) * 0.01
|
||||
b = torch.randn((N, K), device=device, dtype=torch.float32) * 0.01
|
||||
|
||||
# symmetric quantize a
|
||||
a_q, a_scale = sym_quantize_tensor(a)
|
||||
# asymmetric quantize b
|
||||
b_q, b_scale, b_zero = asym_quantize_tensor(b)
|
||||
# convert to qserve format
|
||||
b_q_format, b_scale_format, b_szero_format = convert_to_qserve_format(
|
||||
b_q, b_scale, b_zero
|
||||
)
|
||||
|
||||
# cal sum of every row of a
|
||||
a_sum = a.sum(dim=-1, keepdim=True).to(torch.float16)
|
||||
out = qserve_w4a8_per_chn_gemm(
|
||||
a_q, b_q_format, b_scale_format, a_scale, b_szero_format, a_sum
|
||||
)
|
||||
ref_out = torch_w4a8_per_chn_gemm(a_q, b_q, a_scale, b_scale, b_zero, out_dtype)
|
||||
torch.testing.assert_close(out, ref_out, rtol=1e-3, atol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("M", [1, 16, 32, 64, 128, 512, 1024, 4096, 8192])
|
||||
@pytest.mark.parametrize("N", [128, 512, 1024, 4096, 8192, 16384])
|
||||
@pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 16384])
|
||||
@pytest.mark.parametrize("out_dtype", [torch.float16])
|
||||
def test_accuracy(M, N, K, out_dtype):
|
||||
_test_accuracy_once(M, N, K, out_dtype, "cuda")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
@@ -1,185 +0,0 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from sgl_kernel import qserve_w4a8_per_group_gemm
|
||||
|
||||
|
||||
# Adapted from https://github.com/mit-han-lab/omniserve/blob/main/omniserve/modeling/layers/quantized_linear/w4a8_linear.py
|
||||
def convert_to_qserve_format(qweight, chn_scale, scale_i8, zero_i8, group_size):
|
||||
assert qweight.min() >= 0 and qweight.max() <= 15, "Quantized weight out of range"
|
||||
in_features = qweight.shape[1]
|
||||
out_features = qweight.shape[0]
|
||||
assert in_features % 32 == 0, "Input features must be divisible by 32"
|
||||
assert out_features % 32 == 0, "Output features must be divisible by 32"
|
||||
assert group_size == 128, "Group size must be 128"
|
||||
assert (
|
||||
in_features % group_size == 0
|
||||
), "Input features must be divisible by group_size"
|
||||
|
||||
# ---- Repack the weight ---- #
|
||||
# pack to M // 32, K // 32, (8, 4), ([2], 2, 2, 4)
|
||||
qweight_unpack_reorder = (
|
||||
qweight.reshape(
|
||||
out_features // 32,
|
||||
2,
|
||||
2,
|
||||
8,
|
||||
in_features // 32,
|
||||
2,
|
||||
4,
|
||||
4,
|
||||
)
|
||||
.permute(0, 4, 3, 6, 1, 5, 2, 7)
|
||||
.contiguous()
|
||||
)
|
||||
qweight_unpack_reorder = (
|
||||
qweight_unpack_reorder.permute(0, 1, 2, 3, 5, 6, 7, 4)
|
||||
.contiguous()
|
||||
.to(torch.int8)
|
||||
)
|
||||
# B_fp16_reorder = B_fp16_reorder[:, :, :, :, :, :, [3, 2, 1, 0]].contiguous()
|
||||
# [16, 0, 17, 1, ...]
|
||||
qweigth_unpack_repacked = (
|
||||
qweight_unpack_reorder[..., 1] << 4
|
||||
) + qweight_unpack_reorder[..., 0]
|
||||
qweigth_unpack_repacked = qweigth_unpack_repacked.reshape(
|
||||
out_features // 32, in_features // 32, 32, 16
|
||||
)
|
||||
qweigth_unpack_repacked = qweigth_unpack_repacked.reshape(
|
||||
out_features, in_features // 2
|
||||
)
|
||||
|
||||
# ---- Pack the scales ---- #
|
||||
chn_scale = chn_scale.reshape(out_features)
|
||||
|
||||
scale_i8 = (
|
||||
scale_i8.reshape(out_features, in_features // group_size)
|
||||
.transpose(0, 1)
|
||||
.contiguous()
|
||||
)
|
||||
scale_i8 = scale_i8.reshape(in_features // group_size, out_features // 32, 32)
|
||||
scale_i8 = (
|
||||
scale_i8.reshape(in_features // group_size, out_features // 32, 4, 8)
|
||||
.transpose(-2, -1)
|
||||
.contiguous()
|
||||
)
|
||||
scale_i8 = scale_i8.reshape(in_features // group_size, out_features).contiguous()
|
||||
|
||||
# ---- Pack the zeros ---- #
|
||||
zero_i8 = -zero_i8
|
||||
# zero_i8 = zero_i8.int() # convert to 2-complement
|
||||
|
||||
zero_i8 = (
|
||||
zero_i8.reshape(out_features, in_features // group_size)
|
||||
.transpose(0, 1)
|
||||
.contiguous()
|
||||
)
|
||||
zero_i8 = zero_i8.reshape(in_features // group_size, out_features // 32, 32)
|
||||
# for the last dimension, organize as 0, 8, 16, 24, 1, 9, 17, 25, ... following the requirement of tensor core gemm
|
||||
zero_i8 = (
|
||||
zero_i8.reshape(in_features // group_size, out_features // 32, 4, 8)
|
||||
.transpose(-2, -1)
|
||||
.contiguous()
|
||||
)
|
||||
zero_i8 = (
|
||||
zero_i8.reshape(in_features // group_size, out_features).contiguous() * scale_i8
|
||||
)
|
||||
|
||||
return qweigth_unpack_repacked, chn_scale, scale_i8, zero_i8
|
||||
|
||||
|
||||
# Progressive Group INT4 Quantization
|
||||
def progressive_group_quantize_tensor(tensor, group_size):
|
||||
assert group_size == 128, "Group size must be 128"
|
||||
assert (
|
||||
tensor.shape[-1] % group_size == 0
|
||||
), "Input features must be divisible by group_size"
|
||||
# Channel scale
|
||||
# NOTE(HandH1998): use protective quantization range
|
||||
chn_scale = tensor.abs().max(dim=-1, keepdim=True)[0] / 119
|
||||
tensor_i8 = torch.clamp(torch.round(tensor / chn_scale), -119, 119)
|
||||
|
||||
# Group scale
|
||||
tensor_i8 = tensor_i8.reshape(-1, group_size)
|
||||
tensor_i8_min = tensor_i8.min(dim=-1, keepdim=True)[0]
|
||||
tensor_i8_max = tensor_i8.max(dim=-1, keepdim=True)[0]
|
||||
q_min = 0
|
||||
q_max = 15
|
||||
scale_i8 = torch.round((tensor_i8_max - tensor_i8_min) / (q_max - q_min))
|
||||
zero_i8 = q_min - torch.round(tensor_i8_min / scale_i8)
|
||||
tensor_q = (
|
||||
torch.clamp(torch.round(tensor_i8 / scale_i8) + zero_i8, q_min, q_max)
|
||||
.reshape(tensor.shape[0], -1)
|
||||
.to(torch.int8)
|
||||
)
|
||||
return (
|
||||
tensor_q,
|
||||
chn_scale.to(torch.float16),
|
||||
scale_i8.reshape(tensor.shape[0], -1).to(torch.int8),
|
||||
zero_i8.reshape(tensor.shape[0], -1).to(torch.int8),
|
||||
)
|
||||
|
||||
|
||||
# INT8 Quantization
|
||||
def sym_quantize_tensor(tensor):
|
||||
tensor_scale = tensor.abs().max(dim=-1, keepdim=True)[0] / 127
|
||||
tensor_q = torch.clamp(torch.round(tensor / tensor_scale), -128, 127).to(torch.int8)
|
||||
return tensor_q, tensor_scale.to(torch.float16)
|
||||
|
||||
|
||||
def torch_w4a8_per_group_gemm(
|
||||
a, b, a_scale, b_chn_scale, b_scale_i8, b_zero_i8, group_size, out_dtype
|
||||
):
|
||||
assert group_size == 128, "Group size must be 128"
|
||||
b_dq = (
|
||||
b.reshape(-1, group_size).to(torch.float32)
|
||||
- b_zero_i8.reshape(-1, 1).to(torch.float32)
|
||||
) * b_scale_i8.reshape(-1, 1).to(torch.float32)
|
||||
b_dq = b_dq.reshape(b.shape[0], b.shape[1])
|
||||
o = torch.matmul(a.to(torch.float32), b_dq.t())
|
||||
o = o * a_scale.view(-1, 1) * b_chn_scale.view(1, -1)
|
||||
return o.to(out_dtype)
|
||||
|
||||
|
||||
def _test_accuracy_once(M, N, K, group_size, out_dtype, device):
|
||||
# to avoid overflow, multiply 0.01
|
||||
a = torch.randn((M, K), device=device, dtype=torch.float32) * 0.01
|
||||
b = torch.randn((N, K), device=device, dtype=torch.float32) * 0.01
|
||||
|
||||
# symmetric quantize a
|
||||
a_q, a_scale = sym_quantize_tensor(a)
|
||||
# asymmetric quantize b
|
||||
b_q, b_chn_scale, b_scale_i8, b_zero_i8 = progressive_group_quantize_tensor(
|
||||
b, group_size
|
||||
)
|
||||
# convert to qserve format
|
||||
b_q_format, b_chn_scale_format, b_scale_i8_format, b_zero_i8_format = (
|
||||
convert_to_qserve_format(b_q, b_chn_scale, b_scale_i8, b_zero_i8, group_size)
|
||||
)
|
||||
|
||||
out = qserve_w4a8_per_group_gemm(
|
||||
a_q,
|
||||
b_q_format,
|
||||
b_zero_i8_format,
|
||||
b_scale_i8_format,
|
||||
b_chn_scale_format,
|
||||
a_scale,
|
||||
)
|
||||
ref_out = torch_w4a8_per_group_gemm(
|
||||
a_q, b_q, a_scale, b_chn_scale, b_scale_i8, b_zero_i8, group_size, out_dtype
|
||||
)
|
||||
torch.testing.assert_close(out, ref_out, rtol=1e-3, atol=1e-5)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("M", [1, 16, 32, 64, 128, 512, 1024, 4096, 8192])
|
||||
@pytest.mark.parametrize("N", [128, 512, 1024, 4096, 8192, 16384])
|
||||
@pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 16384])
|
||||
@pytest.mark.parametrize("group_size", [128])
|
||||
@pytest.mark.parametrize("out_dtype", [torch.float16])
|
||||
def test_accuracy(M, N, K, group_size, out_dtype):
|
||||
_test_accuracy_once(M, N, K, group_size, out_dtype, "cuda")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__]))
|
||||
Reference in New Issue
Block a user