"""Unit tests for token embedding override support. Covers: - PositionalEmbeds dataclass (embed_types.py) - convert_embeds_to_tensors (utils.py) - TokenizerManager._resolve_embed_overrides (tokenizer_manager.py) - positional_embed_overrides on GenerateReqInput/EmbeddingReqInput (io_struct.py) - Score mixin override resolution (tokenizer_manager_score_mixin.py) """ import unittest from unittest.mock import AsyncMock, MagicMock import torch from sglang.srt.entrypoints.openai.utils import convert_embeds_to_tensors from sglang.srt.managers.embed_types import PositionalEmbeds from sglang.srt.managers.io_struct import EmbeddingReqInput, GenerateReqInput from sglang.srt.managers.tokenizer_manager import TokenizerManager from sglang.srt.managers.tokenizer_manager_score_mixin import ( TokenizerManagerScoreMixin, ) from sglang.srt.server_args import MIS_DELIMITER_TOKEN_ID from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.test_utils import CustomTestCase register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-small") register_amd_ci(est_time=9, suite="stage-b-test-1-gpu-small-amd") HIDDEN_DIM = 4 def _vec(val: float = 1.0) -> torch.Tensor: """Create a 1-D tensor of size HIDDEN_DIM.""" return torch.full((HIDDEN_DIM,), val, dtype=torch.float32) def _vec2d(val: float = 1.0) -> torch.Tensor: """Create a [1, HIDDEN_DIM] tensor.""" return torch.full((1, HIDDEN_DIM), val, dtype=torch.float32) # ======================================================================== # PositionalEmbeds # ======================================================================== class TestPositionalEmbeds(CustomTestCase): def test_from_list_of_1d_tensors(self): pe = PositionalEmbeds(embeds=[_vec(1), _vec(2)], positions=[0, 5]) self.assertEqual(pe.embeds.shape, (2, HIDDEN_DIM)) self.assertAlmostEqual(pe.embeds[0, 0].item(), 1.0) self.assertAlmostEqual(pe.embeds[1, 0].item(), 2.0) def test_from_list_of_2d_tensors(self): pe = PositionalEmbeds(embeds=[_vec2d(3), _vec2d(4)], positions=[1, 2]) self.assertEqual(pe.embeds.shape, (2, HIDDEN_DIM)) def test_from_pre_stacked_tensor(self): stacked = torch.zeros(3, HIDDEN_DIM) pe = PositionalEmbeds(embeds=stacked, positions=[0, 1, 2]) self.assertIs(pe.embeds, stacked) def test_length_mismatch_raises(self): with self.assertRaises(ValueError): PositionalEmbeds(embeds=[_vec()], positions=[0, 1]) def test_empty(self): pe = PositionalEmbeds(embeds=torch.zeros(0, HIDDEN_DIM), positions=[]) self.assertEqual(pe.embeds.shape[0], 0) # ======================================================================== # convert_embeds_to_tensors # ======================================================================== class TestConvertEmbedsToTensors(CustomTestCase): def test_none_returns_none(self): self.assertIsNone(convert_embeds_to_tensors(None)) def test_empty_list(self): self.assertEqual(convert_embeds_to_tensors([]), []) def test_single_input(self): """[num_replacements][hidden_size] -> [[tensor, ...]]""" result = convert_embeds_to_tensors([[1.0, 2.0], [3.0, 4.0]]) self.assertEqual(len(result), 1) # wrapped in outer list self.assertEqual(len(result[0]), 2) # two replacement vectors self.assertTrue(torch.is_tensor(result[0][0])) self.assertEqual(result[0][0].tolist(), [1.0, 2.0]) self.assertEqual(result[0][0].dtype, torch.float32) self.assertEqual(result[0][0].dim(), 1) # each vector is 1-D def test_batch_input(self): """[num_inputs][num_replacements][hidden_size] -> [[tensor, ...], ...]""" result = convert_embeds_to_tensors( [ [[1.0, 2.0]], [[3.0, 4.0], [5.0, 6.0]], ] ) self.assertEqual(len(result), 2) self.assertEqual(len(result[0]), 1) self.assertEqual(len(result[1]), 2) # ======================================================================== # TokenizerManager._resolve_embed_overrides # ======================================================================== class TestResolveEmbedOverrides(CustomTestCase): def test_basic_resolution(self): embeds = [_vec(1), _vec(2)] pe = TokenizerManager._resolve_embed_overrides( input_ids=[10, 50, 20, 50, 30], token_id=50, embeds=embeds, ) self.assertIsInstance(pe, PositionalEmbeds) self.assertEqual(pe.positions, [1, 3]) self.assertEqual(pe.embeds.shape, (2, HIDDEN_DIM)) def test_no_placeholders_raises(self): with self.assertRaises(ValueError): TokenizerManager._resolve_embed_overrides( input_ids=[10, 20, 30], token_id=50, embeds=[_vec()], ) def test_count_mismatch_raises(self): with self.assertRaises(ValueError): TokenizerManager._resolve_embed_overrides( input_ids=[10, 50, 20], token_id=50, embeds=[_vec(1), _vec(2)], ) # ======================================================================== # io_struct: positional_embed_overrides on GenerateReqInput # ======================================================================== class TestGenerateReqInputEmbedOverride(CustomTestCase): def test_single_override_in_getitem(self): """Single PositionalEmbeds is shared across all items in __getitem__.""" pe = PositionalEmbeds(embeds=[_vec()], positions=[0]) req = GenerateReqInput( input_ids=[[1, 2], [3, 4]], sampling_params=[{}, {}], positional_embed_overrides=pe, ) req.normalize_batch_and_arguments() item = req[0] self.assertIs(item.positional_embed_overrides, pe) def test_batch_override_in_getitem(self): """List[Optional[PositionalEmbeds]] is indexed per-item.""" pe0 = PositionalEmbeds(embeds=[_vec(1)], positions=[0]) pe1 = None req = GenerateReqInput( input_ids=[[1, 2], [3, 4]], sampling_params=[{}, {}], positional_embed_overrides=[pe0, pe1], ) req.normalize_batch_and_arguments() self.assertEqual(req[0].positional_embed_overrides, pe0) self.assertIsNone(req[1].positional_embed_overrides) # ======================================================================== # io_struct: embed override fields on EmbeddingReqInput # ======================================================================== class TestEmbeddingReqInputEmbedOverride(CustomTestCase): def test_override_fields_in_getitem(self): """embed_override_token_id, embed_overrides, and positional_embed_overrides are correctly sliced in __getitem__.""" pe0 = PositionalEmbeds(embeds=[_vec(1)], positions=[0]) pe1 = PositionalEmbeds(embeds=[_vec(2)], positions=[1]) req = EmbeddingReqInput( input_ids=[[50, 10], [20, 50]], sampling_params=[{}, {}], embed_override_token_id=50, embed_overrides=[[_vec(1)], [_vec(2)]], positional_embed_overrides=[pe0, pe1], ) req.normalize_batch_and_arguments() item0 = req[0] item1 = req[1] self.assertEqual(item0.embed_override_token_id, 50) self.assertEqual(len(item0.embed_overrides), 1) self.assertEqual(item0.positional_embed_overrides, pe0) self.assertEqual(item1.positional_embed_overrides, pe1) # ======================================================================== # Score mixin: _resolve_overrides_for_sequence # ======================================================================== class _FakeServerArgs: """Minimal stub for server_args.""" def __init__(self, enable_mis=False): self.enable_mis = enable_mis class _FakeMixin(TokenizerManagerScoreMixin): """Minimal stub to call mixin methods without a full TokenizerManager.""" def __init__(self, enable_mis=False): self.server_args = _FakeServerArgs(enable_mis) self.tokenizer = None self.is_generation = True class TestResolveOverridesForSequence(CustomTestCase): def setUp(self): self.mixin = _FakeMixin() def test_none_embeds_returns_empty(self): embeds, positions = self.mixin._resolve_overrides_for_sequence( token_ids=[10, 50, 20], embeds=None, embed_override_token_id=50, ) self.assertEqual(embeds, []) self.assertEqual(positions, []) def test_basic_resolution(self): e1, e2 = _vec(1), _vec(2) embeds, positions = self.mixin._resolve_overrides_for_sequence( token_ids=[50, 10, 50], embeds=[e1, e2], embed_override_token_id=50, ) self.assertEqual(len(embeds), 2) self.assertEqual(positions, [0, 2]) def test_with_offset(self): embeds, positions = self.mixin._resolve_overrides_for_sequence( token_ids=[10, 50], embeds=[_vec()], embed_override_token_id=50, position_offset=100, ) self.assertEqual(positions, [101]) def test_empty_embeds_list(self): """Empty embeds list with no placeholders succeeds.""" embeds, positions = self.mixin._resolve_overrides_for_sequence( token_ids=[10, 20], embeds=[], embed_override_token_id=50, ) self.assertEqual(embeds, []) self.assertEqual(positions, []) def test_count_mismatch_raises(self): with self.assertRaises(ValueError): self.mixin._resolve_overrides_for_sequence( token_ids=[50, 50], embeds=[_vec()], embed_override_token_id=50, ) # ======================================================================== # Score mixin: _resolve_embed_overrides_for_request # ======================================================================== class TestResolveEmbedOverridesForRequest(CustomTestCase): def setUp(self): self.mixin = _FakeMixin() def test_no_overrides_returns_none(self): result = self.mixin._resolve_embed_overrides_for_request( query=[10, 20], item=[30, 40], embed_override_token_id=50, query_embed_overrides=None, item_embeds=None, item_position_offset=2, item_label="items[0]", ) self.assertIsNone(result) def test_query_only_overrides(self): pe = self.mixin._resolve_embed_overrides_for_request( query=[50, 20], item=[30, 40], embed_override_token_id=50, query_embed_overrides=[_vec(1)], item_embeds=None, item_position_offset=2, item_label="items[0]", ) self.assertIsInstance(pe, PositionalEmbeds) self.assertEqual(pe.positions, [0]) self.assertEqual(pe.embeds.shape, (1, HIDDEN_DIM)) def test_item_only_overrides(self): pe = self.mixin._resolve_embed_overrides_for_request( query=[10, 20], item=[50, 40], embed_override_token_id=50, query_embed_overrides=None, item_embeds=[_vec(2)], item_position_offset=2, item_label="items[0]", ) self.assertEqual(pe.positions, [2]) # offset applied def test_query_and_item_overrides(self): pe = self.mixin._resolve_embed_overrides_for_request( query=[50, 20], item=[30, 50], embed_override_token_id=50, query_embed_overrides=[_vec(1)], item_embeds=[_vec(2)], item_position_offset=2, item_label="items[0]", ) self.assertEqual(pe.positions, [0, 3]) # query pos 0, item pos 1+offset 2 self.assertEqual(pe.embeds.shape, (2, HIDDEN_DIM)) # ======================================================================== # Score mixin: _build_token_id_inputs # ======================================================================== DELIM_TOKEN = MIS_DELIMITER_TOKEN_ID class TestBuildTokenIdInputs(CustomTestCase): def setUp(self): self.mixin = _FakeMixin(enable_mis=True) # --- single-item mode, no embeds --- def test_single_item_no_embeds(self): _, input_ids, positional_embed_overrides, _ = self.mixin._build_token_id_inputs( query=[1, 2], items=[[3, 4], [5, 6]], item_first=False, use_multi_item_scoring=False, embed_override_token_id=None, query_embed_overrides=None, item_embed_overrides=None, ) self.assertEqual(input_ids, [[1, 2, 3, 4], [1, 2, 5, 6]]) self.assertIsNone(positional_embed_overrides) def test_single_item_no_embeds_item_first(self): _, input_ids, positional_embed_overrides, _ = self.mixin._build_token_id_inputs( query=[1, 2], items=[[3, 4]], item_first=True, use_multi_item_scoring=False, embed_override_token_id=None, query_embed_overrides=None, item_embed_overrides=None, ) self.assertEqual(input_ids, [[3, 4, 1, 2]]) self.assertIsNone(positional_embed_overrides) # --- multi-item mode, no embeds --- def test_multi_item_no_embeds(self): _, input_ids, positional_embed_overrides, _ = self.mixin._build_token_id_inputs( query=[1, 2], items=[[3, 4], [5, 6]], item_first=False, use_multi_item_scoring=True, embed_override_token_id=None, query_embed_overrides=None, item_embed_overrides=None, ) # queryitem1item2 self.assertEqual( input_ids, [[1, 2, DELIM_TOKEN, 3, 4, DELIM_TOKEN, 5, 6, DELIM_TOKEN]] ) self.assertIsNone(positional_embed_overrides) # --- single-item mode, with embeds --- def test_single_item_query_embeds(self): """Query placeholder overrides are resolved per item.""" _, input_ids, positional_embed_overrides, _ = self.mixin._build_token_id_inputs( query=[50, 10], items=[[20, 30], [40, 50]], item_first=False, use_multi_item_scoring=False, embed_override_token_id=50, query_embed_overrides=[_vec(1)], item_embed_overrides=None, ) self.assertEqual(input_ids, [[50, 10, 20, 30], [50, 10, 40, 50]]) self.assertIsNotNone(positional_embed_overrides) self.assertEqual(len(positional_embed_overrides), 2) # Each item gets its own PositionalEmbeds with query override at pos 0 self.assertEqual(positional_embed_overrides[0].positions, [0]) self.assertEqual(positional_embed_overrides[1].positions, [0]) def test_single_item_item_embeds(self): """Per-item overrides with correct position offsets.""" _, input_ids, positional_embed_overrides, _ = self.mixin._build_token_id_inputs( query=[10, 20], items=[[50, 30]], item_first=False, use_multi_item_scoring=False, embed_override_token_id=50, query_embed_overrides=None, item_embed_overrides=[[_vec(2)]], ) self.assertEqual(input_ids, [[10, 20, 50, 30]]) self.assertIsNotNone(positional_embed_overrides) # item placeholder at index 0 of item, offset by query length 2 self.assertEqual(positional_embed_overrides[0].positions, [2]) def test_single_item_no_override_positions_returns_none_injection(self): """When no items have placeholders, positional_embed_overrides should be None.""" _, input_ids, positional_embed_overrides, _ = self.mixin._build_token_id_inputs( query=[10, 20], items=[[30, 40]], item_first=False, use_multi_item_scoring=False, embed_override_token_id=50, query_embed_overrides=None, item_embed_overrides=[None], ) self.assertIsNone(positional_embed_overrides) def test_single_item_query_and_item_embeds(self): """Single-item mode with both query and item overrides in one request.""" _, input_ids, positional_embed_overrides, _ = self.mixin._build_token_id_inputs( query=[50, 10], items=[[20, 50]], item_first=False, use_multi_item_scoring=False, embed_override_token_id=50, query_embed_overrides=[_vec(1)], item_embed_overrides=[[_vec(2)]], ) self.assertEqual(input_ids, [[50, 10, 20, 50]]) self.assertIsNotNone(positional_embed_overrides) pe = positional_embed_overrides[0] # query override at pos 0, item override at pos 3 (query_len=2 + idx=1) self.assertEqual(pe.positions, [0, 3]) self.assertEqual(pe.embeds.shape, (2, HIDDEN_DIM)) def test_single_item_empty_query(self): """Empty query with item-only overrides (valid from score_prompts).""" _, input_ids, positional_embed_overrides, _ = self.mixin._build_token_id_inputs( query=[], items=[[50, 10]], item_first=False, use_multi_item_scoring=False, embed_override_token_id=50, query_embed_overrides=None, item_embed_overrides=[[_vec(1)]], ) self.assertEqual(input_ids, [[50, 10]]) self.assertIsNotNone(positional_embed_overrides) # item placeholder at absolute pos 0 (offset=len([])=0) self.assertEqual(positional_embed_overrides[0].positions, [0]) # --- multi-item mode, with embeds --- def test_multi_item_with_query_and_item_embeds(self): """Multi-item mode resolves query overrides once and item overrides per item.""" _, input_ids, positional_embed_overrides, _ = self.mixin._build_token_id_inputs( query=[50, 10], items=[[20, 50], [30, 40]], item_first=False, use_multi_item_scoring=True, embed_override_token_id=50, query_embed_overrides=[_vec(1)], item_embed_overrides=[[_vec(2)], None], ) # queryitem1item2 = [50,10, DELIM, 20,50, DELIM, 30,40, DELIM] self.assertEqual(len(input_ids), 1) self.assertIsNotNone(positional_embed_overrides) self.assertEqual( len(positional_embed_overrides), 1 ) # single PositionalEmbeds for combined sequence pe = positional_embed_overrides[0] # query override at pos 0, item[0] override at pos 4 (query_len=2 + delim=1 + idx=1) self.assertIn(0, pe.positions) self.assertIn(4, pe.positions) self.assertEqual(pe.embeds.shape[0], 2) # ======================================================================== # Score mixin: score_request validation # ======================================================================== class TestScoreRequestValidation(CustomTestCase): """Test validation guards in score_request without running full pipeline.""" def setUp(self): self.mixin = _FakeMixin() def _call(self, **kwargs): """Wrapper to call score_request synchronously.""" import asyncio return asyncio.run(self.mixin.score_request(**kwargs)) def test_generation_requires_label_token_ids(self): self.mixin.is_generation = True with self.assertRaisesRegex(ValueError, "label_token_ids is required"): self._call( query=[1, 2], items=[[3, 4]], label_token_ids=None, ) def test_seq_classification_allows_none_label_token_ids(self): """SequenceClassification models should not require label_token_ids. Verify it passes validation and reaches generate_request.""" self.mixin.is_generation = False mock_result = AsyncMock() mock_result.__anext__ = AsyncMock( return_value=[{"embedding": [0.1, 0.9], "meta_info": {"prompt_tokens": 2}}] ) self.mixin.generate_request = MagicMock(return_value=mock_result) result = self._call( query=[1, 2], items=[[3, 4]], label_token_ids=None, ) self.mixin.generate_request.assert_called_once() self.assertEqual(len(result.scores), 1) def test_items_none_raises(self): with self.assertRaisesRegex(ValueError, "items must be provided"): self._call( query=[1, 2], items=None, label_token_ids=[100], ) def test_empty_items_returns_empty(self): result = self._call( query=[1, 2], items=[], label_token_ids=[100], ) self.assertEqual(result.scores, []) self.assertEqual(result.prompt_tokens, 0) def test_embed_override_token_id_required_with_query_embeds(self): with self.assertRaisesRegex(ValueError, "embed_override_token_id is required"): self._call( query=[1, 2], items=[[3, 4]], label_token_ids=[100], query_embed_overrides=[_vec(1)], embed_override_token_id=None, ) def test_embed_override_token_id_required_with_item_embeds(self): with self.assertRaisesRegex(ValueError, "embed_override_token_id is required"): self._call( query=[1, 2], items=[[3, 4]], label_token_ids=[100], item_embed_overrides=[[_vec(1)]], embed_override_token_id=None, ) def test_item_first_with_embeds_raises(self): with self.assertRaisesRegex(ValueError, "item_first is not supported"): self._call( query=[1, 2], items=[[3, 4]], label_token_ids=[100], item_first=True, embed_override_token_id=50, query_embed_overrides=[_vec(1)], ) def test_item_embed_overrides_length_mismatch_raises(self): with self.assertRaisesRegex(ValueError, "must match items length"): self._call( query=[1, 2], items=[[3, 4], [5, 6]], label_token_ids=[100], embed_override_token_id=50, item_embed_overrides=[[_vec(1)]], # 1 override for 2 items ) if __name__ == "__main__": unittest.main()