Handle FlashInfer workspace size arrays

Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com>

Author: lesj0610

tests/v1/worker/test_cudagraph_memory_profiling.py

@@ -51,6 +51,17 @@ def reset_workspace_buffer(
         self.reset_calls += 1
 
 
+class _FakeWorkspaceSizeArray:
+    def __init__(self, values: list[int]) -> None:
+        self.values = values
+
+    def __len__(self) -> int:
+        return len(self.values)
+
+    def __getitem__(self, index: int) -> int:
+        return self.values[index]
+
+
 def _make_flashinfer_builder(flashinfer_backend):
     FlashInferMetadataBuilder = flashinfer_backend.FlashInferMetadataBuilder
     builder = FlashInferMetadataBuilder.__new__(FlashInferMetadataBuilder)
@@ -61,6 +72,44 @@ def _make_flashinfer_builder(flashinfer_backend):
     return builder
 
 
+def test_flashinfer_workspace_size_parses_sequence_like_array():
+    pytest.importorskip("flashinfer")
+    from vllm.v1.attention.backends import flashinfer as flashinfer_backend
+
+    WorkspaceSizes = flashinfer_backend.WorkspaceSizes
+
+    assert flashinfer_backend._parse_workspace_sizes(
+        _FakeWorkspaceSizeArray([1024, 64])
+    ) == WorkspaceSizes(1024, 64)
+
+
+def test_flashinfer_workspace_size_float_only_keeps_default_int_workspace():
+    pytest.importorskip("flashinfer")
+    from vllm.v1.attention.backends import flashinfer as flashinfer_backend
+
+    builder = _make_flashinfer_builder(flashinfer_backend)
+    wrapper = _FakeFlashInferWrapper(int_workspace_bytes=64)
+    sizes = flashinfer_backend._parse_workspace_sizes(_FakeWorkspaceSizeArray([1024]))
+
+    reset_workspace_manager()
+    init_workspace_manager(torch.device("cpu"))
+    try:
+        builder._ensure_flashinfer_wrapper_workspace(wrapper, sizes)
+    finally:
+        reset_workspace_manager()
+
+    assert wrapper._float_workspace_buffer.numel() == 1024
+    assert wrapper._int_workspace_buffer.numel() == 64
+
+
+def test_flashinfer_workspace_size_rejects_invalid_sequence_length():
+    pytest.importorskip("flashinfer")
+    from vllm.v1.attention.backends import flashinfer as flashinfer_backend
+
+    with pytest.raises(ValueError, match="workspace_size"):
+        flashinfer_backend._parse_workspace_sizes(_FakeWorkspaceSizeArray([1, 2, 3]))
+
+
 def test_flashinfer_separate_cudagraph_memory_profile_gate():
     pytest.importorskip("flashinfer")
     from vllm.v1.attention.backends.flashinfer import FlashInferMetadataBuilder

vllm/v1/attention/backends/flashinfer.py

@@ -112,10 +112,23 @@ def total_bytes(self) -> int:
 
 
 def _parse_workspace_sizes(workspace_size: Any) -> WorkspaceSizes:
-    if isinstance(workspace_size, (tuple, list)):
-        float_bytes = int(workspace_size[0])
-        int_bytes = int(workspace_size[1]) if len(workspace_size) > 1 else 0
-        return WorkspaceSizes(float_bytes, int_bytes)
+    if not isinstance(workspace_size, (str, bytes)):
+        try:
+            workspace_size_len = len(workspace_size)
+        except TypeError:
+            pass
+        else:
+            if workspace_size_len == 0:
+                raise ValueError("FlashInfer workspace_size returned an empty result")
+            if workspace_size_len > 2:
+                raise ValueError(
+                    "FlashInfer workspace_size must return a scalar or a "
+                    "(float_bytes, int_bytes) pair"
+                )
+            float_bytes = int(workspace_size[0])
+            int_bytes = int(workspace_size[1]) if workspace_size_len == 2 else 0
+            return WorkspaceSizes(float_bytes, int_bytes)
+
     return WorkspaceSizes(int(workspace_size), 0)