Unverified Commit 9ab4388c authored by Woosuk Kwon's avatar Woosuk Kwon Committed by GitHub
Browse files

[Model Runner V2] Support FLASHINFER_MLA backend (#32709)


Signed-off-by: default avatarWoosuk Kwon <woosuk.kwon@berkeley.edu>
parent 04a9e064
......@@ -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:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment