[Model Runner V2] Support FLASHINFER_MLA backend (#32709)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
@@ -57,7 +57,7 @@ def init_attn_backend(
|
||||
)
|
||||
attn_metadata_builders.append(attn_metadata_builder) # type: ignore
|
||||
|
||||
if "FLASHINFER" in attn_backend.get_name():
|
||||
if attn_backend.get_name() == "FLASHINFER":
|
||||
if flashinfer_workspace is None:
|
||||
flashinfer_workspace = attn_metadata_builder._get_workspace_buffer()
|
||||
else:
|
||||
|
||||
@@ -248,7 +248,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
|
||||
)
|
||||
|
||||
# TODO(woosuk): Support other backends.
|
||||
supported_backends = ("FLASH_ATTN", "FLASHINFER")
|
||||
supported_backends = ("FLASH_ATTN", "FLASHINFER", "FLASHINFER_MLA")
|
||||
for backend in self.attn_backends.values():
|
||||
backend_name = backend.get_name()
|
||||
if backend_name not in supported_backends:
|
||||
|
||||
Reference in New Issue
Block a user