[Bugfix][Qwen4Exp] Keep pinned PLE prefetch ids out of the CUDA graph pool (#58489)
Signed-off-by: Juntian Liu <Juntianl777@gmail.com> Co-authored-by: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Juntian Liu committed
Sep 24, 2026 at 19:51 UTC
48d8880d098c567c64150203110b6a6a4221d44f
2 files changed
+42
tests/models/qwen4_exp/test_ple.py
+35
@@ -35,6 +35,7 @@ from vllm.models.qwen4_exp.nvidia.ngram_embedding import (
35
Qwen4ExpPLEUnquantizedEmbeddingMethod,
36
)
37
from vllm.models.qwen4_exp.nvidia.ple_layer import Qwen4ExpPLELayer
38
+from vllm.utils.torch_utils import weak_ref_tensor
39
from vllm.v1.attention.backends.short_conv_attn import (
40
PleShortConvAttentionMetadata,
41
)
@@ -885,6 +886,40 @@ def test_fused_ngram_ids_correctness(
886
assert torch.all((actual >= offsets) & (actual < offsets + sizes))
887
888
889
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="fused PLE needs CUDA")
890
+def test_ngram_prefetch_ids_outlive_start_prefetch() -> None:
891
+ """Eager breaks hand the side-stream lookup weak refs, so the ids must stay
892
+ valid after start_prefetch returns rather than be reused from the pool."""
893
+ device = torch.device("cuda")
894
+ query_start_loc = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
895
+ input_ids = torch.arange(20, 25, dtype=torch.int32, device=device)
896
+ ngram_context = torch.tensor([[11, 12], [13, 14]], dtype=torch.int32, device=device)
897
+ params = _ngram_hash_params(device, ngram_context.shape[1])
898
+ module = Qwen4ExpNGramEmbedding.__new__(Qwen4ExpNGramEmbedding)
899
+ nn.Module.__init__(module)
900
+ for name in ("layer_multipliers", "ngram_heads_vocab_sizes", "ngram_heads_offsets"):
901
+ module.register_buffer(name, params[name])
902
+ module.eos_token_id = params["eos_token_id"]
903
+ module.heads_per_ngram = params["heads_per_ngram"]
904
+ num_heads = params["ngram_heads_vocab_sizes"].numel()
905
+ module._prefetch_ids = torch.empty(8, num_heads, dtype=torch.long, device=device)
906
+ prefetched: list[torch.Tensor] = []
907
+ module.ngram_embedding = SimpleNamespace(
908
+ supports_prefetch=True,
909
+ start_prefetch=lambda _, ids: prefetched.append(weak_ref_tensor(ids)),
910
+ )
911
+ expected = _reference_ngram_ids(input_ids, query_start_loc, ngram_context, **params)
912
+
913
+ # A fresh pool makes the next same-size allocation reuse any freed ids,
914
+ # like a later CUDA graph segment would.
915
+ pool = torch.cuda.MemPool()
916
+ with torch.cuda.use_mem_pool(pool):
917
+ module.start_prefetch(None, input_ids, query_start_loc, ngram_context)
918
+ torch.full_like(expected, -1)
919
+
920
+ assert torch.equal(prefetched[0], expected)
921
+
922
+
923
def _short_conv_dilated_decode_pytorch(
924
x_d: torch.Tensor,
925
conv_state: torch.Tensor,
vllm/models/qwen4_exp/nvidia/ngram_embedding.py
+7
@@ -716,6 +716,12 @@ class Qwen4ExpNGramEmbedding(nn.Module):
716
max_total_tokens=max_total_tokens,
717
data_parallel_rank=data_parallel_rank,
718
)
719
+ if self.ngram_embedding.supports_prefetch:
720
+ # The side-stream lookup outlives eager-break args, whose
721
+ # graph-pool storage later segments may reuse.
722
+ self._prefetch_ids = torch.empty(
723
+ max_total_tokens, self.ngram_heads, dtype=torch.long
724
+ )
725
weight = self.ngram_embedding.weight
726
logger.info(
727
"Initialized PLE embedding %s: quantization_method=%s, "
@@ -864,6 +870,7 @@ class Qwen4ExpNGramEmbedding(nn.Module):
870
input_ids,
871
query_start_loc,
872
ngram_context,
873
+ output=self._prefetch_ids[: input_ids.numel()],
874
)
875
embedding.start_prefetch(hidden_states, ngram_ids)
876