Change the FP8 per-tensor GEMM backend on SM120 to cuBLAS (#31961)
Co-authored-by: Brayden Zhong <brayden.zhong@radixark.ai>
This commit is contained in:
co-authored by
Brayden Zhong
parent
cb324cc958
commit
0eae9423d8
@@ -1710,7 +1710,7 @@ def apply_fp8_linear_bmm_flashinfer(
|
|||||||
input_scale: torch.Tensor,
|
input_scale: torch.Tensor,
|
||||||
bias: Optional[torch.Tensor] = None,
|
bias: Optional[torch.Tensor] = None,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
"""Per-tensor static fp8 linear via flashinfer bmm_fp8 (SM10X only)."""
|
"""Per-tensor static fp8 linear via flashinfer bmm_fp8 (SM100/SM120 Blackwell)."""
|
||||||
output_shape = [*input.shape[:-1], weight.shape[1]]
|
output_shape = [*input.shape[:-1], weight.shape[1]]
|
||||||
input_2d = input.view(-1, input.shape[-1])
|
input_2d = input.view(-1, input.shape[-1])
|
||||||
qinput, x_scale = static_quant_fp8(input_2d, input_scale, repeat_scale=False)
|
qinput, x_scale = static_quant_fp8(input_2d, input_scale, repeat_scale=False)
|
||||||
|
|||||||
@@ -499,7 +499,9 @@ class ModelOptFp8LinearMethod(LinearMethodBase):
|
|||||||
super().__init__()
|
super().__init__()
|
||||||
self.quant_config = quant_config
|
self.quant_config = quant_config
|
||||||
self.cutlass_fp8_supported = cutlass_fp8_supported()
|
self.cutlass_fp8_supported = cutlass_fp8_supported()
|
||||||
self.enable_flashinfer_bmm = is_sm100_supported() and is_flashinfer_available()
|
self.enable_flashinfer_bmm = (
|
||||||
|
is_sm100_supported() or is_sm120_supported()
|
||||||
|
) and is_flashinfer_available()
|
||||||
|
|
||||||
def create_weights(
|
def create_weights(
|
||||||
self,
|
self,
|
||||||
|
|||||||
Reference in New Issue
Block a user