[plugin] enable OOT platforms to provide custom quant configs (#25347)
Signed-off-by: Devashish Lal <devcode@fb.com> Co-authored-by: Devashish Lal <devcode@fb.com> Co-authored-by: Devashish Lal <laldevashish@gmail.com>
This commit is contained in:
co-authored by
Devashish Lal
Devashish Lal
parent
857ecb2dbc
commit
52a5c01eba
@@ -521,6 +521,11 @@ python -c "from sglang.srt.platforms import current_platform; print(current_plat
|
||||
<td><code>raise NotImplementedError</code></td>
|
||||
<td>Paged allocator class</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>get_quantization_config(quantization)</code></td>
|
||||
<td><code>raise NotImplementedError</code></td>
|
||||
<td>Return hardware-specific quantization config for the specific quantization scheme, raise an error if not supported or return None to use the default config.</td>
|
||||
</tr>
|
||||
<tr>
|
||||
<td><code>get_piecewise_backend_cls()</code></td>
|
||||
<td><code>raise NotImplementedError</code></td>
|
||||
|
||||
@@ -51,6 +51,7 @@ from sglang.srt.layers.quantization.quark_int4fp8_moe import QuarkInt4Fp8Config
|
||||
from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config
|
||||
from sglang.srt.layers.quantization.w8a8_fp8 import W8A8Fp8Config
|
||||
from sglang.srt.layers.quantization.w8a8_int8 import W8A8Int8Config
|
||||
from sglang.srt.platforms import current_platform
|
||||
from sglang.srt.utils import (
|
||||
cpu_has_amx_support,
|
||||
is_cpu,
|
||||
@@ -151,6 +152,13 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
|
||||
else:
|
||||
return CPU_QUANTIZATION_METHODS[quantization]
|
||||
|
||||
if current_platform.is_out_of_tree():
|
||||
config = current_platform.get_quantization_config(quantization)
|
||||
|
||||
# If the platform has a quantization config, use it else use the default
|
||||
if config is not None:
|
||||
return config
|
||||
|
||||
return QUANTIZATION_METHODS[quantization]
|
||||
|
||||
|
||||
|
||||
@@ -10,8 +10,15 @@ Out-of-tree platforms register via setuptools entry_points under the
|
||||
"sglang.srt.platforms" group and should subclass SRTPlatform.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional, Type
|
||||
|
||||
from sglang.srt.platforms.device_mixin import DeviceMixin, PlatformEnum
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
# Re-export for convenience
|
||||
__all__ = ["SRTPlatform", "PlatformEnum"]
|
||||
|
||||
@@ -80,6 +87,14 @@ class SRTPlatform(DeviceMixin):
|
||||
"""Return the piecewise compilation backend class for this platform."""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_quantization_config(
|
||||
self, quantization: str
|
||||
) -> Optional[Type[QuantizationConfig]]:
|
||||
"""Return hardware-specific quantization config for the specific
|
||||
quantization scheme, raise an error if not supported or return None
|
||||
to use the default config."""
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# Capability flags (safe conservative defaults)
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user