@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 91dab0eb7f

[ROCm][Perf] Allocate the pinned PLE prefetch buffer lazily (#58797)

Signed-off-by: Mathew Odden <modden@redhat.com> Co-authored-by: opencode+glm-5.3+vllm <opencode+glm-5.3+vllm@example.com>

Mathew Odden committed Sep 30, 2026 at 04:43 UTC 91dab0eb7f42efd8580086935a7c0768a594f575

2 files changed +155 -15

tests/models/qwen4_exp/test_ple.py

+106

@@ -609,6 +609,32 @@ def test_pinned_embedding_forward_finalizes_prefetched_output(

609 assert torch.equal(output, expected)

610

611

612 +def test_pinned_embedding_forward_requires_prior_start_prefetch() -> None:

613 + """Forward before any start_prefetch must fail loudly, not read garbage."""

614 + embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding)

615 + nn.Module.__init__(embedding)

616 + embedding._prefetch_buffer = None

617 + embedding._prefetch_stream = None

618 + embedding._output_dim = 6

619 + hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16)

620 +

621 + with pytest.raises(RuntimeError, match="prior start_prefetch"):

622 + embedding(hidden_states)

623 +

624 +

625 +def test_pinned_embedding_finalize_requires_prior_start_prefetch() -> None:

626 + """_finalize_prefetch before any start_prefetch must fail loudly."""

627 + embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding)

628 + nn.Module.__init__(embedding)

629 + embedding._prefetch_buffer = None

630 + embedding._prefetch_stream = None

631 + embedding.tp_size = 1

632 + output = torch.zeros(2, 6, dtype=torch.bfloat16)

633 +

634 + with pytest.raises(RuntimeError, match="prior start_prefetch"):

635 + embedding._finalize_prefetch(torch.empty(4, 2, 3), output)

636 +

637 +

638 def test_pinned_fp8_embedding_uses_int8_for_parallel_reduce() -> None:

639 embedding = Qwen4ExpPLEPinnedHostEmbedding.__new__(Qwen4ExpPLEPinnedHostEmbedding)

640 nn.Module.__init__(embedding)

@@ -723,6 +749,67 @@ def test_ple_pinned_embedding_loads_on_cpu_and_looks_up_through_uva(

749 )

750

751

752 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")

753 +def test_pinned_embedding_start_prefetch_allocates_lazily_and_feeds_forward(

754 + monkeypatch: pytest.MonkeyPatch,

755 +) -> None:

756 + """start_prefetch lazily allocates the stream/buffer, then forward consumes it.

757 +

758 + Exercises the full pinned-prefetch contract end to end: first call

759 + allocates the side stream and buffer with the constructor geometry,

760 + and forward finalizes the lookup into the returned output.

761 + """

762 + _mock_etp_group(monkeypatch)

763 + monkeypatch.setattr(embedding_module, "get_tensor_model_parallel_rank", lambda: 0)

764 + monkeypatch.setattr(

765 + embedding_module, "get_tensor_model_parallel_world_size", lambda: 1

766 + )

767 + monkeypatch.setattr(parameter_module, "get_tensor_model_parallel_rank", lambda: 0)

768 + monkeypatch.setattr(

769 + parameter_module, "get_tensor_model_parallel_world_size", lambda: 1

770 + )

771 +

772 + with torch.device("cuda:0"):

773 + embedding = Qwen4ExpPLEPinnedHostEmbedding(

774 + 4,

775 + 3,

776 + params_dtype=torch.bfloat16,

777 + padding_size=1,

778 + prefix="test.ple_embedding",

779 + embedding_method=Qwen4ExpPLEUnquantizedEmbeddingMethod(),

780 + num_ngram_heads=2,

781 + max_total_tokens=4,

782 + )

783 + assert embedding._prefetch_buffer is None

784 + assert embedding._prefetch_stream is None

785 +

786 + loaded_weight = torch.arange(12, dtype=torch.float32).reshape(4, 3)

787 + loaded_weight = loaded_weight.to(torch.bfloat16)

788 + copy_ple_embedding_shard_(

789 + embedding.weight,

790 + loaded_weight,

791 + checkpoint_start=0,

792 + tp_start=0,

793 + tp_end=4,

794 + )

795 +

796 + ngram_ids = torch.tensor([[3, 0], [1, 2]], device="cuda:0")

797 + hidden_states = torch.zeros(2, 4, dtype=torch.bfloat16, device="cuda:0")

798 + embedding.start_prefetch(hidden_states, ngram_ids)

799 +

800 + buffer = embedding._prefetch_buffer

801 + assert buffer is not None

802 + assert buffer.shape == (4, 2, 3)

803 + assert buffer.dtype == torch.bfloat16

804 + assert buffer.device.type == "cuda"

805 + assert embedding._prefetch_stream is not None

806 +

807 + output = embedding(hidden_states)

808 +

809 + expected = loaded_weight[ngram_ids.cpu()].to(device="cuda:0").flatten(-2)

810 + torch.testing.assert_close(output, expected, rtol=0, atol=0)

811 +

812 +

813 def test_ple_fp8_embedding_supports_mixed_precision_config() -> None:

814 prefix = "model.language_model.layers.1.ple.ple_embedding.ngram_embedding"

815 quant_config = ModelOptMixedPrecisionConfig.from_config(

@@ -2004,6 +2091,25 @@ def test_amd_fp8_embedding_loads_checkpoint_shards_and_global_scale(

2091 )

2092

2093

2094 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA or ROCm")

2095 +def test_amd_pinned_embedding_prefetch_machinery_is_lazy(

2096 + monkeypatch: pytest.MonkeyPatch,

2097 +) -> None:

2098 + """The pinned embedding must not allocate the prefetch stream/buffer in init.

2099 +

2100 + Paths that never prefetch (AMD uses the synchronous pinned-lookup op)

2101 + keep the stream/buffer unallocated for the process lifetime. NVIDIA's

2102 + first ``start_prefetch`` — the eager profile run, always before cudagraph

2103 + capture — allocates them lazily.

2104 + """

2105 + module, _ = _build_amd_ngram_embedding(

2106 + monkeypatch, device="cuda:0", cpu_offload=True, fp8_checkpoint=False

2107 + )

2108 + embedding = module.ngram_embedding

2109 + assert embedding._prefetch_buffer is None

2110 + assert embedding._prefetch_stream is None

2111 +

2112 +

2113 @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA or ROCm")

2114 @pytest.mark.parametrize("fp8_checkpoint", [False, True], ids=["bf16", "fp8"])

2115 def test_amd_pinned_embedding_output_written_under_compile(

vllm/models/qwen4_exp/common/ngram_embedding.py

+49 -15

@@ -7,6 +7,7 @@ n-gram embedding table can be kept in pinned host memory and looked up through

7 Unified Virtual Addressing on any CUDA-alike platform.

8 """

10 +import threading

11 from abc import ABC, abstractmethod

12 from typing import ClassVar

@@ -415,14 +416,11 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding):

416 )

417 self._uva_weight = get_accelerator_view_from_cpu_tensor(self.weight)

418 self._block_d = triton.next_power_of_2(self.embedding_dim)

418 - self._prefetch_stream = torch.cuda.Stream(device=self._uva_weight.device)

419 - self._prefetch_buffer = torch.empty(

420 - max_total_tokens * self.etp_data_parallel_size,

421 - num_ngram_heads,

422 - self.embedding_dim,

423 - dtype=self.weight.dtype,

424 - device=self._uva_weight.device,

425 - )

419 + self._prefetch_stream: torch.cuda.Stream | None = None

420 + self._prefetch_buffer: torch.Tensor | None = None

421 + self._prefetch_alloc_lock = threading.Lock()

422 + self._prefetch_rows = max_total_tokens * self.etp_data_parallel_size

423 + self._num_ngram_heads = num_ngram_heads

424 self._output_dim = num_ngram_heads * self.embedding_dim

425

426 def allocate_embedding_weight(

@@ -506,10 +504,42 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding):

504 ngram_ids: torch.Tensor,

505 ) -> None:

506 """Gather ETP IDs and launch their UVA lookup on the side stream."""

507 + buffer = self._prefetch_buffer

508 + if buffer is None:

509 + # First use allocates. The eager profile run always precedes

510 + # cudagraph capture, so allocation never happens mid-capture;

511 + # the lock keeps concurrent first callers from tearing the

512 + # stream/buffer pair.

513 + with self._prefetch_alloc_lock:

514 + buffer = self._prefetch_buffer

515 + if buffer is None:

516 + if torch.cuda.is_current_stream_capturing():

517 + raise RuntimeError(

518 + "pinned PLE prefetch buffer must be allocated "

519 + "eagerly, before cudagraph capture"

520 + )

521 + self._prefetch_stream = torch.cuda.Stream(

522 + device=self._uva_weight.device

523 + )

524 + buffer = torch.empty(

525 + self._prefetch_rows,

526 + self._num_ngram_heads,

527 + self.embedding_dim,

528 + dtype=self.weight.dtype,

529 + device=self._uva_weight.device,

530 + )

531 + self._prefetch_buffer = buffer

532 + prefetch_stream = self._prefetch_stream

533 + if prefetch_stream is None:

534 + raise RuntimeError("pinned PLE prefetch stream was not allocated")

535 slot_size, _ = self._get_dp_gather_slot(ngram_ids.shape[0])

536 gathered_ids = self._gather_dp_ids(ngram_ids, slot_size)

511 - active_output = self._prefetch_buffer[: gathered_ids.shape[0]]

512 - prefetch_stream = self._prefetch_stream

537 + if gathered_ids.shape[0] > buffer.shape[0]:

538 + raise ValueError(

539 + f"pinned PLE prefetch buffer holds {buffer.shape[0]} rows, "

540 + f"but the batch needs {gathered_ids.shape[0]}"

541 + )

542 + active_output = buffer[: gathered_ids.shape[0]]

543 prefetch_stream.wait_stream(torch.cuda.current_stream())

544 gathered_ids.record_stream(prefetch_stream)

545 with torch.cuda.stream(prefetch_stream):

@@ -522,7 +552,10 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding):

552 output: torch.Tensor,

553 ) -> None:

554 """Join the side stream, reduce ETP shards, and select local rows."""

525 - torch.cuda.current_stream().wait_stream(self._prefetch_stream)

555 + prefetch_stream = self._prefetch_stream

556 + if prefetch_stream is None:

557 + raise RuntimeError("pinned PLE finalize requires a prior start_prefetch")

558 + torch.cuda.current_stream().wait_stream(prefetch_stream)

559 slot_size, slot_offset = self._get_dp_gather_slot(output.shape[0])

560 active_output = prefetch_output[: slot_size * self.etp_data_parallel_size]

561 embeddings = self._reduce_etp_embeddings(active_output)

@@ -535,8 +568,9 @@ class Qwen4ExpPLEPinnedHostEmbedding(Qwen4ExpPLEEmbedding):

568

569 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:

570 """Finish the pinned lookup into graph-owned output storage."""

538 - output = self._prefetch_buffer.new_empty(

539 - (hidden_states.shape[0], self._output_dim)

540 - )

541 - self._finalize_prefetch(self._prefetch_buffer, output)

571 + buffer = self._prefetch_buffer

572 + if buffer is None:

573 + raise RuntimeError("pinned PLE lookup requires a prior start_prefetch")

574 + output = buffer.new_empty((hidden_states.shape[0], self._output_dim))

575 + self._finalize_prefetch(buffer, output)

576 return output