[CPU][INT4] Add INT4 kernels for CPU (#8226)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
co-authored by
gemini-code-assist[bot]
parent
88f7759402
commit
c35aa0238c
@@ -15,7 +15,10 @@ from sglang.srt.distributed.device_communicators.pynccl_allocator import (
|
||||
use_symmetric_memory,
|
||||
)
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
|
||||
from sglang.srt.layers.amx_utils import (
|
||||
CPUQuantMethod,
|
||||
_amx_process_weight_after_loading,
|
||||
)
|
||||
from sglang.srt.layers.dp_attention import is_allocation_symmetric
|
||||
from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.deep_gemm import DeepGemmMoeQuantInfo
|
||||
@@ -1355,13 +1358,12 @@ class Fp8MoEMethod(FusedMoEMethodBase):
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
False, # inplace See [Note] inplace should be False in fused_experts.
|
||||
False, # use_int8_w8a8
|
||||
True, # use_fp8_w8a16
|
||||
CPUQuantMethod.FP8_W8A16,
|
||||
layer.w13_weight_scale_inv, # w1_scale
|
||||
layer.w2_weight_scale_inv, # w2_scale
|
||||
None, # w1_zp
|
||||
None, # w2_zp
|
||||
self.quant_config.weight_block_size, # block_size
|
||||
None, # a1_scale
|
||||
None, # a2_scale
|
||||
True, # is_vnni
|
||||
)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
@@ -6,7 +6,10 @@ import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
|
||||
from sglang.srt.layers.amx_utils import (
|
||||
CPUQuantMethod,
|
||||
_amx_process_weight_after_loading,
|
||||
)
|
||||
from sglang.srt.layers.moe import (
|
||||
MoeRunner,
|
||||
MoeRunnerBackend,
|
||||
@@ -447,13 +450,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
False, # inplace # See [Note] inplace should be False in fused_experts.
|
||||
False, # use_int8_w8a8
|
||||
False, # use_fp8_w8a16
|
||||
CPUQuantMethod.UNQUANT,
|
||||
None, # w1_scale
|
||||
None, # w2_scale
|
||||
None, # w1_zp
|
||||
None, # w2_zp
|
||||
None, # block_size
|
||||
None, # a1_scale
|
||||
None, # a2_scale
|
||||
True, # is_vnni
|
||||
)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
@@ -8,7 +8,10 @@ import torch
|
||||
from torch.nn.parameter import Parameter
|
||||
|
||||
from sglang.srt.distributed import get_tensor_model_parallel_world_size
|
||||
from sglang.srt.layers.amx_utils import _amx_process_weight_after_loading
|
||||
from sglang.srt.layers.amx_utils import (
|
||||
CPUQuantMethod,
|
||||
_amx_process_weight_after_loading,
|
||||
)
|
||||
from sglang.srt.layers.moe import MoeRunner, MoeRunnerBackend, MoeRunnerConfig
|
||||
from sglang.srt.layers.moe.moe_runner.triton import TritonMoeQuantInfo
|
||||
from sglang.srt.layers.parameter import ChannelQuantScaleParameter, ModelWeightParameter
|
||||
@@ -352,13 +355,12 @@ class W8A8Int8MoEMethod(FusedMoEMethodBase):
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
False, # inplace See [Note] inplace should be False in fused_experts.
|
||||
True, # use_int8_w8a8
|
||||
False, # use_fp8_w8a16
|
||||
CPUQuantMethod.INT8_W8A8,
|
||||
layer.w13_weight_scale, # w1_scale
|
||||
layer.w2_weight_scale, # w2_scale
|
||||
None, # w1_zp
|
||||
None, # w2_zp
|
||||
None, # block_size
|
||||
layer.w13_input_scale, # a1_scale
|
||||
layer.w2_input_scale, # a2_scale
|
||||
True, # is_vnni
|
||||
)
|
||||
return StandardCombineInput(hidden_states=output)
|
||||
|
||||
Reference in New Issue
Block a user