@tduncan / Nvidia-Optimization-Stack-VLLM / commits / e5fce7b56b

[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>

Misha Goin committed Sep 20, 2026 at 05:28 UTC e5fce7b56b07cdf992febd0862bfd599d6dca3fd

21 files changed +598 -211

tests/kernels/moe/test_moe.py

+17 -3

@@ -1384,24 +1384,33 @@ def test_humming_global_valid_shape_m(

1384 assert result == expected

1385

1386

1387 -def test_humming_permute_scratch_is_keyed_by_runtime_topk(

1387 +def test_humming_permute_scratch_is_shared_by_config(

1388 monkeypatch: pytest.MonkeyPatch,

1389 ):

1390 from types import SimpleNamespace

1391 from unittest.mock import Mock

1392

1393 import vllm.model_executor.layers.fused_moe.experts.fused_humming_moe as humming

1394 + import vllm.model_executor.layers.fused_moe.moe_permute_unpermute as permute

1395 + import vllm.v1.worker.workspace as workspace

1396 +

1397 + manager = workspace.WorkspaceManager(torch.device("cpu"))

1398 + monkeypatch.setattr(workspace, "_manager", manager)

1399

1400 scratch_topk6 = Mock()

1401 scratch_topk1 = Mock()

1402 scratch_type = Mock(side_effect=[scratch_topk6, scratch_topk1])

1403 monkeypatch.setattr(humming, "moe_permute_unpermute_supported", lambda: True)

1399 - monkeypatch.setattr(humming, "MoEPermuteScratch", scratch_type)

1404 + monkeypatch.setattr(permute, "MoEPermuteScratch", scratch_type)

1405 moe_config = make_dummy_moe_config(max_num_tokens=512, experts_per_token=6)

1406 moe_config.moe_parallel_config.dp_size = 2

1402 - experts = SimpleNamespace(_permute_scratch={}, moe_config=moe_config)

1407 + experts = SimpleNamespace(moe_config=moe_config)

1408 + other_layer = SimpleNamespace(moe_config=moe_config)

1409

1410 assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6

1411 + assert (

1412 + humming.HummingExpertsBase._get_permute_scratch(other_layer, 6) is scratch_topk6

1413 + )

1414 assert humming.HummingExpertsBase._get_permute_scratch(experts, 1) is scratch_topk1

1415 assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6

1416

@@ -1412,6 +1421,11 @@ def test_humming_permute_scratch_is_keyed_by_runtime_topk(

1421 assert second_call.kwargs["max_num_tokens"] == 6144

1422 assert second_call.kwargs["topk"] == 1

1423

1424 + manager.lock()

1425 + assert humming.HummingExpertsBase._get_permute_scratch(experts, 6) is scratch_topk6

1426 + with pytest.raises(AssertionError, match="was not allocated during warmup"):

1427 + humming.HummingExpertsBase._get_permute_scratch(experts, 2)

1428 +

1429

1430 def test_humming_delegates_to_instance_activation():

1431 from types import SimpleNamespace

tests/kernels/moe/test_moe_permute_unpermute.py

+163 -12

@@ -15,12 +15,14 @@ from vllm.model_executor.layers.fused_moe.expert_map_manager import (

15 )

16 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (

17 MoEPermuteScratch,

18 + get_moe_permute_scratch,

19 moe_permute,

20 moe_permute_unpermute_supported,

21 moe_unpermute,

22 )

23 from vllm.platforms import current_platform

24 from vllm.utils.torch_utils import set_random_seed

25 +from vllm.v1.worker.workspace import current_workspace_manager

27 NUM_EXPERTS = [16, 64, 256]

28 TOP_KS = [2, 6, 8]

@@ -120,6 +122,7 @@ def torch_unpermute(

122 @pytest.mark.parametrize("topk", TOP_KS)

123 @pytest.mark.parametrize("dtype", [torch.bfloat16])

124 @pytest.mark.parametrize("ep_size", EP_SIZE)

125 +@pytest.mark.parametrize("use_scratch", [False, True])

126 def test_moe_permute_unpermute(

127 n_token: int,

128 n_hidden: int,

@@ -127,6 +130,8 @@ def test_moe_permute_unpermute(

130 n_expert: int,

131 ep_size: int,

132 dtype: torch.dtype,

133 + use_scratch: bool,

134 + workspace_init,

135 ):

136 if not moe_permute_unpermute_supported():

137 pytest.skip("moe_permute_unpermute is not supported on this platform.")

@@ -160,6 +165,18 @@ def test_moe_permute_unpermute(

165 expert_map=expert_map,

166 )

167

168 + scratch = None

169 + if use_scratch:

170 + scratch = get_moe_permute_scratch(

171 + max_num_tokens=n_token,

172 + topk=topk,

173 + num_experts=n_expert,

174 + num_local_experts=n_local_expert,

175 + device=hidden_states.device,

176 + hidden_size=n_hidden,

177 + hidden_dtype=dtype,

178 + )

179 +

180 (

181 permuted_hidden_states,

182 _,

@@ -173,6 +190,7 @@ def test_moe_permute_unpermute(

190 n_expert=n_expert,

191 n_local_expert=n_local_expert,

192 expert_map=expert_map,

193 + scratch=scratch,

194 )

195

196 # check expert_first_token_offset

@@ -212,20 +230,22 @@ def test_moe_permute_unpermute(

230 torch.testing.assert_close(result4, gold4, atol=2e-2, rtol=0)

231

232

215 -@pytest.mark.parametrize("dtype", [torch.bfloat16])

216 -def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):

233 +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])

234 +@pytest.mark.parametrize("n_token", [1, 33, 128])

235 +@pytest.mark.parametrize("topk", [1, 6, 8])

236 +def test_moe_permute_reuses_scratch_buffers(

237 + dtype: torch.dtype, n_token: int, topk: int, workspace_init

238 +):

239 if not moe_permute_unpermute_supported():

240 pytest.skip("moe_permute_unpermute is not supported on this platform.")

241

220 - n_token = 64

242 n_hidden = 2048

243 n_expert = 16

223 - topk = 4

244 hidden_states = torch.randn((n_token, n_hidden), device="cuda").to(dtype)

245 gating_output = torch.randn((n_token, n_expert), device="cuda").to(dtype)

246 _, topk_ids, _ = fused_topk(hidden_states, gating_output, topk, False)

247

228 - scratch = MoEPermuteScratch(

248 + scratch_config = dict(

249 max_num_tokens=n_token,

250 topk=topk,

251 num_experts=n_expert,

@@ -234,6 +254,7 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):

254 hidden_size=n_hidden,

255 hidden_dtype=hidden_states.dtype,

256 )

257 + scratch = get_moe_permute_scratch(**scratch_config)

258

259 first = moe_permute(

260 hidden_states=hidden_states,

@@ -242,13 +263,26 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):

263 n_expert=n_expert,

264 scratch=scratch,

265 )

245 - second = moe_permute(

246 - hidden_states=hidden_states,

247 - a1q_scale=None,

248 - topk_ids=topk_ids,

249 - n_expert=n_expert,

250 - scratch=scratch,

251 - )

266 + current_workspace_manager().lock()

267 + assert get_moe_permute_scratch(**scratch_config) is scratch

268 + graph = torch.cuda.CUDAGraph()

269 + with torch.cuda.graph(graph):

270 + second = moe_permute(

271 + hidden_states=hidden_states,

272 + a1q_scale=None,

273 + topk_ids=topk_ids,

274 + n_expert=n_expert,

275 + scratch=get_moe_permute_scratch(**scratch_config),

276 + )

277 +

278 + for _ in range(2):

279 + hidden_states.add_(1)

280 + topk_ids.copy_(topk_ids.roll(1, 0))

281 + graph.replay()

282 + expected = moe_permute(hidden_states, None, topk_ids, n_expert)

283 + for actual, reference in zip(second, expected):

284 + if actual is not None:

285 + torch.testing.assert_close(actual, reference)

286

287 (

288 permuted_hidden_states_1,

@@ -288,6 +322,123 @@ def test_moe_permute_reuses_scratch_buffers(dtype: torch.dtype):

322 )

323

324

325 +def test_moe_permute_scratch_reused_across_graph_sizes(workspace_init) -> None:

326 + """Switching graph sizes must not expose metadata left by a larger batch."""

327 + if not moe_permute_unpermute_supported():

328 + pytest.skip("moe_permute_unpermute is not supported on this platform.")

329 +

330 + device = torch.device("cuda:0")

331 + scratch = get_moe_permute_scratch(

332 + max_num_tokens=128,

333 + topk=6,

334 + num_experts=16,

335 + num_local_experts=16,

336 + device=device,

337 + hidden_size=2048,

338 + hidden_dtype=torch.bfloat16,

339 + )

340 + current_workspace_manager().lock()

341 + runs = []

342 + for n_token in (128, 1, 33):

343 + hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device)

344 + topk_ids = torch.randint(16, (n_token, 6), device=device)

345 + moe_permute(hidden, None, topk_ids, 16, scratch=scratch)

346 + graph = torch.cuda.CUDAGraph()

347 + with torch.cuda.graph(graph):

348 + output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch)

349 + runs.append((graph, hidden, topk_ids, output))

350 +

351 + for graph, hidden, topk_ids, output in runs + runs[::-1]:

352 + hidden.add_(1)

353 + topk_ids.copy_((topk_ids + 1) % 16)

354 + graph.replay()

355 + expected = moe_permute(hidden, None, topk_ids, 16)

356 + for actual, reference in zip(output, expected):

357 + if actual is not None:

358 + torch.testing.assert_close(actual, reference)

359 +

360 +

361 +def test_moe_permute_scratch_isolated_across_execution_slots(monkeypatch) -> None:

362 + """Concurrent graph replays in different ubatches/lanes cannot share scratch."""

363 + if not moe_permute_unpermute_supported():

364 + pytest.skip("moe_permute_unpermute is not supported on this platform.")

365 +

366 + from vllm.v1.worker import workspace

367 +

368 + device = torch.device("cuda:0")

369 + manager = workspace.WorkspaceManager(device, num_ubatches=2, num_lanes=2)

370 + monkeypatch.setattr(workspace, "_manager", manager)

371 + ubatch = 0

372 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: ubatch)

373 + scratch_config = dict(

374 + max_num_tokens=64,

375 + topk=6,

376 + num_experts=16,

377 + num_local_experts=16,

378 + device=device,

379 + hidden_size=2048,

380 + hidden_dtype=torch.bfloat16,

381 + )

382 + runs = []

383 + for i, n_token in enumerate((1, 7, 33, 64)):

384 + ubatch = i // 2

385 + stream = torch.cuda.Stream()

386 + stream.wait_stream(torch.cuda.current_stream())

387 + with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream):

388 + hidden = torch.randn((n_token, 2048), dtype=torch.bfloat16, device=device)

389 + topk_ids = torch.randint(16, (n_token, 6), device=device)

390 + scratch = get_moe_permute_scratch(**scratch_config)

391 + moe_permute(hidden, None, topk_ids, 16, scratch=scratch)

392 + graph = torch.cuda.CUDAGraph()

393 + with torch.cuda.graph(graph, stream=stream):

394 + output = moe_permute(hidden, None, topk_ids, 16, scratch=scratch)

395 + runs.append((stream, graph, hidden, topk_ids, output, scratch))

396 + torch.cuda.current_stream().wait_stream(stream)

397 +

398 + assert len({run[-1].permuted_hidden_states.data_ptr() for run in runs}) == 4

399 + manager.lock()

400 + for _ in range(3):

401 + for i, (stream, graph, hidden, topk_ids, _, scratch) in enumerate(runs):

402 + ubatch = i // 2

403 + with workspace.use_workspace_lane(i % 2), torch.cuda.stream(stream):

404 + assert get_moe_permute_scratch(**scratch_config) is scratch

405 + hidden.add_(1)

406 + topk_ids.copy_((topk_ids + 1) % 16)

407 + graph.replay()

408 + for stream, _, hidden, topk_ids, output, _ in runs:

409 + torch.cuda.current_stream().wait_stream(stream)

410 + expected = moe_permute(hidden, None, topk_ids, 16)

411 + for actual, reference in zip(output, expected):

412 + if actual is not None:

413 + torch.testing.assert_close(actual, reference)

414 +

415 +

416 +def test_moe_permute_scratch_without_manager(monkeypatch) -> None:

417 + """Standalone calls get independent scratch with initialized row indices."""

418 + if not moe_permute_unpermute_supported():

419 + pytest.skip("moe_permute_unpermute is not supported on this platform.")

420 +

421 + from vllm.v1.worker import workspace

422 +

423 + monkeypatch.setattr(workspace, "_manager", None)

424 + config = dict(

425 + max_num_tokens=4,

426 + topk=2,

427 + num_experts=4,

428 + num_local_experts=4,

429 + device=torch.device("cuda"),

430 + )

431 + first = get_moe_permute_scratch(**config)

432 + second = get_moe_permute_scratch(**config)

433 +

434 + assert (

435 + first.token_expert_indices.data_ptr() != second.token_expert_indices.data_ptr()

436 + )

437 + torch.testing.assert_close(

438 + first.token_expert_indices, torch.arange(8, dtype=torch.int32, device="cuda")

439 + )

440 +

441 +

442 def test_moe_permute_ignores_invalid_expert_ids_with_scratch() -> None:

443 if not moe_permute_unpermute_supported():

444 pytest.skip("moe_permute_unpermute is not supported on this platform.")

tests/kernels/quantization/test_marlin_gemm.py

+58

@@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization.utils.int8_utils import (

17 per_token_quant_int8,

18 )

19 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

20 + get_marlin_workspace,

21 marlin_make_empty,

22 marlin_make_workspace_new,

23 marlin_permute_bias,

@@ -43,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

44 )

45 from vllm.platforms import current_platform

46 from vllm.scalar_type import scalar_types

47 +from vllm.v1.worker.workspace import current_workspace_manager

49 if current_platform.is_rocm():

50 pytest.skip(

@@ -477,6 +479,62 @@ def test_marlin_gemm(

479 assert max_diff < 0.04

480

481

482 +def test_marlin_persistent_locks_cuda_graph_streams(workspace_init):

483 + """Compiled GEMMs select stream-local locks at runtime and survive replay."""

484 + size_m, size_k, size_n = 32, 1024, 2048

485 + inputs = [rand_data((size_m, size_k)) for _ in range(2)]

486 + weight = rand_data((size_k, size_n))

487 + quant_type = scalar_types.uint4b8

488 + w_ref, qweight, scales = marlin_quantize(weight, quant_type, 128)

489 + zeros = marlin_make_empty(scales.device)

490 + streams = [torch.cuda.Stream() for _ in inputs]

491 + graphs = [torch.cuda.CUDAGraph() for _ in inputs]

492 + locks = []

493 + outputs = []

494 +

495 + @torch.compile(backend="eager", fullgraph=True)

496 + def run(x):

497 + return torch.ops.vllm.marlin_gemm(

498 + x,

499 + qweight,

500 + None,

501 + scales,

502 + None,

503 + None,

504 + zeros,

505 + None,

506 + quant_type.id,

507 + size_m,

508 + size_n,

509 + size_k,

510 + use_atomic_add=False,

511 + use_fp32_reduce=True,

512 + )

513 +

514 + for stream, graph, x in zip(streams, graphs, inputs):

515 + stream.wait_stream(torch.cuda.current_stream())

516 + with torch.cuda.stream(stream):

517 + locks.append(get_marlin_workspace(x.device))

518 + run(x)

519 + with torch.cuda.graph(graph, stream=stream):

520 + outputs.append(run(x))

521 + torch.cuda.current_stream().wait_stream(stream)

522 +

523 + assert locks[0].data_ptr() != locks[1].data_ptr()

524 + current_workspace_manager().lock()

525 + for _ in range(3):

526 + for stream, graph, x, lock in zip(streams, graphs, inputs, locks):

527 + with torch.cuda.stream(stream):

528 + assert get_marlin_workspace(x.device) is lock

529 + x.add_(0.1)

530 + graph.replay()

531 + for stream in streams:

532 + torch.cuda.current_stream().wait_stream(stream)

533 + for x, output, lock in zip(inputs, outputs, locks):

534 + assert compute_max_diff(output, x @ w_ref) < 0.04

535 + assert torch.count_nonzero(lock) == 0

536 +

537 +

538 def test_marlin_gemm_subset_input():

539 quant_type = scalar_types.uint4b8

540 group_size = 128

tests/kernels/quantization/test_marlin_tile_padding.py

+7 -6

@@ -14,6 +14,7 @@ from vllm import _custom_ops as ops

14 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

15 GPTQ_MARLIN_TILE,

16 apply_gptq_marlin_linear,

17 + get_marlin_workspace,

18 marlin_make_empty,

19 marlin_make_workspace_new,

20 marlin_moe_padded_intermediate,

@@ -225,7 +226,7 @@ def test_fp8_marlin_padded_round_trip(shape, use_bias):

226 input=x,

227 weight=layer.weight,

228 weight_scale=layer.weight_scale,

228 - workspace=layer.workspace,

229 + workspace=get_marlin_workspace(layer.weight.device),

230 size_n=size_n,

231 size_k=size_k,

232 bias=layer.bias if use_bias else None,

@@ -288,7 +289,7 @@ def test_nvfp4_marlin_padded_round_trip(shape):

289 weight=layer.weight,

290 weight_scale=layer.weight_scale,

291 weight_global_scale=layer.weight_global_scale,

291 - workspace=layer.workspace,

292 + workspace=get_marlin_workspace(layer.weight.device),

293 size_n=size_n,

294 size_k=size_k,

295 )

@@ -392,7 +393,7 @@ def test_fp8_block_marlin_padded_round_trip(shape):

393 input=x,

394 weight=layer.weight,

395 weight_scale=layer.weight_scale_inv,

395 - workspace=layer.workspace,

396 + workspace=get_marlin_workspace(layer.weight.device),

397 size_n=size_n,

398 size_k=size_k,

399 bias=None,

@@ -440,7 +441,7 @@ def test_mxfp8_marlin_padded_round_trip(shape):

441 input=x,

442 weight=layer.weight,

443 weight_scale=layer.weight_scale,

443 - workspace=layer.workspace,

444 + workspace=get_marlin_workspace(layer.weight.device),

445 size_n=size_n,

446 size_k=size_k,

447 )

@@ -756,7 +757,7 @@ def test_fp8_marlin_moe_padded_round_trip(shape, quant):

757 topk_ids,

758 quant_type_id=scalar_types.float8_e4m3fn.id,

759 global_num_experts=e,

759 - workspace=layer.workspace,

760 + workspace=get_marlin_workspace(device),

761 )

762 with set_current_vllm_config(VllmConfig()):

763 ref = torch_experts(

@@ -833,7 +834,7 @@ def test_mxfp8_marlin_moe_padded_round_trip(shape):

834 topk_ids,

835 quant_type_id=scalar_types.float8_e4m3fn.id,

836 global_num_experts=e,

836 - workspace=layer.workspace,

837 + workspace=get_marlin_workspace(device),

838 )

839 with set_current_vllm_config(VllmConfig()):

840 ref = torch_experts(

tests/model_executor/model_loader/test_reload.py

+44 -23

@@ -330,8 +330,8 @@ def _load_marlin_checkpoint_format_weights(layer):

330 )

331

332

333 -def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_init):

334 - """Marlin must reuse the workspace storage captured by CUDA graphs."""

333 +def test_marlin_post_load_does_not_own_workspace(monkeypatch, dist_init):

334 + """Weight reload must not create layer-owned Marlin lock storage."""

335 _stub_marlin_ops(monkeypatch)

336 kernel = _make_marlin_kernel()

337

@@ -339,19 +339,17 @@ def test_marlin_post_load_preserves_runtime_tensor_addresses(monkeypatch, dist_i

339 _load_marlin_checkpoint_format_weights(layer)

340 kernel.process_weights_after_loading(layer)

341

342 - workspace_ptr = kernel.workspace.data_ptr()

342 + assert not hasattr(kernel, "workspace")

343

344 _load_marlin_checkpoint_format_weights(layer)

345 kernel.process_weights_after_loading(layer)

346

347 - assert kernel.workspace.data_ptr() == workspace_ptr

348 - assert torch.all(kernel.workspace == 0)

347 + assert not hasattr(kernel, "workspace")

348

349

350 @pytest.mark.parametrize("variant", ["fp8", "mxfp8", "nvfp4"])

352 -def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant):

353 - """The Marlin fallback prepare_* functions rerun on weight reload and must

354 - reuse the workspace storage whose address captured CUDA graphs hold."""

351 +def test_marlin_prepare_layer_does_not_own_workspace(monkeypatch, variant):

352 + """Weight preparation must not attach runtime workspace to model layers."""

353 from vllm import _custom_ops as ops

354 from vllm.model_executor.layers.quantization.utils import (

355 marlin_utils,

@@ -417,34 +415,57 @@ def test_marlin_prepare_layer_preserves_workspace_address(monkeypatch, variant):

415

416 load_checkpoint_format_weights()

417 prepare(layer)

420 - workspace_ptr = layer.workspace.data_ptr()

418 + assert not hasattr(layer, "workspace")

419

420 # Reload: fresh checkpoint-format tensors, prepare runs again

421 load_checkpoint_format_weights()

422 prepare(layer)

423

426 - assert layer.workspace.data_ptr() == workspace_ptr

427 - assert torch.all(layer.workspace == 0)

424 + assert not hasattr(layer, "workspace")

425

426

430 -def test_marlin_make_workspace_new_rejects_incompatible_existing(monkeypatch):

431 - """An incompatible existing workspace means the address captured by CUDA

432 - graphs is already unusable; allocating a replacement would hide that."""

427 +def test_marlin_workspace_uses_persistent_workspace_manager(monkeypatch):

428 + """Calls reuse initialized locks; independent streams get separate storage."""

429 from vllm.model_executor.layers.quantization.utils import marlin_utils

430 + from vllm.utils import torch_utils

431 + from vllm.v1.worker import workspace as workspace_module

432

433 monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4)

434 device = torch.device("cpu")

435 + manager = workspace_module.WorkspaceManager(device)

436 + monkeypatch.setattr(workspace_module, "_manager", manager)

437 + stream = "main"

438 + monkeypatch.setattr(torch_utils, "current_stream", lambda: stream)

439

438 - workspace = marlin_utils.marlin_make_workspace_new(device)

439 - reused = marlin_utils.marlin_make_workspace_new(device, existing=workspace)

440 - assert reused is workspace

440 + workspace = marlin_utils.get_marlin_workspace(device)

441 + assert workspace.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,)

442 + assert workspace.dtype == torch.int32

443 + assert torch.count_nonzero(workspace) == 0

444

442 - with pytest.raises(ValueError, match="incompatible"):

443 - marlin_utils.marlin_make_workspace_new(device, 4, existing=workspace)

444 - with pytest.raises(ValueError, match="incompatible"):

445 - marlin_utils.marlin_make_workspace_new(

446 - device, existing=workspace.to(torch.int64)

447 - )

445 + workspace.fill_(1)

446 + assert marlin_utils.get_marlin_workspace(device) is workspace

447 + assert torch.all(workspace == 1)

448 +

449 + stream = "aux"

450 + aux_workspace = marlin_utils.get_marlin_workspace(device)

451 + assert aux_workspace.data_ptr() != workspace.data_ptr()

452 + assert torch.count_nonzero(aux_workspace) == 0

453 +

454 + manager.lock()

455 + assert marlin_utils.get_marlin_workspace(device) is aux_workspace

456 +

457 +

458 +def test_marlin_workspace_without_manager(monkeypatch):

459 + from vllm.model_executor.layers.quantization.utils import marlin_utils

460 + from vllm.v1.worker import workspace as workspace_module

461 +

462 + monkeypatch.setattr(workspace_module, "_manager", None)

463 + monkeypatch.setattr(marlin_utils, "num_compute_units", lambda _: 4)

464 + first = marlin_utils.get_marlin_workspace(torch.device("cpu"))

465 + second = marlin_utils.get_marlin_workspace(torch.device("cpu"))

466 + assert first.data_ptr() != second.data_ptr()

467 + assert first.shape == (4 * marlin_utils.MARLIN_MAX_BLOCKS_PER_SM,)

468 + assert torch.count_nonzero(first) == 0

469

470

471 def test_model_cleanup(dist_init, default_vllm_config):

tests/v1/worker/test_workspace.py

+63

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from typing import cast

5 +from unittest.mock import Mock

7 import pytest

8 import torch

@@ -124,3 +125,65 @@ def test_workspace_lane_validation(monkeypatch) -> None:

125

126 with pytest.raises(ValueError, match="at least one"):

127 workspace.WorkspaceManager(torch.device("cpu"), num_lanes=0)

128 +

129 +

130 +def test_persistent_resources_are_lazy_and_isolated(monkeypatch) -> None:

131 + """Only the requesting slot allocates; each ubatch/lane owns its resource."""

132 + active_ubatch = 0

133 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: active_ubatch)

134 + manager = workspace.WorkspaceManager(

135 + torch.device("cpu"), num_ubatches=2, num_lanes=2

136 + )

137 + factory = Mock(side_effect=lambda: {"buffer": torch.zeros(8)})

138 + resources = []

139 + for ubatch in range(2):

140 + active_ubatch = ubatch

141 + for lane in range(2):

142 + with workspace.use_workspace_lane(lane):

143 + resource = manager.get_persistent_resource("scratch", factory)

144 + resources.append(resource)

145 + resource["buffer"].fill_(len(resources))

146 + assert manager.get_persistent_resource("scratch", factory) is resource

147 + assert factory.call_count == len(resources)

148 +

149 + assert len({r["buffer"].data_ptr() for r in resources}) == 4

150 + for i, resource in enumerate(resources, 1):

151 + assert torch.all(resource["buffer"] == i)

152 +

153 +

154 +def test_persistent_resource_lock_applies_to_each_slot(monkeypatch) -> None:

155 + monkeypatch.setattr(workspace, "dbo_current_ubatch_id", lambda: 0)

156 + manager = workspace.WorkspaceManager(torch.device("cpu"), num_lanes=2)

157 + factory = Mock(side_effect=object)

158 + first = manager.get_persistent_resource("scratch", factory)

159 + manager.lock()

160 + assert manager.get_persistent_resource("scratch", factory) is first

161 + with pytest.raises(AssertionError, match="was not allocated during warmup"):

162 + manager.get_persistent_resource("new", factory)

163 + with workspace.use_workspace_lane(1):

164 + with pytest.raises(AssertionError, match="was not allocated during warmup"):

165 + manager.get_persistent_resource("scratch", factory)

166 + assert factory.call_count == 1

167 + manager.unlock()

168 + assert manager.get_persistent_resource("scratch", factory) is not first

169 + assert manager.get_persistent_resource("scratch", factory) is first

170 +

171 +

172 +def test_persistent_tensor_preserves_contents_and_rejects_changes() -> None:

173 + """Initialization happens once and transient scratch cannot overwrite locks."""

174 + manager = workspace.WorkspaceManager(torch.device("cpu"))

175 + first = manager.get_persistent("locks", (8,), torch.int32, zero_init=True)

176 + assert torch.count_nonzero(first) == 0

177 + first.fill_(1)

178 + (scratch,) = manager.get_simultaneous(((1024,), torch.uint8))

179 + scratch.zero_()

180 + assert manager.get_persistent("locks", (8,), torch.int32, zero_init=True) is first

181 + assert torch.all(first == 1)

182 + with pytest.raises(ValueError, match="requested shape"):

183 + manager.get_persistent("locks", (16,), torch.int32)

184 + with pytest.raises(ValueError, match="requested shape"):

185 + manager.get_persistent("locks", (8,), torch.int64)

186 + manager.lock()

187 + assert manager.get_persistent("locks", (8,), torch.int32) is first

188 + with pytest.raises(AssertionError, match="was not allocated during warmup"):

189 + manager.get_persistent("new", (8,), torch.int32)

vllm/model_executor/kernels/linear/mixed_precision/marlin.py

+1 -7

@@ -10,7 +10,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (

10 apply_gptq_marlin_linear,

11 marlin_act_int8_process_scales,

12 marlin_make_empty,

13 - marlin_make_workspace_new,

13 marlin_pad_dim,

14 marlin_pad_qweight,

15 marlin_pad_scales,

@@ -91,11 +90,6 @@ class MarlinLinearKernel(MPLinearKernel):

90 size_k, size_n = c.partition_weight_shape

91 padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size)

94 - # Allocate marlin workspace, reusing existing storage on reload.

95 - self.workspace = marlin_make_workspace_new(

96 - device, existing=getattr(self, "workspace", None)

97 - )

98 -

93 # Default name since marlin requires empty parameter for zp,

94 # TODO: remove this requirement from marlin (allow optional tensors)

95 if self.w_zp_name is None:

@@ -197,7 +191,7 @@ class MarlinLinearKernel(MPLinearKernel):

191 weight=w_q,

192 weight_scale=w_s,

193 weight_zp=w_zp, # type: ignore

200 - workspace=self.workspace,

194 + workspace=None,

195 wtype=c.weight_type,

196 input_size_per_partition=c.partition_weight_shape[0],

197 output_size_per_partition=c.partition_weight_shape[1],

vllm/model_executor/kernels/linear/mxfp4/marlin.py

+1 -1

@@ -58,7 +58,7 @@ class MarlinMxFp4LinearKernel(MxFp4LinearKernel):

58 weight=layer.weight,

59 weight_scale=layer.weight_scale,

60 weight_global_scale=None,

61 - workspace=layer.workspace,

61 + workspace=None,

62 size_n=layer.output_size_per_partition,

63 size_k=layer.input_size_per_partition,

64 bias=bias,

vllm/model_executor/kernels/linear/mxfp8/marlin.py

+1 -1

@@ -46,7 +46,7 @@ class MarlinMxfp8LinearKernel(Mxfp8LinearKernel):

46 input=x,

47 weight=layer.weight,

48 weight_scale=layer.weight_scale,

49 - workspace=layer.workspace,

49 + workspace=None,

50 size_n=layer.output_size_per_partition,

51 size_k=layer.input_size_per_partition,

52 bias=bias,

vllm/model_executor/kernels/linear/nvfp4/marlin.py

+1 -1

@@ -50,7 +50,7 @@ class MarlinNvFp4LinearKernel(NvFp4LinearKernel):

50 weight=layer.weight,

51 weight_scale=layer.weight_scale,

52 weight_global_scale=layer.weight_global_scale,

53 - workspace=layer.workspace,

53 + workspace=None,

54 size_n=layer.output_size_per_partition,

55 size_k=layer.input_size_per_partition,

56 bias=bias,

vllm/model_executor/kernels/linear/scaled_mm/marlin.py

+1 -1

@@ -93,7 +93,7 @@ class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):

93 input=x,

94 weight=layer.weight,

95 weight_scale=weight_scale,

96 - workspace=layer.workspace,

96 + workspace=None,

97 size_n=layer.output_size_per_partition,

98 size_k=layer.input_size_per_partition,

99 input_dtype=self.marlin_input_dtype,

vllm/model_executor/layers/fused_moe/experts/fused_humming_moe.py

+15 -19

@@ -27,6 +27,7 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import (

27 from vllm.model_executor.layers.fused_moe.moe_fused_mul_sum import moe_fused_mul_sum

28 from vllm.model_executor.layers.fused_moe.moe_permute_unpermute import (

29 MoEPermuteScratch,

30 + get_moe_permute_scratch,

31 moe_permute,

32 moe_permute_unpermute_supported,

33 moe_unpermute,

@@ -176,7 +177,6 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

177 max_num_tokens=max_num_tokens,

178 num_dispatchers=num_dispatchers,

179 )

179 - self._permute_scratch: dict[int, MoEPermuteScratch] = {}

180

181 def init_humming_moe(self):

182 from vllm.utils.humming import get_heuristics_config

@@ -259,24 +259,20 @@ class HummingExpertsBase(mk.FusedMoEExpertsModular):

259 if not moe_permute_unpermute_supported():

260 return None

261

262 - scratch = self._permute_scratch.get(topk)

263 - if scratch is None:

264 - max_expanded_rows = (

265 - self.moe_config.max_num_tokens

266 - * self.moe_config.dp_size

267 - * self.moe_config.experts_per_token

268 - )

269 - scratch = MoEPermuteScratch(

270 - max_num_tokens=math.ceil(max_expanded_rows / topk),

271 - topk=topk,

272 - num_experts=self.moe_config.num_experts,

273 - num_local_experts=self.moe_config.num_local_experts,

274 - device=torch.device(self.moe_config.device),

275 - hidden_size=self.moe_config.hidden_dim,

276 - hidden_dtype=self.moe_config.in_dtype,

277 - )

278 - self._permute_scratch[topk] = scratch

279 - return scratch

262 + max_expanded_rows = (

263 + self.moe_config.max_num_tokens

264 + * self.moe_config.dp_size

265 + * self.moe_config.experts_per_token

266 + )

267 + return get_moe_permute_scratch(

268 + max_num_tokens=math.ceil(max_expanded_rows / topk),

269 + topk=topk,

270 + num_experts=self.moe_config.num_experts,

271 + num_local_experts=self.moe_config.num_local_experts,

272 + device=torch.device(self.moe_config.device),

273 + hidden_size=self.moe_config.hidden_dim,

274 + hidden_dtype=self.moe_config.in_dtype,

275 + )

276

277 def get_global_valid_shape_m(self, topk_ids: torch.Tensor):

278 ctx = get_forward_context()

vllm/model_executor/layers/fused_moe/experts/marlin_moe.py

+2 -2

@@ -35,7 +35,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

35 from vllm.model_executor.layers.fused_moe.utils import _resize_cache

36 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

37 get_marlin_input_dtype,

38 - marlin_make_workspace_new,

38 + get_marlin_workspace,

39 marlin_moe_intermediate_size,

40 marlin_quant_input,

41 )

@@ -95,7 +95,7 @@ def _fused_marlin_moe(

95 N = marlin_moe_intermediate_size(w1, w2)

96 w13_num_shards = 2 if activation.is_gated else 1

97 if workspace is None:

98 - workspace = marlin_make_workspace_new(hidden_states.device, 4)

98 + workspace = get_marlin_workspace(hidden_states.device)

100 if intermediate_cache13 is None:

101 intermediate_cache13 = torch.empty(

vllm/model_executor/layers/fused_moe/moe_permute_unpermute.py

+47

@@ -5,6 +5,11 @@ from dataclasses import dataclass, field

6 import torch

8 +from vllm.v1.worker.workspace import (

9 + current_workspace_manager,

10 + is_workspace_manager_initialized,

11 +)

12 +

14 @dataclass

15 class MoEPermuteScratch:

@@ -102,6 +107,48 @@ class MoEPermuteScratch:

107 return topk_ids_int32

108

109

110 +def get_moe_permute_scratch(

111 + *,

112 + max_num_tokens: int,

113 + topk: int,

114 + num_experts: int,

115 + num_local_experts: int,

116 + device: torch.device,

117 + hidden_size: int | None = None,

118 + hidden_dtype: torch.dtype | None = None,

119 +) -> MoEPermuteScratch:

120 + """Share scratch across sequential layers in the current ubatch and lane.

121 +

122 + Without a workspace manager, allocate new scratch for each call.

123 + """

124 +

125 + def create_scratch() -> MoEPermuteScratch:

126 + return MoEPermuteScratch(

127 + max_num_tokens=max_num_tokens,

128 + topk=topk,

129 + num_experts=num_experts,

130 + num_local_experts=num_local_experts,

131 + device=device,

132 + hidden_size=hidden_size,

133 + hidden_dtype=hidden_dtype,

134 + )

135 +

136 + if not is_workspace_manager_initialized():

137 + return create_scratch()

138 +

139 + key = (

140 + MoEPermuteScratch,

141 + max_num_tokens,

142 + topk,

143 + num_experts,

144 + num_local_experts,

145 + device,

146 + hidden_size,

147 + hidden_dtype,

148 + )

149 + return current_workspace_manager().get_persistent_resource(key, create_scratch)

150 +

151 +

152 def moe_permute(

153 hidden_states: torch.Tensor,

154 a1q_scale: torch.Tensor | None,

vllm/model_executor/layers/quantization/auto_awq.py

-4

@@ -48,7 +48,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (

48 check_marlin_supports_layer,

49 check_moe_marlin_supports_layer,

50 get_marlin_input_dtype,

51 - marlin_make_workspace_new,

51 verify_marlin_supported,

52 )

53 from vllm.model_executor.layers.quantization.utils.quant_utils import (

@@ -642,9 +641,6 @@ class AutoAWQMoEMethod(FusedMoEMethodBase):

641 layer.register_parameter("w2_qzeros", w2_qzeros)

642 set_weight_attrs(w2_qzeros, extra_weight_attrs)

643

645 - device = layer.w13_qweight.device

646 - layer.workspace = marlin_make_workspace_new(device, 4)

647 -

644 def process_weights_after_loading(self, layer: RoutedExperts) -> None:

645 converted = convert_to_wna16_moe_kernel_format(

646 backend=self.wna16_moe_backend,

vllm/model_executor/layers/quantization/auto_gptq.py

-8

@@ -16,7 +16,6 @@ from vllm.model_executor.kernels.linear import (

16 )

17 from vllm.model_executor.layers.fused_moe import (

18 FusedMoEConfig,

19 - FusedMoEExpertsModular,

19 FusedMoEMethodBase,

20 FusedMoEQuantConfig,

21 FusedMoeWeightScaleSupported,

@@ -46,7 +45,6 @@ from vllm.model_executor.layers.quantization.utils.gptq_utils import (

45 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

46 check_moe_marlin_supports_layer,

47 get_marlin_input_dtype,

49 - marlin_make_workspace_new,

48 marlin_repeat_scales_on_all_ranks,

49 verify_marlin_supported,

50 )

@@ -601,12 +599,6 @@ class AutoGPTQMoEMethod(FusedMoEMethodBase):

599 layer.register_parameter("w2_bias", w2_bias)

600 set_weight_attrs(w2_bias, extra_weight_attrs)

601

604 - if self.experts_cls is not None and issubclass(

605 - self.experts_cls, FusedMoEExpertsModular

606 - ):

607 - device = layer.w13_qweight.device

608 - layer.workspace = marlin_make_workspace_new(device, 4)

609 -

602 def process_weights_after_loading(self, layer: RoutedExperts) -> None:

603 def replace_or_register(name: str, val: torch.Tensor | None):

604 if val is None:

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py

-14

@@ -11,7 +11,6 @@ from compressed_tensors.quantization import (

12 from vllm.logger import init_logger

13 from vllm.model_executor.layers.fused_moe import (

14 - FusedMoEExpertsModular,

14 RoutedExperts,

15 SharedExperts,

16 )

@@ -36,7 +35,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compress

35 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

36 check_moe_marlin_supports_config,

37 get_marlin_input_dtype,

39 - marlin_make_workspace_new,

38 )

39 from vllm.model_executor.layers.quantization.utils.quant_utils import (

40 GroupShape,

@@ -466,18 +464,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

464 torch.nn.Parameter(w2_input_global_scale, requires_grad=False),

465 )

466

469 - # Marlin workspace — only needed for Marlin-family backends, not emulation.

470 - if (

471 - self.experts_cls is not None

472 - and issubclass(self.experts_cls, FusedMoEExpertsModular)

473 - and self.wna16_backend != WNA16MoEBackend.EMULATION

474 - ):

475 - layer.workspace = marlin_make_workspace_new(

476 - layer.w13_weight_packed.device,

477 - 4,

478 - existing=getattr(layer, "workspace", None),

479 - )

480 -

467 # Alias packed weights to w13_weight/w2_weight for the modular kernel interface

468 layer.w13_weight = layer.w13_weight_packed

469 layer.w2_weight = layer.w2_weight_packed

vllm/model_executor/layers/quantization/utils/marlin_utils.py

+91 -24

@@ -8,7 +8,6 @@ import numpy

8 import torch

10 import vllm.envs as envs

11 -from vllm import _custom_ops as ops

11 from vllm.distributed.utils import verify_group_size_divides_partition

12 from vllm.logger import init_logger

13 from vllm.model_executor.layers.fused_moe import RoutedExperts

@@ -23,6 +22,7 @@ from vllm.platforms import current_platform

22 from vllm.scalar_type import ScalarType, scalar_types

23 from vllm.utils.math_utils import round_up

24 from vllm.utils.platform_utils import num_compute_units

25 +from vllm.utils.torch_utils import direct_register_custom_op

27 from .quant_utils import pack_cols, unpack_cols

@@ -32,6 +32,7 @@ GPTQ_MARLIN_TILE = 16

32 GPTQ_MARLIN_MIN_THREAD_N = 64

33 GPTQ_MARLIN_MIN_THREAD_K = 128

34 GPTQ_MARLIN_MAX_PARALLEL = 16

35 +MARLIN_MAX_BLOCKS_PER_SM = 4

37 MARLIN_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128]

@@ -407,29 +408,96 @@ def marlin_moe_intermediate_size(w1_packed: torch.Tensor, w2_packed: torch.Tenso

408 def marlin_make_workspace_new(

409 device: torch.device,

410 max_blocks_per_sm: int = 1,

410 - existing: torch.Tensor | None = None,

411 ) -> torch.Tensor:

412 # In the new marlin kernel, we use the num of threadblocks as workspace

413 # size. The num of threadblocks is sms_count * max_blocks_per_sm.

414 sms = num_compute_units(device.index)

415 - size = sms * max_blocks_per_sm

416 - # On weight reload, reuse the existing storage so the workspace address

417 - # captured by CUDA graphs stays valid.

418 - if existing is not None:

419 - if (

420 - existing.device != device

421 - or existing.dtype != torch.int

422 - or existing.numel() != size

423 - ):

424 - raise ValueError(

425 - f"Existing Marlin workspace is incompatible "

426 - f"(device={existing.device}, dtype={existing.dtype}, "

427 - f"numel={existing.numel()}; expected device={device}, "

428 - f"dtype={torch.int}, numel={size}). Reload must reuse the "

429 - f"workspace storage captured by CUDA graphs."

430 - )

431 - return existing.zero_()

432 - return torch.zeros(size, dtype=torch.int, device=device, requires_grad=False)

415 + return torch.zeros(sms * max_blocks_per_sm, dtype=torch.int, device=device)

416 +

417 +

418 +def get_marlin_workspace(device: torch.device) -> torch.Tensor:

419 + """Get zero-initialized locks for the current stream, ubatch, and lane.

420 +

421 + Marlin resets its locks after each call, so ordered calls can reuse them.

422 + """

423 + from vllm.utils.torch_utils import current_stream

424 + from vllm.v1.worker.workspace import (

425 + current_workspace_manager,

426 + is_workspace_manager_initialized,

427 + )

428 +

429 + if not is_workspace_manager_initialized():

430 + return marlin_make_workspace_new(device, MARLIN_MAX_BLOCKS_PER_SM)

431 +

432 + return current_workspace_manager().get_persistent(

433 + ("marlin_locks", current_stream()),

434 + (num_compute_units(device.index) * MARLIN_MAX_BLOCKS_PER_SM,),

435 + torch.int,

436 + zero_init=True,

437 + )

438 +

439 +

440 +def _marlin_gemm(

441 + a: torch.Tensor,

442 + b_q_weight: torch.Tensor,

443 + b_bias: torch.Tensor | None,

444 + b_scales: torch.Tensor,

445 + a_scales: torch.Tensor | None,

446 + global_scale: torch.Tensor | None,

447 + b_zeros: torch.Tensor | None,

448 + workspace: torch.Tensor | None,

449 + b_q_type_id: int,

450 + size_m: int,

451 + size_n: int,

452 + size_k: int,

453 + use_atomic_add: bool = False,

454 + use_fp32_reduce: bool = False,

455 + is_zp_float: bool = False,

456 +) -> torch.Tensor:

457 + if workspace is None:

458 + workspace = get_marlin_workspace(a.device)

459 + return torch.ops._C.marlin_gemm(

460 + a,

461 + None,

462 + b_q_weight,

463 + b_bias,

464 + b_scales,

465 + a_scales,

466 + global_scale,

467 + b_zeros,

468 + workspace,

469 + b_q_type_id,

470 + size_m,

471 + size_n,

472 + size_k,

473 + use_atomic_add,

474 + use_fp32_reduce,

475 + is_zp_float,

476 + )

477 +

478 +

479 +def _marlin_gemm_fake(

480 + a: torch.Tensor,

481 + b_q_weight: torch.Tensor,

482 + b_bias: torch.Tensor | None,

483 + b_scales: torch.Tensor,

484 + a_scales: torch.Tensor | None,

485 + global_scale: torch.Tensor | None,

486 + b_zeros: torch.Tensor | None,

487 + workspace: torch.Tensor | None,

488 + b_q_type_id: int,

489 + size_m: int,

490 + size_n: int,

491 + size_k: int,

492 + use_atomic_add: bool = False,

493 + use_fp32_reduce: bool = False,

494 + is_zp_float: bool = False,

495 +) -> torch.Tensor:

496 + dtype = a.dtype if a.dtype in (torch.float16, torch.bfloat16) else b_scales.dtype

497 + return torch.empty((size_m, size_n), device=a.device, dtype=dtype)

498 +

499 +

500 +direct_register_custom_op("marlin_gemm", _marlin_gemm, fake_impl=_marlin_gemm_fake)

501

502

503 def marlin_repeat_scales_on_all_ranks(group_size: int, is_row_parallel: bool) -> bool:

@@ -686,7 +754,7 @@ def apply_gptq_marlin_linear(

754 weight: torch.Tensor,

755 weight_scale: torch.Tensor,

756 weight_zp: torch.Tensor,

689 - workspace: torch.Tensor,

757 + workspace: torch.Tensor | None,

758 wtype: ScalarType,

759 output_size_per_partition: int,

760 input_size_per_partition: int,

@@ -723,9 +791,8 @@ def apply_gptq_marlin_linear(

791

792 reshaped_x, a_scales = marlin_quant_input(reshaped_x, input_dtype)

793

726 - output = ops.marlin_gemm(

794 + output = torch.ops.vllm.marlin_gemm(

795 reshaped_x,

728 - None,

796 weight,

797 bias,

798 weight_scale,

@@ -733,7 +800,7 @@ def apply_gptq_marlin_linear(

800 None,

801 weight_zp,

802 workspace,

736 - wtype,

803 + wtype.id,

804 size_m=reshaped_x.shape[0],

805 size_n=padded_n,

806 size_k=padded_k,

vllm/model_executor/layers/quantization/utils/marlin_utils_fp4.py

+3 -23

@@ -10,7 +10,6 @@ from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts

10 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

11 USE_FP32_REDUCE_DEFAULT,

12 get_marlin_input_dtype,

13 - marlin_make_workspace_new,

13 marlin_pad_dim,

14 marlin_pad_qweight,

15 marlin_pad_scales,

@@ -161,7 +160,7 @@ def apply_fp4_marlin_linear(

160 weight: torch.Tensor,

161 weight_scale: torch.Tensor,

162 weight_global_scale: torch.Tensor | None,

164 - workspace: torch.Tensor,

163 + workspace: torch.Tensor | None,

164 size_n: int,

165 size_k: int,

166 bias: torch.Tensor | None = None,

@@ -196,9 +195,8 @@ def apply_fp4_marlin_linear(

195

196 inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn)

197

199 - output = ops.marlin_gemm(

198 + output = torch.ops.vllm.marlin_gemm(

199 a=inputs,

201 - c=None,

200 b_q_weight=weight,

201 b_bias=bias,

202 b_scales=weight_scale,

@@ -206,7 +204,7 @@ def apply_fp4_marlin_linear(

204 global_scale=weight_global_scale,

205 b_zeros=None,

206 workspace=workspace,

209 - b_q_type=scalar_types.float4_e2m1f,

207 + b_q_type_id=scalar_types.float4_e2m1f.id,

208 size_m=reshaped_x.size(0),

209 size_n=padded_n,

210 size_k=padded_k,

@@ -237,13 +235,6 @@ def prepare_fp4_layer_for_marlin(

235

236 assert layer.weight.shape == (part_size_n, part_size_k // 2)

237

240 - device = layer.weight.device

241 -

242 - # WORKSPACE

243 - layer.workspace = marlin_make_workspace_new(

244 - device, existing=getattr(layer, "workspace", None)

245 - )

246 -

238 # WEIGHT

239 # Repack weights to marlin format

240 qweight = layer.weight.view(torch.int32).T.contiguous()

@@ -389,15 +380,9 @@ def prepare_nvfp4_moe_layer_for_marlin(

380 return x

381 return torch.nn.functional.pad(x, (0, (padded_N - N) // packing))

382

392 - device = w13.device

383 param_dtype = layer.params_dtype

384 is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1

385

396 - # WORKSPACE

397 - layer.workspace = marlin_make_workspace_new(

398 - device, 4, existing=getattr(layer, "workspace", None)

399 - )

400 -

386 # WEIGHT

387 # Repack weights to marlin format

388 def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:

@@ -481,12 +466,7 @@ def prepare_moe_fp4_layer_for_marlin(

466 k = layer.moe_config.hidden_dim

467 n = layer.moe_config.intermediate_size_per_partition

468

484 - # WORKSPACE

485 - device = layer.w13_weight.device

469 param_dtype = layer.params_dtype

487 - layer.workspace = marlin_make_workspace_new(

488 - device, 4, existing=getattr(layer, "workspace", None)

489 - )

470 is_a_8bit = input_dtype is not None and input_dtype.itemsize == 1

471

472 # WEIGHT

vllm/model_executor/layers/quantization/utils/marlin_utils_fp8.py

+8 -41

@@ -9,7 +9,6 @@ from vllm.logger import init_logger

9 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

10 USE_FP32_REDUCE_DEFAULT,

11 get_marlin_input_dtype,

12 - marlin_make_workspace_new,

12 marlin_moe_padded_intermediate,

13 marlin_pad_dim,

14 marlin_pad_qweight,

@@ -50,7 +49,7 @@ def apply_fp8_marlin_linear(

49 input: torch.Tensor,

50 weight: torch.Tensor,

51 weight_scale: torch.Tensor,

53 - workspace: torch.Tensor,

52 + workspace: torch.Tensor | None,

53 size_n: int,

54 size_k: int,

55 bias: torch.Tensor | None,

@@ -80,9 +79,8 @@ def apply_fp8_marlin_linear(

79 # inputs, a_scales = marlin_quant_input(inputs, torch.float8_e4m3fn)

80 raise RuntimeError("Marlin W8A8 is not supported.")

83 - output = ops.marlin_gemm(

82 + output = torch.ops.vllm.marlin_gemm(

83 a=inputs,

85 - c=None,

84 b_q_weight=weight,

85 b_bias=bias,

86 b_scales=weight_scale,

@@ -90,7 +88,7 @@ def apply_fp8_marlin_linear(

88 global_scale=None,

89 b_zeros=None,

90 workspace=workspace,

93 - b_q_type=scalar_types.float8_e4m3fn,

91 + b_q_type_id=scalar_types.float8_e4m3fn.id,

92 size_m=reshaped_x.size(0),

93 size_n=padded_n,

94 size_k=padded_k,

@@ -127,13 +125,6 @@ def prepare_fp8_layer_for_marlin(

125 else:

126 assert layer.weight.shape == (part_size_n, part_size_k)

127

130 - device = layer.weight.device

131 -

132 - # WORKSPACE

133 - layer.workspace = marlin_make_workspace_new(

134 - device, existing=getattr(layer, "workspace", None)

135 - )

136 -

128 # WEIGHT

129 # Repack weights to marlin format

130 qweight = pack_fp8_to_int32(layer.weight, size_k_first)

@@ -241,12 +232,7 @@ def prepare_fp8_moe_layer_for_marlin(

232 w13_weight_scale: torch.Tensor,

233 w2_weight_scale: torch.Tensor,

234 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:

244 - """Shuffle weights and scales into marlin format.

245 -

246 - Note that this function has the side effect of adding a `workspace`

247 - attribute to the layer. This `workspace` does not need to be

248 - registered as a Parameter as it is not used during weight reloading.

249 - """

235 + """Shuffle weights and scales into marlin format."""

236 logger.warning_once(

237 "Your GPU does not have native support for FP8 computation but "

238 "FP8 quantization is being used. Weight-only FP8 compression will "

@@ -272,12 +258,6 @@ def prepare_fp8_moe_layer_for_marlin(

258 w13_weight = _moe_pad_shard_rows(w13_weight, n, padded_n)

259 w2_weight = _moe_pad_last(w2_weight, n, padded_n)

260

275 - # WORKSPACE

276 - device = layer.w13_weight.device

277 - layer.workspace = marlin_make_workspace_new(

278 - device, 4, existing=getattr(layer, "workspace", None)

279 - )

280 -

261 # WEIGHT

262 # Repack weights to marlin format

263 def repack_weight(name: str, weight: torch.Tensor) -> torch.Tensor:

@@ -396,7 +376,7 @@ def apply_mxfp8_marlin_linear(

376 input: torch.Tensor,

377 weight: torch.Tensor,

378 weight_scale: torch.Tensor,

399 - workspace: torch.Tensor,

379 + workspace: torch.Tensor | None,

380 size_n: int,

381 size_k: int,

382 bias: torch.Tensor | None = None,

@@ -416,9 +396,8 @@ def apply_mxfp8_marlin_linear(

396 dtype=input.dtype,

397 )

398

419 - output = ops.marlin_gemm(

399 + output = torch.ops.vllm.marlin_gemm(

400 a=reshaped_x,

421 - c=None,

401 b_q_weight=weight,

402 b_bias=bias,

403 b_scales=weight_scale,

@@ -426,7 +405,7 @@ def apply_mxfp8_marlin_linear(

405 global_scale=None,

406 b_zeros=None,

407 workspace=workspace,

429 - b_q_type=scalar_types.float8_e4m3fn,

408 + b_q_type_id=scalar_types.float8_e4m3fn.id,

409 size_m=reshaped_x.size(0),

410 size_n=padded_n,

411 size_k=padded_k,

@@ -451,13 +430,6 @@ def prepare_mxfp8_layer_for_marlin(layer: torch.nn.Module) -> None:

430 group_size = 32 # MX standard block size

431 padded_n, padded_k = marlin_padded_nk(part_size_n, part_size_k, group_size)

432

454 - device = layer.weight.device

455 -

456 - # WORKSPACE

457 - layer.workspace = marlin_make_workspace_new(

458 - device, existing=getattr(layer, "workspace", None)

459 - )

460 -

433 # WEIGHT - repack FP8 weights to Marlin format

434 qweight = pack_fp8_to_int32(layer.weight, size_k_first=False)

435 qweight = qweight.T.contiguous()

@@ -512,7 +484,7 @@ def prepare_mxfp8_moe_layer_for_marlin(

484 """Repack MXFP8 MoE weights and scales into Marlin kernel format.

485

486 Args:

515 - layer: MoE layer (used to read params_dtype and attach workspace).

487 + layer: MoE layer.

488 w13: [E, 2*N, K] float8_e4m3fn weights.

489 w2: [E, K, N] float8_e4m3fn weights.

490 w13_scale: [E, 2*N, K//32] uint8 e8m0 scales.

@@ -537,13 +509,8 @@ def prepare_mxfp8_moe_layer_for_marlin(

509 n = padded_n

510 w13_n = w13.shape[1]

511

540 - device = w13.device

512 param_dtype = torch.get_default_dtype()

513

543 - layer.workspace = marlin_make_workspace_new(

544 - device, 4, existing=getattr(layer, "workspace", None)

545 - )

546 -

514 def repack_weight(weight: torch.Tensor, name: str) -> torch.Tensor:

515 if "w13" in name:

516 size_n, size_k = w13_n, k

vllm/v1/worker/workspace.py

+75 -21

@@ -3,11 +3,12 @@

4 import inspect

5 import os

6 -from collections.abc import Iterator

6 +from collections.abc import Callable, Hashable, Iterator

7 from contextlib import contextmanager

8 from contextvars import ContextVar

9 from itertools import accumulate

10 from math import prod

11 +from typing import TypeVar, cast

13 import torch

@@ -17,6 +18,7 @@ from vllm.utils.math_utils import round_up

18 from vllm.v1.worker.ubatching import dbo_current_ubatch_id

20 logger = init_logger(__name__)

21 +T = TypeVar("T")

24 def _compute_bytes(shape: tuple[int, ...], dtype: torch.dtype) -> int:

@@ -47,8 +49,8 @@ def use_workspace_lane(lane: int) -> Iterator[None]:

49 class WorkspaceManager:

50 """Manager for workspace allocation.

50 - Manages one workspace buffer per active ``(ubatch, lane)`` slot.

51 - Can be locked to prevent further growth during execution.

52 + Manages shared scratch and keyed persistent resources per ``(ubatch, lane)``.

53 + Can be locked to prevent growth or new resources during execution.

54 """

56 def __init__(

@@ -66,6 +68,9 @@ class WorkspaceManager:

68 self._current_workspaces: list[torch.Tensor | None] = [None] * (

69 self._num_ubatches * self._num_lanes

70 )

71 + self._persistent_resources: list[dict[Hashable, object]] = [

72 + {} for _ in self._current_workspaces

73 + ]

74 self._locked: bool = False

76 @staticmethod

@@ -76,10 +81,10 @@ class WorkspaceManager:

81 return workspace.numel() * workspace.element_size()

83 def lock(self) -> None:

79 - """Lock the workspace to prevent further growth.

84 + """Lock the workspace to prevent growth or new persistent resources.

81 - After locking, any attempt to allocate a larger workspace will raise

82 - an assertion error. This ensures workspace size is fixed during execution.

86 + Larger scratch requests and new persistent resources raise an assertion

87 + error. Existing resources retain their storage throughout execution.

88 """

89 self._locked = True

90 if envs.VLLM_DEBUG_WORKSPACE:

@@ -113,6 +118,63 @@ class WorkspaceManager:

118 """Check if workspace is locked."""

119 return self._locked

120

121 + def _get_workspace_id(self) -> int:

122 + lane = _workspace_lane.get()

123 + if lane >= self._num_lanes:

124 + raise RuntimeError(

125 + f"Workspace lane {lane} is not configured; manager has "

126 + f"{self._num_lanes} lane(s)."

127 + )

128 + return dbo_current_ubatch_id() * self._num_lanes + lane

129 +

130 + def get_persistent_resource(self, key: Hashable, factory: Callable[[], T]) -> T:

131 + """Return the resource cached for ``key`` in the current ubatch and lane.

132 +

133 + The first request calls ``factory`` and stores its result for the lifetime

134 + of this manager. Later requests return the same object without calling

135 + the factory again. Each ubatch/lane pair has a separate cache. Once the

136 + manager is locked, requesting a key missing from that cache raises an

137 + AssertionError.

138 +

139 + Keys must identify the resource type and configuration. Callers must not

140 + resize or replace tensor storage held by a cached resource. Uses of the

141 + same resource must not overlap; include the CUDA stream in the key when

142 + streams can execute concurrently.

143 +

144 + With torch.compile, call this inside a runtime custom op so resource

145 + selection happens during execution.

146 + """

147 + resources = self._persistent_resources[self._get_workspace_id()]

148 + if key not in resources:

149 + if self._locked:

150 + raise AssertionError(

151 + f"Workspace is locked but {key!r} was not allocated during "

152 + "warmup for the current ubatch and lane."

153 + )

154 + resources[key] = factory()

155 + return cast(T, resources[key])

156 +

157 + def get_persistent(

158 + self,

159 + key: Hashable,

160 + shape: tuple[int, ...],

161 + dtype: torch.dtype,

162 + *,

163 + zero_init: bool = False,

164 + ) -> torch.Tensor:

165 + """Get a fixed-shape tensor, optionally zeroed on first allocation only."""

166 + factory = torch.zeros if zero_init else torch.empty

167 + workspace = self.get_persistent_resource(

168 + (torch.Tensor, key),

169 + lambda: factory(shape, dtype=dtype, device=self._device),

170 + )

171 + if workspace.shape != shape or workspace.dtype != dtype:

172 + raise ValueError(

173 + f"Persistent workspace {key!r} has shape={workspace.shape}, "

174 + f"dtype={workspace.dtype}; requested shape={shape}, dtype={dtype}."

175 + )

176 + return workspace

177 +

178 def get_simultaneous(

179 self, *shapes_and_dtypes: tuple[tuple[int, ...], torch.dtype]

180 ) -> list[torch.Tensor]:

@@ -151,14 +213,7 @@ class WorkspaceManager:

213 The current workspace tensor.

214

215 """

154 - ubatch_id = dbo_current_ubatch_id()

155 - lane = _workspace_lane.get()

156 - if lane >= self._num_lanes:

157 - raise RuntimeError(

158 - f"Workspace lane {lane} is not configured; manager has "

159 - f"{self._num_lanes} lane(s)."

160 - )

161 - workspace_id = ubatch_id * self._num_lanes + lane

216 + workspace_id = self._get_workspace_id()

217 current_workspace = self._current_workspaces[workspace_id]

218 current_size = self._workspace_size_bytes(current_workspace)

219

@@ -218,8 +273,8 @@ class WorkspaceManager:

273 get_caller_info(),

274 current_size / _MB,

275 required_bytes / _MB,

221 - ubatch_id,

222 - lane,

276 + workspace_id // self._num_lanes,

277 + workspace_id % self._num_lanes,

278 )

279

280 return current_workspace

@@ -277,12 +332,11 @@ def init_workspace_manager(

332

333

334 def lock_workspace() -> None:

280 - """Lock the workspace to prevent further growth.

335 + """Lock the workspace to prevent growth or new persistent resources.

336

282 - After calling this function, any attempt to allocate a workspace larger

283 - than the current size will raise an AssertionError. This ensures that

284 - workspace size is fixed during execution and prevents unexpected memory

285 - allocations in the hot path.

337 + Larger scratch requests and new persistent resources raise an AssertionError.

338 + All required resources must be initialized in each execution slot during

339 + warmup, so their storage remains fixed during execution.

340

341 Example:

342 # During initialization