[Hybrid] Stage the postprocess inputs with a single loop over the request list (#48120)
Signed-off-by: Francesco Fusco <ffu@zurich.ibm.com> Co-authored-by: Claude <noreply@anthropic.com>
Francesco Fusco committed
Aug 3, 2026 at 22:24 UTC
755513d9a24fb639a8985172e645c7fe5a263154
2 files changed
+147
-76
tests/v1/worker/test_mamba_utils.py
+118
@@ -21,6 +21,7 @@ from vllm.v1.worker.mamba_utils import (
21
collect_mamba_copy_meta,
22
do_mamba_copy_block,
23
preprocess_mamba,
24
+ stage_postprocess_inputs_to_gpu,
25
)
26
27
MambaStateCopyFunc = Callable[..., Any]
@@ -371,6 +372,123 @@ def _make_gpu_ctx(
372
)
373
374
375
+# -----------------------------------------------------------------------------
376
+# stage_postprocess_inputs_to_gpu: single-pass staging into pinned views
377
+# -----------------------------------------------------------------------------
378
+
379
+
380
+def _make_staging_ctx(max_num_reqs: int, device: torch.device) -> MagicMock:
381
+ """Build a MambaSpecDecodeGPUContext stand-in exposing only the four
382
+ per-request staging buffers touched by stage_postprocess_inputs_to_gpu."""
383
+ ctx = MagicMock()
384
+ ctx.mamba_state_idx_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device)
385
+ ctx.num_scheduled_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device)
386
+ ctx.num_computed_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device)
387
+ ctx.num_draft_tokens_buf = _MockCpuGpuBuffer(max_num_reqs, torch.int32, device)
388
+ return ctx
389
+
390
+
391
+def test_stage_postprocess_inputs_to_gpu_fills_pinned_views():
392
+ """stage_postprocess_inputs_to_gpu writes each request's state_idx and
393
+ scheduled/computed/draft counts into slot `i` for `req_ids[i]`, leaves
394
+ slots past `num_reqs` untouched, and mirrors the values to the GPU buffers."""
395
+ device = torch.device("cpu")
396
+ max_num_reqs = 8
397
+ ctx = _make_staging_ctx(max_num_reqs, device)
398
+
399
+ # Any negative int32 works as a sentinel: all staged values (state_idx,
400
+ # scheduled/computed/draft token counts) are non-negative, so a negative
401
+ # entry surviving in slots [num_reqs:] proves the loop didn't overrun.
402
+ sentinel = np.int32(-1)
403
+ bufs = (
404
+ ctx.mamba_state_idx_buf,
405
+ ctx.num_scheduled_tokens_buf,
406
+ ctx.num_computed_tokens_buf,
407
+ ctx.num_draft_tokens_buf,
408
+ )
409
+ for buf in bufs:
410
+ buf.np[:] = sentinel
411
+
412
+ req_ids = ["req_a", "req_b", "req_c"]
413
+ num_reqs = len(req_ids)
414
+ scheduler_output = _make_postprocess_scheduler_output(
415
+ req_ids=req_ids,
416
+ num_scheduled_tokens={"req_a": 5, "req_b": 12, "req_c": 3},
417
+ # req_b intentionally absent -> defaults to 0 drafts.
418
+ scheduled_spec_decode_tokens={
419
+ "req_a": [1, 2],
420
+ "req_c": [1, 2, 3, 4],
421
+ },
422
+ )
423
+ requests = _make_requests(
424
+ req_ids=req_ids,
425
+ num_computed_tokens=[10, 20, 30],
426
+ block_ids_per_req=[[0], [0], [0]],
427
+ )
428
+ mamba_state_idx = {"req_a": 100, "req_b": 200, "req_c": 300}
429
+ # A trailing entry past num_reqs must not be read.
430
+ req_ids_padded = req_ids + ["not_scheduled"]
431
+
432
+ stage_postprocess_inputs_to_gpu(
433
+ ctx,
434
+ scheduler_output,
435
+ req_ids_padded,
436
+ num_reqs,
437
+ requests,
438
+ mamba_state_idx,
439
+ )
440
+
441
+ np.testing.assert_array_equal(
442
+ ctx.mamba_state_idx_buf.np[:num_reqs], [100, 200, 300]
443
+ )
444
+ np.testing.assert_array_equal(
445
+ ctx.num_scheduled_tokens_buf.np[:num_reqs], [5, 12, 3]
446
+ )
447
+ np.testing.assert_array_equal(
448
+ ctx.num_computed_tokens_buf.np[:num_reqs], [10, 20, 30]
449
+ )
450
+ np.testing.assert_array_equal(ctx.num_draft_tokens_buf.np[:num_reqs], [2, 0, 4])
451
+ for buf in bufs:
452
+ assert (buf.np[num_reqs:] == sentinel).all()
453
+
454
+ assert torch.equal(
455
+ ctx.mamba_state_idx_buf.gpu[:num_reqs],
456
+ torch.tensor([100, 200, 300], dtype=torch.int32),
457
+ )
458
+ assert torch.equal(
459
+ ctx.num_draft_tokens_buf.gpu[:num_reqs],
460
+ torch.tensor([2, 0, 4], dtype=torch.int32),
461
+ )
462
+
463
+
464
+def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx():
465
+ """If preprocess_mamba didn't populate mamba_state_idx for a req in the
466
+ batch, staging must fail loudly rather than silently writing a stale index."""
467
+ device = torch.device("cpu")
468
+ ctx = _make_staging_ctx(max_num_reqs=4, device=device)
469
+ scheduler_output = _make_postprocess_scheduler_output(
470
+ req_ids=["req_a"],
471
+ num_scheduled_tokens={"req_a": 1},
472
+ )
473
+ requests = _make_requests(["req_a"], [0], [[0]])
474
+
475
+ # mamba_state_idx has an entry for a *different* request but not for
476
+ # req_a (the one being staged). This is the realistic failure mode:
477
+ # preprocess_mamba ran, populated some reqs, but missed this one.
478
+ # The safety assert in stage_postprocess_inputs_to_gpu must fire on
479
+ # req_a's missing entry rather than silently writing None.
480
+ mamba_state_idx: dict[str, int] = {"other_req": 42}
481
+ with pytest.raises(AssertionError, match="mamba_state_idx missing entry"):
482
+ stage_postprocess_inputs_to_gpu(
483
+ ctx,
484
+ scheduler_output,
485
+ ["req_a"],
486
+ 1,
487
+ requests,
488
+ mamba_state_idx,
489
+ )
490
+
491
+
492
def _run_gpu_postprocess(
493
gpu_ctx: MambaSpecDecodeGPUContext,
494
*,
vllm/v1/worker/mamba_utils.py
+29
-76
@@ -1238,97 +1238,50 @@ def postprocess_mamba_align_gpu(
1238
)
1239
1240
1241
-def stage_postprocess_metadata_to_gpu(
1241
+def stage_postprocess_inputs_to_gpu(
1242
+ ctx: MambaSpecDecodeGPUContext,
1243
scheduler_output: SchedulerOutput,
1244
req_ids: list[str],
1245
num_reqs: int,
1246
requests: dict[str, CachedRequestState],
1246
- num_scheduled_tokens_buf: CpuGpuBuffer,
1247
- num_computed_tokens_buf: CpuGpuBuffer,
1248
- num_draft_tokens_buf: CpuGpuBuffer,
1249
-) -> None:
1250
- """Stage per-request postprocess metadata into GPU buffers (non-blocking).
1251
-
1252
- Walks ``req_ids[:num_reqs]`` in batch order and writes each request's
1253
- scheduled/computed/draft token counts into the matching pinned numpy
1254
- views, then issues three non-blocking H→D copies. These values don't
1255
- change between ``_prepare_inputs`` and ``_update_states_after_model_execute``.
1256
- The fused postprocess kernel indexes the resulting GPU tensors
1257
- by ``req_idx``.
1258
- """
1259
- scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
1260
- num_scheduled = scheduler_output.num_scheduled_tokens
1261
- scheduled_np = num_scheduled_tokens_buf.np
1262
- computed_np = num_computed_tokens_buf.np
1263
- draft_np = num_draft_tokens_buf.np
1264
- for i in range(num_reqs):
1265
- req_id = req_ids[i]
1266
- scheduled_np[i] = num_scheduled[req_id]
1267
- computed_np[i] = requests[req_id].num_computed_tokens
1268
- draft_np[i] = len(scheduled_spec_tokens.get(req_id, []))
1269
- num_scheduled_tokens_buf.copy_to_gpu(num_reqs)
1270
- num_computed_tokens_buf.copy_to_gpu(num_reqs)
1271
- num_draft_tokens_buf.copy_to_gpu(num_reqs)
1272
-
1273
-
1274
-def stage_mamba_state_idx_to_gpu(
1247
mamba_state_idx: dict[str, int],
1276
- req_ids: list[str],
1277
- num_reqs: int,
1278
- gpu_buf: CpuGpuBuffer,
1248
) -> None:
1280
- """Materialize ``mamba_state_idx`` into ``gpu_buf`` and copy to GPU.
1249
+ """Stage all per-request inputs the fused mamba postprocess kernel reads.
1250
1282
- Walks ``req_ids[:num_reqs]`` in batch order, writing each request's block
1283
- index into the buffer's pinned numpy view, then issues a non-blocking H→D
1284
- copy. The fused kernel indexes the resulting GPU tensor by ``req_idx``.
1251
+ Walks ``req_ids[:num_reqs]`` once, writing each request's mamba block
1252
+ index and scheduled/computed/draft token counts into the matching pinned
1253
+ numpy views, then issues four non-blocking H→D copies. The fused kernel
1254
+ indexes the resulting GPU tensors by ``req_idx``. Buffers live on ``ctx``
1255
+ and only exist when the postprocess kernel is enabled.
1256
1257
Invariant: ``preprocess_mamba`` must have run first for the same batch so
1258
that every ``req_ids[i]`` has an entry in ``mamba_state_idx``.
1259
"""
1289
- np_view = gpu_buf.np
1260
+ assert ctx.mamba_state_idx_buf is not None
1261
+ assert ctx.num_scheduled_tokens_buf is not None
1262
+ assert ctx.num_computed_tokens_buf is not None
1263
+ assert ctx.num_draft_tokens_buf is not None
1264
+
1265
+ scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens
1266
+ num_scheduled = scheduler_output.num_scheduled_tokens
1267
+ state_idx_np = ctx.mamba_state_idx_buf.np
1268
+ scheduled_np = ctx.num_scheduled_tokens_buf.np
1269
+ computed_np = ctx.num_computed_tokens_buf.np
1270
+ draft_np = ctx.num_draft_tokens_buf.np
1271
+
1272
for i in range(num_reqs):
1273
req_id = req_ids[i]
1274
state_idx = mamba_state_idx.get(req_id)
1275
assert state_idx is not None, (
1276
f"mamba_state_idx missing entry for {req_id!r}; "
1295
- "preprocess_mamba must run before stage_mamba_state_idx_to_gpu"
1277
+ "preprocess_mamba must run before stage_postprocess_inputs_to_gpu"
1278
)
1297
- np_view[i] = state_idx
1298
- gpu_buf.copy_to_gpu(num_reqs)
1299
-
1300
-
1301
-def stage_postprocess_inputs_to_gpu(
1302
- ctx: MambaSpecDecodeGPUContext,
1303
- scheduler_output: SchedulerOutput,
1304
- req_ids: list[str],
1305
- num_reqs: int,
1306
- requests: dict[str, CachedRequestState],
1307
- mamba_state_idx: dict[str, int],
1308
-) -> None:
1309
- """Stage all per-request inputs the fused mamba postprocess kernel reads.
1279
+ state_idx_np[i] = state_idx
1280
+ scheduled_np[i] = num_scheduled[req_id]
1281
+ computed_np[i] = requests[req_id].num_computed_tokens
1282
+ draft_np[i] = len(scheduled_spec_tokens.get(req_id, []))
1283
1311
- Bundles ``stage_mamba_state_idx_to_gpu`` and
1312
- ``stage_postprocess_metadata_to_gpu`` into a single call so the runner
1313
- has one entry point for postprocess staging. Buffers live on ``ctx``
1314
- and only exist when the postprocess kernel is enabled.
1315
- """
1316
- assert ctx.mamba_state_idx_buf is not None
1317
- assert ctx.num_scheduled_tokens_buf is not None
1318
- assert ctx.num_computed_tokens_buf is not None
1319
- assert ctx.num_draft_tokens_buf is not None
1320
- stage_mamba_state_idx_to_gpu(
1321
- mamba_state_idx,
1322
- req_ids,
1323
- num_reqs,
1324
- ctx.mamba_state_idx_buf,
1325
- )
1326
- stage_postprocess_metadata_to_gpu(
1327
- scheduler_output,
1328
- req_ids,
1329
- num_reqs,
1330
- requests,
1331
- ctx.num_scheduled_tokens_buf,
1332
- ctx.num_computed_tokens_buf,
1333
- ctx.num_draft_tokens_buf,
1334
- )
1284
+ ctx.mamba_state_idx_buf.copy_to_gpu(num_reqs)
1285
+ ctx.num_scheduled_tokens_buf.copy_to_gpu(num_reqs)
1286
+ ctx.num_computed_tokens_buf.copy_to_gpu(num_reqs)
1287
+ ctx.num_draft_tokens_buf.copy_to_gpu(num_reqs)