[Core][Kernel] Share persistent workspaces for Marlin and Humming (#57421)

Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: ErinYin <liang.yin@daocloud.io>

tests/kernels/quantization/test_marlin_gemm.py

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (
     per_token_quant_int8,
 )
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+    get_marlin_workspace,
     marlin_make_empty,
     marlin_make_workspace_new,
     marlin_permute_bias,
@@ -43,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
 )
 from vllm.platforms import current_platform
 from vllm.scalar_type import scalar_types
+from vllm.v1.worker.workspace import current_workspace_manager
 
 if current_platform.is_rocm():
     pytest.skip(
@@ -477,6 +479,62 @@ def test_marlin_gemm(
     assert max_diff < 0.04
 
 
+def test_marlin_persistent_locks_cuda_graph_streams(workspace_init):
+    """Compiled GEMMs select stream-local locks at runtime and survive replay."""
+    size_m, size_k, size_n = 32, 1024, 2048
+    inputs = [rand_data((size_m, size_k)) for _ in range(2)]
+    weight = rand_data((size_k, size_n))
+    quant_type = scalar_types.uint4b8
+    w_ref, qweight, scales = marlin_quantize(weight, quant_type, 128)
+    zeros = marlin_make_empty(scales.device)
+    streams = [torch.cuda.Stream() for _ in inputs]
+    graphs = [torch.cuda.CUDAGraph() for _ in inputs]
+    locks = []
+    outputs = []
+
+    @torch.compile(backend="eager", fullgraph=True)
+    def run(x):
+        return torch.ops.vllm.marlin_gemm(
+            x,
+            qweight,
+            None,
+            scales,
+            None,
+            None,
+            zeros,
+            None,
+            quant_type.id,
+            size_m,
+            size_n,
+            size_k,
+            use_atomic_add=False,
+            use_fp32_reduce=True,
+        )
+
+    for stream, graph, x in zip(streams, graphs, inputs):
+        stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(stream):
+            locks.append(get_marlin_workspace(x.device))
+            run(x)
+            with torch.cuda.graph(graph, stream=stream):
+                outputs.append(run(x))
+        torch.cuda.current_stream().wait_stream(stream)
+
+    assert locks[0].data_ptr() != locks[1].data_ptr()
+    current_workspace_manager().lock()
+    for _ in range(3):
+        for stream, graph, x, lock in zip(streams, graphs, inputs, locks):
+            with torch.cuda.stream(stream):
+                assert get_marlin_workspace(x.device) is lock
+                x.add_(0.1)
+                graph.replay()
+        for stream in streams:
+            torch.cuda.current_stream().wait_stream(stream)
+        for x, output, lock in zip(inputs, outputs, locks):
+            assert compute_max_diff(output, x @ w_ref) < 0.04
+            assert torch.count_nonzero(lock) == 0
+
+
 def test_marlin_gemm_subset_input():
     quant_type = scalar_types.uint4b8
     group_size = 128

tests/kernels/quantization/test_marlin_tile_padding.py

@@ -14,6 +14,7 @@ from vllm import _custom_ops as ops
 from vllm.model_executor.layers.quantization.utils.marlin_utils import (
     GPTQ_MARLIN_TILE,
     apply_gptq_marlin_linear,
+    get_marlin_workspace,
     marlin_make_empty,
     marlin_make_workspace_new,
     marlin_moe_padded_intermediate,
@@ -225,7 +226,7 @@ def test_fp8_marlin_padded_round_trip(shape, use_bias):
         input=x,
         weight=layer.weight,
         weight_scale=layer.weight_scale,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
         bias=layer.bias if use_bias else None,
@@ -288,7 +289,7 @@ def test_nvfp4_marlin_padded_round_trip(shape):
         weight=layer.weight,
         weight_scale=layer.weight_scale,
         weight_global_scale=layer.weight_global_scale,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
     )
@@ -392,7 +393,7 @@ def test_fp8_block_marlin_padded_round_trip(shape):
         input=x,
         weight=layer.weight,
         weight_scale=layer.weight_scale_inv,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
         bias=None,
@@ -440,7 +441,7 @@ def test_mxfp8_marlin_padded_round_trip(shape):
         input=x,
         weight=layer.weight,
         weight_scale=layer.weight_scale,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(layer.weight.device),
         size_n=size_n,
         size_k=size_k,
     )
@@ -756,7 +757,7 @@ def test_fp8_marlin_moe_padded_round_trip(shape, quant):
         topk_ids,
         quant_type_id=scalar_types.float8_e4m3fn.id,
         global_num_experts=e,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(device),
     )
     with set_current_vllm_config(VllmConfig()):
         ref = torch_experts(
@@ -833,7 +834,7 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape):
         topk_ids,
         quant_type_id=scalar_types.float8_e4m3fn.id,
         global_num_experts=e,
-        workspace=layer.workspace,
+        workspace=get_marlin_workspace(device),
     )
     with set_current_vllm_config(VllmConfig()):
         ref = torch_experts(