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:
Brayden Zhong
2026-07-21 14:09:01 -07:00
committed by GitHub
co-authored by Brayden Zhong
parent cb324cc958
commit 0eae9423d8
2 changed files with 4 additions and 2 deletions
@@ -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,