Fix GPU kernel ordering and MXFP8 quantization dispatch (#37331)
Co-authored-by: Pranjal Shankhdhar <pranjalssh@meta.com> Co-authored-by: Chengze Fan <fancz2002@gmail.com>
This commit is contained in:
co-authored by
Pranjal Shankhdhar
Chengze Fan
parent
221a6273ce
commit
33428d3dae
@@ -167,7 +167,10 @@ template <bool kUsePDL>
|
||||
SGL_DEVICE void PDLTriggerSecondary() {
|
||||
#if SGL_ARCH_HOPPER_OR_GREATER
|
||||
if constexpr (kUsePDL) {
|
||||
asm volatile("griddepcontrol.launch_dependents;" :::);
|
||||
// The "memory" clobber is load-bearing: without it the compiler may sink
|
||||
// this kernel's stores past the trigger, and the dependent grid's
|
||||
// griddepcontrol.wait only covers writes issued BEFORE launch_dependents.
|
||||
asm volatile("griddepcontrol.launch_dependents;" ::: "memory");
|
||||
}
|
||||
#endif
|
||||
}
|
||||
|
||||
@@ -1293,6 +1293,12 @@ def mxfp8_group_quantize(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
assert x.is_contiguous(), "MXFP8 quantization requires a contiguous 2D tensor."
|
||||
_, k = x.shape
|
||||
assert k % 32 == 0, f"{k=} must be divisible by 32"
|
||||
if _is_hip and _is_gfx95_supported:
|
||||
from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import (
|
||||
mxfp8_e4m3_quantize,
|
||||
)
|
||||
|
||||
return mxfp8_e4m3_quantize(x)
|
||||
downcast_to_mxfp = _get_triton_mxfp8_downcast()
|
||||
q_input, scale_u8 = downcast_to_mxfp(x, torch.float8_e4m3fn, axis=1)
|
||||
return q_input.contiguous(), scale_u8.contiguous()
|
||||
|
||||
Reference in New Issue
Block a user