[FP8] SM120: route FP8 linear to per-tensor (cudnn/nvjet) instead of channelwise cutlass (#38006)
Co-authored-by: BBuf <bbuf@example.com>
This commit is contained in:
@@ -87,6 +87,7 @@ from sglang.srt.runtime_context import (
|
|||||||
from sglang.srt.utils import (
|
from sglang.srt.utils import (
|
||||||
cpu_has_amx_support,
|
cpu_has_amx_support,
|
||||||
get_bool_env_var,
|
get_bool_env_var,
|
||||||
|
get_device_capability,
|
||||||
is_cpu,
|
is_cpu,
|
||||||
is_cuda,
|
is_cuda,
|
||||||
is_flashinfer_available,
|
is_flashinfer_available,
|
||||||
@@ -911,9 +912,20 @@ class Fp8LinearMethod(LinearMethodBase):
|
|||||||
layer.input_scale.data, requires_grad=False
|
layer.input_scale.data, requires_grad=False
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# On SM120 (Blackwell RTX 50) the per-tensor FP8 path dispatches
|
||||||
|
# to the fast cudnn/nvjet SM120 kernel, which beats the
|
||||||
|
# channelwise cutlass GemmUniversal by ~1.2-2.7x across the
|
||||||
|
# decode/prefill M range (matches vLLM's per-tensor SM120
|
||||||
|
# choice). Requantize per-channel -> per-tensor (max scale)
|
||||||
|
# there instead.
|
||||||
|
use_sm120_fp8_pertensor = (
|
||||||
|
self.cutlass_fp8_supported
|
||||||
|
and not self.use_marlin
|
||||||
|
and get_device_capability()[0] == 12
|
||||||
|
)
|
||||||
# cutlass sgl-kernel and marlin only support per-channel scale; aiter supports per-channel scale
|
# cutlass sgl-kernel and marlin only support per-channel scale; aiter supports per-channel scale
|
||||||
if (
|
if (
|
||||||
self.cutlass_fp8_supported
|
(self.cutlass_fp8_supported and not use_sm120_fp8_pertensor)
|
||||||
or self.use_marlin
|
or self.use_marlin
|
||||||
or (_use_aiter and self.use_aiter_fp8_per_token)
|
or (_use_aiter and self.use_aiter_fp8_per_token)
|
||||||
):
|
):
|
||||||
|
|||||||
Reference in New Issue
Block a user