5502c3b52d · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -0,0 +1,499 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +"""Unit tests for sequence and token pooling methods and their factories."""

4 +

5 +from dataclasses import dataclass

6 +from unittest.mock import patch

7 +

8 +import pytest

9 +import torch

10 +

11 +from vllm.model_executor.layers.pooler.seqwise.methods import (

12 + CLSPool,

13 + LastPool,

14 + MeanPool,

15 + get_seq_pooling_method,

16 +)

17 +from vllm.model_executor.layers.pooler.tokwise.methods import (

18 + AllPool,

19 + StepPool,

20 + get_tok_pooling_method,

21 +)

22 +from vllm.pooling_params import PoolingParams

23 +from vllm.v1.pool.metadata import PoolingCursor, PoolingMetadata, PoolingStates

24 +

25 +_CPU = torch.device("cpu")

26 +

27 +

28 +def _make_pooling_cursor(

29 + prompt_lens: list[int],

30 + *,

31 + num_scheduled_tokens: list[int] | None = None,

32 + seq_lens: list[int] | None = None,

33 + device: torch.device = _CPU,

34 +) -> PoolingCursor:

35 + """Build a PoolingCursor from a list of per-sequence prompt lengths."""

36 + prompt_lens_cpu = torch.tensor(prompt_lens, dtype=torch.long)

37 + if num_scheduled_tokens is None:

38 + num_scheduled_tokens_cpu = prompt_lens_cpu.clone()

39 + else:

40 + num_scheduled_tokens_cpu = torch.tensor(num_scheduled_tokens, dtype=torch.long)

41 + if seq_lens is None:

42 + seq_lens_cpu = prompt_lens_cpu.clone()

43 + else:

44 + seq_lens_cpu = torch.tensor(seq_lens, dtype=torch.long)

45 +

46 + cumsum = torch.zeros(len(prompt_lens) + 1, dtype=torch.long, device=device)

47 + torch.cumsum(num_scheduled_tokens_cpu, dim=0, out=cumsum[1:])

48 +

49 + return PoolingCursor(

50 + first_token_indices_gpu=cumsum[: len(prompt_lens)].to(device),

51 + last_token_indices_gpu=(cumsum[1:] - 1).to(device),

52 + prompt_lens_cpu=prompt_lens_cpu,

53 + seq_lens_cpu=seq_lens_cpu,

54 + num_scheduled_tokens_cpu=num_scheduled_tokens_cpu,

55 + )

56 +

57 +

58 +def _make_metadata(

59 + prompt_lens: list[int],

60 + *,

61 + tasks: list[str] | None = None,

62 + token_ids: list[list[int]] | None = None,

63 + pooling_params: list[PoolingParams] | None = None,

64 + num_scheduled_tokens: list[int] | None = None,

65 + seq_lens: list[int] | None = None,

66 + device: torch.device = _CPU,

67 +) -> PoolingMetadata:

68 + """Build a minimal PoolingMetadata for testing pooling methods."""

69 + n_seqs = len(prompt_lens)

70 + if tasks is None:

71 + tasks = ["embed"] * n_seqs

72 + if pooling_params is None:

73 + pooling_params = [PoolingParams(task=t) for t in tasks]

74 +

75 + prompt_lens_tensor = torch.tensor(prompt_lens, dtype=torch.long)

76 +

77 + prompt_token_ids_cpu = None

78 + prompt_token_ids = None

79 + if token_ids is not None:

80 + max_len = max(len(t) for t in token_ids)

81 + padded = [t + [0] * (max_len - len(t)) for t in token_ids]

82 + prompt_token_ids_cpu = torch.tensor(padded, dtype=torch.long)

83 + prompt_token_ids = prompt_token_ids_cpu.to(device)

84 +

85 + cursor = _make_pooling_cursor(

86 + prompt_lens,

87 + num_scheduled_tokens=num_scheduled_tokens,

88 + seq_lens=seq_lens,

89 + device=device,

90 + )

91 +

92 + pooling_states = [PoolingStates() for _ in range(n_seqs)]

93 +

94 + return PoolingMetadata(

95 + prompt_lens=prompt_lens_tensor,

96 + prompt_token_ids=prompt_token_ids,

97 + prompt_token_ids_cpu=prompt_token_ids_cpu,

98 + pooling_params=pooling_params,

99 + pooling_states=pooling_states,

100 + pooling_cursor=cursor,

101 + )

102 +

103 +

104 +# ---------------------------------------------------------------------------

105 +# CLSPool

106 +# ---------------------------------------------------------------------------

107 +class TestCLSPool:

108 + def test_extracts_first_token(self):

109 + hidden = torch.tensor(

110 + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]

111 + )

112 + metadata = _make_metadata([2, 3])

113 + pooler = CLSPool()

114 + out = pooler(hidden, metadata)

115 + expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])

116 + assert torch.equal(out, expected)

117 +

118 + def test_rejects_partial_prefill(self):

119 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

120 + metadata = _make_metadata([3], num_scheduled_tokens=[2])

121 + pooler = CLSPool()

122 + with pytest.raises(AssertionError, match="partial prefill"):

123 + pooler(hidden, metadata)

124 +

125 +

126 +# ---------------------------------------------------------------------------

127 +# LastPool

128 +# ---------------------------------------------------------------------------

129 +class TestLastPool:

130 + def test_extracts_last_token(self):

131 + hidden = torch.tensor(

132 + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]

133 + )

134 + metadata = _make_metadata([2, 3])

135 + pooler = LastPool()

136 + out = pooler(hidden, metadata)

137 + expected = torch.tensor([[3.0, 4.0], [9.0, 10.0]])

138 + assert torch.equal(out, expected)

139 +

140 + def test_partial_prefill_extracts_last_scheduled(self):

141 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

142 + metadata = _make_metadata([4], num_scheduled_tokens=[2])

143 + pooler = LastPool()

144 + out = pooler(hidden, metadata)

145 + expected = torch.tensor([[3.0, 4.0]])

146 + assert torch.equal(out, expected)

147 +

148 +

149 +# ---------------------------------------------------------------------------

150 +# MeanPool

151 +# ---------------------------------------------------------------------------

152 +class TestMeanPool:

153 + def test_computes_mean(self):

154 + hidden = torch.tensor(

155 + [[1.0, 2.0], [3.0, 4.0], [10.0, 20.0]], dtype=torch.float32

156 + )

157 + metadata = _make_metadata([2, 1])

158 + pooler = MeanPool()

159 + out = pooler(hidden, metadata)

160 + expected = torch.tensor([[2.0, 3.0], [10.0, 20.0]], dtype=torch.float32)

161 + assert torch.allclose(out, expected, atol=1e-5)

162 +

163 + def test_single_token_is_identity(self):

164 + hidden = torch.tensor([[5.0, 10.0]], dtype=torch.float32)

165 + metadata = _make_metadata([1])

166 + pooler = MeanPool()

167 + out = pooler(hidden, metadata)

168 + assert torch.allclose(out, hidden, atol=1e-5)

169 +

170 + def test_uniform_values_return_same(self):

171 + hidden = torch.full((4, 3), 7.0, dtype=torch.float32)

172 + metadata = _make_metadata([4])

173 + pooler = MeanPool()

174 + out = pooler(hidden, metadata)

175 + expected = torch.full((1, 3), 7.0, dtype=torch.float32)

176 + assert torch.allclose(out, expected, atol=1e-5)

177 +

178 + def test_multiple_sequences(self):

179 + hidden = torch.tensor(

180 + [

181 + [0.0, 0.0],

182 + [2.0, 4.0],

183 + [4.0, 8.0],

184 + [10.0, 10.0],

185 + ],

186 + dtype=torch.float32,

187 + )

188 + metadata = _make_metadata([3, 1])

189 + pooler = MeanPool()

190 + out = pooler(hidden, metadata)

191 + expected = torch.tensor([[2.0, 4.0], [10.0, 10.0]], dtype=torch.float32)

192 + assert torch.allclose(out, expected, atol=1e-5)

193 +

194 + def test_empty_batch(self):

195 + hidden = torch.empty((0, 8), dtype=torch.float32)

196 + metadata = _make_metadata([])

197 + pooler = MeanPool()

198 + out = pooler(hidden, metadata)

199 + assert out.shape == (0, 8)

200 +

201 + def test_rejects_partial_prefill(self):

202 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float32)

203 + metadata = _make_metadata([3], num_scheduled_tokens=[2])

204 + pooler = MeanPool()

205 + with pytest.raises(AssertionError, match="partial prefill"):

206 + pooler(hidden, metadata)

207 +

208 + def test_chunked_accumulation(self):

209 + hidden = torch.arange(20, dtype=torch.float32).reshape(5, 4)

210 + metadata = _make_metadata([3, 2])

211 + pooler = MeanPool()

212 + with patch(

213 + "vllm.model_executor.layers.pooler.seqwise.methods"

214 + "._MEAN_POOL_ACCUMULATION_CHUNK_BYTES",

215 + 16,

216 + ):

217 + out = pooler(hidden, metadata)

218 + expected_seq0 = hidden[:3].float().mean(dim=0, keepdim=True)

219 + expected_seq1 = hidden[3:].float().mean(dim=0, keepdim=True)

220 + expected = torch.cat([expected_seq0, expected_seq1], dim=0)

221 + assert torch.allclose(out, expected, atol=1e-5)

222 +

223 + def test_upcasts_to_float32(self):

224 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]], dtype=torch.float16)

225 + metadata = _make_metadata([2])

226 + pooler = MeanPool()

227 + out = pooler(hidden, metadata)

228 + assert out.dtype == torch.float32

229 + expected = torch.tensor([[2.0, 3.0]], dtype=torch.float32)

230 + assert torch.allclose(out, expected, atol=1e-2)

231 +

232 +

233 +# ---------------------------------------------------------------------------

234 +# get_seq_pooling_method factory

235 +# ---------------------------------------------------------------------------

236 +class TestGetSeqPoolingMethod:

237 + def test_cls(self):

238 + assert isinstance(get_seq_pooling_method("CLS"), CLSPool)

239 +

240 + def test_last(self):

241 + assert isinstance(get_seq_pooling_method("LAST"), LastPool)

242 +

243 + def test_mean(self):

244 + assert isinstance(get_seq_pooling_method("MEAN"), MeanPool)

245 +

246 + def test_unknown_raises(self):

247 + with pytest.raises(NotImplementedError, match="UNKNOWN"):

248 + get_seq_pooling_method("UNKNOWN")

249 +

250 +

251 +# ---------------------------------------------------------------------------

252 +# AllPool

253 +# ---------------------------------------------------------------------------

254 +

255 +

256 +@dataclass

257 +class _FakeSchedulerConfig:

258 + enable_chunked_prefill: bool = False

259 +

260 +

261 +@dataclass

262 +class _FakeVllmConfig:

263 + scheduler_config: _FakeSchedulerConfig

264 +

265 +

266 +class TestAllPool:

267 + @staticmethod

268 + def _make_all_pool(*, chunked: bool = False) -> AllPool:

269 + fake_config = _FakeVllmConfig(

270 + scheduler_config=_FakeSchedulerConfig(

271 + enable_chunked_prefill=chunked,

272 + ),

273 + )

274 + with patch(

275 + "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",

276 + return_value=fake_config,

277 + ):

278 + return AllPool()

279 +

280 + def test_splits_by_sequence(self):

281 + pooler = self._make_all_pool()

282 + hidden = torch.tensor(

283 + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]

284 + )

285 + metadata = _make_metadata([2, 3])

286 + out = pooler(hidden, metadata)

287 + assert len(out) == 2

288 + assert torch.equal(out[0], hidden[:2])

289 + assert torch.equal(out[1], hidden[2:])

290 +

291 + def test_single_sequence(self):

292 + pooler = self._make_all_pool()

293 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

294 + metadata = _make_metadata([3])

295 + out = pooler(hidden, metadata)

296 + assert len(out) == 1

297 + assert torch.equal(out[0], hidden)

298 +

299 + def test_chunked_prefill_returns_none_for_unfinished(self):

300 + pooler = self._make_all_pool(chunked=True)

301 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

302 + metadata = _make_metadata(

303 + [4],

304 + num_scheduled_tokens=[2],

305 + seq_lens=[2],

306 + )

307 + out = pooler(hidden, metadata)

308 + assert len(out) == 1

309 + assert out[0] is None

310 +

311 + def test_chunked_prefill_returns_concat_when_finished(self):

312 + pooler = self._make_all_pool(chunked=True)

313 +

314 + chunk1 = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

315 + metadata1 = _make_metadata(

316 + [4],

317 + num_scheduled_tokens=[2],

318 + seq_lens=[2],

319 + )

320 + out1 = pooler(chunk1, metadata1)

321 + assert out1[0] is None

322 +

323 + chunk2 = torch.tensor([[5.0, 6.0], [7.0, 8.0]])

324 + metadata2 = _make_metadata(

325 + [4],

326 + num_scheduled_tokens=[2],

327 + seq_lens=[4],

328 + )

329 + metadata2.pooling_states = metadata1.pooling_states

330 + out2 = pooler(chunk2, metadata2)

331 + assert out2[0] is not None

332 + expected = torch.cat([chunk1, chunk2], dim=0)

333 + assert torch.equal(out2[0], expected)

334 +

335 + def test_chunked_prefill_single_shot_matches_non_chunked(self):

336 + pooler = self._make_all_pool(chunked=True)

337 + hidden = torch.tensor(

338 + [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0], [9.0, 10.0]]

339 + )

340 + metadata = _make_metadata([2, 3])

341 + out = pooler(hidden, metadata)

342 + assert len(out) == 2

343 + assert torch.equal(out[0], hidden[:2])

344 + assert torch.equal(out[1], hidden[2:])

345 +

346 + def test_chunked_prefill_mixed_finished_unfinished(self):

347 + pooler = self._make_all_pool(chunked=True)

348 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])

349 + metadata = _make_metadata(

350 + [2, 4],

351 + num_scheduled_tokens=[2, 1],

352 + seq_lens=[2, 1],

353 + )

354 + out = pooler(hidden, metadata)

355 + assert len(out) == 2

356 + assert torch.equal(out[0], hidden[:2])

357 + assert out[1] is None

358 +

359 +

360 +# ---------------------------------------------------------------------------

361 +# StepPool

362 +# ---------------------------------------------------------------------------

363 +class TestStepPool:

364 + @staticmethod

365 + def _make_step_pool(*, chunked: bool = False) -> StepPool:

366 + fake_config = _FakeVllmConfig(

367 + scheduler_config=_FakeSchedulerConfig(

368 + enable_chunked_prefill=chunked,

369 + ),

370 + )

371 + with patch(

372 + "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",

373 + return_value=fake_config,

374 + ):

375 + return StepPool()

376 +

377 + def test_filters_by_step_tag_id(self):

378 + pooler = self._make_step_pool()

379 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])

380 + token_ids = [[10, 99, 10, 20]]

381 + params = [PoolingParams(task="token_classify", step_tag_id=10)]

382 + metadata = _make_metadata([4], token_ids=token_ids, pooling_params=params)

383 + out = pooler(hidden, metadata)

384 + assert len(out) == 1

385 + expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])

386 + assert torch.equal(out[0], expected)

387 +

388 + def test_filters_by_returned_token_ids(self):

389 + pooler = self._make_step_pool()

390 + hidden = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])

391 + token_ids = [[10, 20]]

392 + params = [PoolingParams(task="token_classify", returned_token_ids=[0, 2])]

393 + metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)

394 + out = pooler(hidden, metadata)

395 + assert len(out) == 1

396 + expected = torch.tensor([[1.0, 3.0], [4.0, 6.0]])

397 + assert torch.equal(out[0], expected)

398 +

399 + def test_no_filtering_without_params(self):

400 + pooler = self._make_step_pool()

401 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

402 + token_ids = [[10, 20]]

403 + params = [PoolingParams(task="token_classify")]

404 + metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)

405 + out = pooler(hidden, metadata)

406 + assert len(out) == 1

407 + assert torch.equal(out[0], hidden)

408 +

409 + def test_combined_step_tag_and_returned_token_ids(self):

410 + pooler = self._make_step_pool()

411 + hidden = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]])

412 + token_ids = [[99, 10, 99]]

413 + params = [

414 + PoolingParams(

415 + task="token_classify",

416 + step_tag_id=10,

417 + returned_token_ids=[0, 2],

418 + )

419 + ]

420 + metadata = _make_metadata([3], token_ids=token_ids, pooling_params=params)

421 + out = pooler(hidden, metadata)

422 + assert len(out) == 1

423 + expected = torch.tensor([[4.0, 6.0]])

424 + assert torch.equal(out[0], expected)

425 +

426 + def test_step_tag_id_no_match_returns_empty(self):

427 + pooler = self._make_step_pool()

428 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

429 + token_ids = [[10, 20]]

430 + params = [PoolingParams(task="token_classify", step_tag_id=999)]

431 + metadata = _make_metadata([2], token_ids=token_ids, pooling_params=params)

432 + out = pooler(hidden, metadata)

433 + assert len(out) == 1

434 + assert out[0].shape == (0, 2)

435 +

436 + def test_chunked_prefill_propagates_none_for_unfinished(self):

437 + pooler = self._make_step_pool(chunked=True)

438 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])

439 + token_ids = [[10, 20, 30, 40]]

440 + params = [PoolingParams(task="token_classify", step_tag_id=10)]

441 + metadata = _make_metadata(

442 + [4],

443 + token_ids=token_ids,

444 + pooling_params=params,

445 + num_scheduled_tokens=[2],

446 + seq_lens=[2],

447 + )

448 + out = pooler(hidden, metadata)

449 + assert len(out) == 1

450 + assert out[0] is None

451 +

452 + def test_chunked_prefill_filters_when_finished(self):

453 + pooler = self._make_step_pool(chunked=True)

454 + hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0], [7.0, 8.0]])

455 + token_ids = [[10, 99, 10, 20]]

456 + params = [PoolingParams(task="token_classify", step_tag_id=10)]

457 + metadata = _make_metadata([4], token_ids=token_ids, pooling_params=params)

458 + out = pooler(hidden, metadata)

459 + assert len(out) == 1

460 + expected = torch.tensor([[1.0, 2.0], [5.0, 6.0]])

461 + assert torch.equal(out[0], expected)

462 +

463 + def test_requires_token_ids_update(self):

464 + pooler = self._make_step_pool()

465 + update = pooler.get_pooling_updates("token_classify")

466 + assert update.requires_token_ids is True

467 +

468 +

469 +# ---------------------------------------------------------------------------

470 +# get_tok_pooling_method factory

471 +# ---------------------------------------------------------------------------

472 +class TestGetTokPoolingMethod:

473 + def test_all(self):

474 + fake_config = _FakeVllmConfig(

475 + scheduler_config=_FakeSchedulerConfig(

476 + enable_chunked_prefill=False,

477 + ),

478 + )

479 + with patch(

480 + "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",

481 + return_value=fake_config,

482 + ):

483 + assert isinstance(get_tok_pooling_method("ALL"), AllPool)

484 +

485 + def test_step(self):

486 + fake_config = _FakeVllmConfig(

487 + scheduler_config=_FakeSchedulerConfig(

488 + enable_chunked_prefill=False,

489 + ),

490 + )

491 + with patch(

492 + "vllm.model_executor.layers.pooler.tokwise.methods.get_current_vllm_config",

493 + return_value=fake_config,

494 + ):

495 + assert isinstance(get_tok_pooling_method("STEP"), StepPool)

496 +

497 + def test_unknown_raises(self):

498 + with pytest.raises(NotImplementedError, match="UNKNOWN"):

499 + get_tok_pooling_method("UNKNOWN")

[Misc] added unit tests for the core pooling methods (#43818)