@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 153 lines 5.7 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """Tests for the Model Runner V2 prompt-embeds overlay (PromptEmbedsState).
4
5 The overlay kernel reads each request's GPU-resident prompt embeddings through
6 a per-request pointer table and writes the rows scheduled this step into
7 `inputs_embeds`, honoring chunked prefill (`num_computed_tokens` offset), the
8 prompt/decode boundary (rows past the embeds length untouched), and the
9 mixed-mode `prompt_is_token_ids` mask (token-id rows keep the base embedding).
10 """
11
12 from dataclasses import dataclass
13
14 import pytest
15 import torch
16
17 pytest.importorskip("triton")
18
19 from vllm.platforms import current_platform
20
21 if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
22 pytest.skip(
23 "An accelerator is required for prompt-embeds overlay tests",
24 allow_module_level=True,
25 )
26
27 from vllm.v1.worker.gpu.model_states.prompt_embeds import PromptEmbedsState
28
29 HIDDEN = 24
30 MAX_NUM_REQS = 8
31 DEVICE = torch.device(current_platform.device_type)
32
33
34 @dataclass
35 class _NewReqData:
36 req_id: str
37 prompt_embeds: torch.Tensor | None
38 prompt_is_token_ids: list[bool] | None = None
39
40
41 @dataclass
42 class _Batch:
43 num_reqs: int
44 num_scheduled_tokens: torch.Tensor # np-like, only .max() is used
45 idx_mapping: torch.Tensor
46 query_start_loc: torch.Tensor
47
48
49 def _make_state() -> PromptEmbedsState:
50 return PromptEmbedsState(MAX_NUM_REQS, HIDDEN, torch.float32, DEVICE)
51
52
53 def _batch(num_scheduled: list[int], idx_mapping: list[int]) -> _Batch:
54 query_start_loc = [0]
55 for n in num_scheduled:
56 query_start_loc.append(query_start_loc[-1] + n)
57 return _Batch(
58 num_reqs=len(num_scheduled),
59 num_scheduled_tokens=torch.tensor(num_scheduled, dtype=torch.int32),
60 idx_mapping=torch.tensor(idx_mapping, dtype=torch.int64, device=DEVICE),
61 query_start_loc=torch.tensor(query_start_loc, dtype=torch.int32, device=DEVICE),
62 )
63
64
65 def _apply(
66 state: PromptEmbedsState,
67 batch: _Batch,
68 num_computed: list[int],
69 num_tokens: int,
70 ) -> tuple[torch.Tensor, torch.Tensor]:
71 """Run the overlay on a fresh base buffer; return (result, base)."""
72 num_computed_tokens = torch.zeros(MAX_NUM_REQS, dtype=torch.int32, device=DEVICE)
73 for batch_idx, req_index in enumerate(batch.idx_mapping.tolist()):
74 num_computed_tokens[req_index] = num_computed[batch_idx]
75 base = torch.randn(num_tokens, HIDDEN, dtype=torch.float32, device=DEVICE)
76 inputs_embeds = base.clone()
77 state.apply(batch, num_computed_tokens, inputs_embeds)
78 torch.accelerator.synchronize()
79 return inputs_embeds, base
80
81
82 def test_overlay_chunked_prefill_and_decode():
83 """Rows within the embeds range come from prompt_embeds at the
84 num_computed offset; requests without embeds and requests past their
85 embeds length (decode) keep the base embedding."""
86 state = _make_state()
87 embeds_a = torch.randn(6, HIDDEN, dtype=torch.float32)
88 embeds_b = torch.randn(5, HIDDEN, dtype=torch.float32)
89 state.add_request(0, _NewReqData("a", embeds_a))
90 state.add_request(1, _NewReqData("b", embeds_b))
91 state.add_request(2, _NewReqData("c", None))
92 state.apply_staged_writes()
93
94 # a: chunk [2, 6) of its embeds; b: fully decoded; c: no embeds.
95 batch = _batch(num_scheduled=[4, 1, 3], idx_mapping=[0, 1, 2])
96 out, base = _apply(state, batch, num_computed=[2, 7, 1], num_tokens=8)
97
98 torch.testing.assert_close(out[0:4], embeds_a[2:6].to(DEVICE))
99 torch.testing.assert_close(out[4:8], base[4:8])
100
101
102 def test_overlay_clamps_to_embeds_length():
103 """A window straddling the end of the prompt embeds writes only the
104 in-range rows (e.g. final prefill chunk + sampled token)."""
105 state = _make_state()
106 embeds = torch.randn(4, HIDDEN, dtype=torch.float32)
107 state.add_request(3, _NewReqData("a", embeds))
108 state.apply_staged_writes()
109
110 batch = _batch(num_scheduled=[3], idx_mapping=[3])
111 out, base = _apply(state, batch, num_computed=[2], num_tokens=3)
112
113 torch.testing.assert_close(out[0:2], embeds[2:4].to(DEVICE))
114 torch.testing.assert_close(out[2:3], base[2:3])
115
116
117 def test_overlay_respects_is_token_ids_mask():
118 """Mixed mode: positions marked as real token ids keep the base
119 embedding; only embed positions are overwritten."""
120 state = _make_state()
121 embeds = torch.randn(5, HIDDEN, dtype=torch.float32)
122 is_token_ids = [True, False, False, True, False]
123 state.add_request(0, _NewReqData("a", embeds, is_token_ids))
124 state.apply_staged_writes()
125
126 batch = _batch(num_scheduled=[5], idx_mapping=[0])
127 out, base = _apply(state, batch, num_computed=[0], num_tokens=5)
128
129 embeds_gpu = embeds.to(DEVICE)
130 for pos, is_token in enumerate(is_token_ids):
131 expected = base[pos] if is_token else embeds_gpu[pos]
132 torch.testing.assert_close(out[pos], expected)
133
134
135 def test_index_reuse_clears_stale_entry():
136 """A request added at a previously-used index without embeds must not
137 inherit the prior occupant's pointer-table entry."""
138 state = _make_state()
139 state.add_request(0, _NewReqData("a", torch.randn(4, HIDDEN, dtype=torch.float32)))
140 # A second live embeds request so the kernel actually launches (the
141 # overlay is skipped entirely when no request holds embeds).
142 other = torch.randn(2, HIDDEN, dtype=torch.float32)
143 state.add_request(1, _NewReqData("other", other))
144 state.apply_staged_writes()
145 state.remove_request("a")
146 state.add_request(0, _NewReqData("b", None))
147 state.apply_staged_writes()
148
149 batch = _batch(num_scheduled=[2, 2], idx_mapping=[0, 1])
150 out, base = _apply(state, batch, num_computed=[0, 0], num_tokens=4)
151
152 torch.testing.assert_close(out[0:2], base[0:2])
153 torch.testing.assert_close(out[2:4], other.to(DEVICE))