Support CuteDSL mm_fp4 backend (#18801)

This commit is contained in:
Brayden Zhong
2026-03-19 17:20:01 -04:00
committed by GitHub
parent d8ece7fb22
commit b42b9f6e1a
5 changed files with 109 additions and 219 deletions

View File

@@ -18,6 +18,7 @@ class Fp4GemmRunnerBackend(Enum):
AUTO = "auto"
FLASHINFER_CUDNN = "flashinfer_cudnn"
FLASHINFER_CUTEDSL = "flashinfer_cutedsl"
FLASHINFER_CUTLASS = "flashinfer_cutlass"
FLASHINFER_TRTLLM = "flashinfer_trtllm"
@@ -33,6 +34,9 @@ class Fp4GemmRunnerBackend(Enum):
def is_flashinfer_trtllm(self) -> bool:
return self == Fp4GemmRunnerBackend.FLASHINFER_TRTLLM
def is_flashinfer_cutedsl(self) -> bool:
return self == Fp4GemmRunnerBackend.FLASHINFER_CUTEDSL
def get_flashinfer_backend(self) -> str:
"""Get the backend string to pass to FlashInfer's mm_fp4 API.
@@ -41,7 +45,10 @@ class Fp4GemmRunnerBackend(Enum):
'flashinfer_trtllm' -> 'trtllm'
'flashinfer_cutlass' -> 'cutlass'
'flashinfer_cudnn' -> 'cudnn'
'flashinfer_cutedsl' -> 'cute-dsl'
"""
if self == Fp4GemmRunnerBackend.FLASHINFER_CUTEDSL:
return "cute-dsl"
if self.value.startswith("flashinfer_"):
return self.value.removeprefix("flashinfer_")
else: