"""Unit tests for aiter greedy_sample kernel and Sampler integration. Validates that: 1. aiter.greedy_sample produces identical results to torch.argmax (kernel level) 2. Sampler.forward() correctly dispatches to aiter when _use_aiter=True 3. The fallback to torch.argmax works when _use_aiter=False 4. return_logprob path works with the aiter greedy branch The kernel is designed for production LLM inference (large vocab, bf16) and is used when SGLANG_USE_AITER=1 on ROCm. """ import unittest from unittest import mock import torch from sglang.srt.utils.common import is_hip from sglang.test.ci.ci_register import register_amd_ci register_amd_ci(est_time=60, suite="stage-b-test-1-gpu-small-amd") def _mock_global_server_args(backend="pytorch"): from sglang.srt.layers import sampler as sampler_mod from sglang.srt.server_args import ( ServerArgs, set_global_server_args_for_scheduler, ) # Publish for real: the sampler reads the context slot through # get_server_args(), which a module-attribute rebinding cannot intercept. set_global_server_args_for_scheduler( ServerArgs(model_path="dummy", sampling_backend=backend) ) class _DummyTPGroup: device_group = None sampler_mod.get_tp_group = lambda: _DummyTPGroup() from sglang.srt.runtime_context import get_flags get_flags().dp.enabled = False def _make_sampling_info(batch_size, vocab_size, device="cuda"): from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo return SamplingBatchInfo( temperatures=torch.ones(batch_size, 1, device=device, dtype=torch.float), top_ps=torch.ones(batch_size, device=device), top_ks=torch.zeros(batch_size, device=device, dtype=torch.int32), min_ps=torch.zeros(batch_size, device=device), is_all_greedy=True, is_any_greedy=True, need_top_p_sampling=False, need_top_k_sampling=False, need_min_p_sampling=False, vocab_size=vocab_size, device=device, ) @unittest.skipUnless(is_hip(), "aiter greedy_sample requires ROCm") class TestAiterGreedySample(unittest.TestCase): """Kernel-level correctness: aiter.greedy_sample vs torch.argmax.""" @classmethod def setUpClass(cls): try: from aiter import greedy_sample cls.greedy_sample = staticmethod(greedy_sample) except ImportError: raise unittest.SkipTest("aiter not installed") cls.device = "cuda" def setUp(self): torch.manual_seed(42) torch.cuda.manual_seed_all(42) def _run_and_compare(self, batch_size, vocab_size): logits = torch.randn( batch_size, vocab_size, device=self.device, dtype=torch.bfloat16 ) expected = torch.argmax(logits, dim=-1) actual = torch.empty(logits.shape[0], device=logits.device, dtype=torch.int32) self.greedy_sample(actual, logits) self.assertTrue( torch.equal(actual.to(expected.dtype), expected), f"Mismatch for shape ({batch_size}, {vocab_size}): " f"expected={expected[:8].tolist()}, got={actual[:8].tolist()}", ) def test_single_request(self): self._run_and_compare(1, 32000) def test_small_batch(self): self._run_and_compare(4, 32000) def test_medium_batch(self): self._run_and_compare(32, 32000) def test_large_batch(self): self._run_and_compare(128, 32000) def test_realistic_vocab_deepseek(self): self._run_and_compare(64, 129280) def test_realistic_vocab_llama3(self): self._run_and_compare(64, 128256) def test_realistic_vocab_qwen(self): self._run_and_compare(64, 151936) def test_various_batch_sizes(self): configs = [ (1, 128256), (2, 128256), (8, 128256), (16, 128256), (32, 129280), (64, 129280), (128, 129280), (256, 129280), ] for batch_size, vocab_size in configs: with self.subTest(batch_size=batch_size, vocab_size=vocab_size): self._run_and_compare(batch_size, vocab_size) def test_tied_values(self): vocab_size = 32000 logits = torch.zeros(8, vocab_size, device=self.device, dtype=torch.bfloat16) logits[:, 0] = 1.0 expected = torch.argmax(logits, dim=-1) actual = torch.empty(8, device=self.device, dtype=torch.int32) self.greedy_sample(actual, logits) self.assertTrue(torch.equal(actual.to(expected.dtype), expected)) def test_negative_logits(self): vocab_size = 32000 logits = ( torch.randn(16, vocab_size, device=self.device, dtype=torch.bfloat16) - 5.0 ) expected = torch.argmax(logits, dim=-1) actual = torch.empty(16, device=self.device, dtype=torch.int32) self.greedy_sample(actual, logits) self.assertTrue(torch.equal(actual.to(expected.dtype), expected)) def test_extreme_values(self): vocab_size = 32000 logits = torch.randn(16, vocab_size, device=self.device, dtype=torch.bfloat16) logits[0, 42] = 1e4 logits[1, 100] = -1e4 expected = torch.argmax(logits, dim=-1) actual = torch.empty(16, device=self.device, dtype=torch.int32) self.greedy_sample(actual, logits) self.assertTrue(torch.equal(actual.to(expected.dtype), expected)) @unittest.skipUnless(is_hip(), "aiter greedy_sample requires ROCm") class TestAiterGreedyIntegration(unittest.TestCase): """Integration: Sampler.forward() with _use_aiter on/off.""" @classmethod def setUpClass(cls): try: from aiter import greedy_sample cls._greedy_sample_fn = staticmethod(greedy_sample) except ImportError: raise unittest.SkipTest("aiter not installed") cls.device = "cuda" def setUp(self): torch.manual_seed(42) torch.cuda.manual_seed_all(42) def _run_sampler(self, use_aiter, logits, sampling_info, return_logprob=False): from sglang.srt.layers import sampler as sampler_mod from sglang.srt.layers.logits_processor import LogitsProcessorOutput _mock_global_server_args() patches = {"_use_aiter": use_aiter} if use_aiter: patches["_aiter_greedy_sample"] = self._greedy_sample_fn with mock.patch.multiple(sampler_mod, **patches): sampler = sampler_mod.Sampler() batch_size = logits.shape[0] positions = torch.arange(batch_size, device=self.device, dtype=torch.int32) return sampler.forward( logits_output=LogitsProcessorOutput(next_token_logits=logits.clone()), sampling_info=sampling_info, return_logprob=return_logprob, top_logprobs_nums=[0] * batch_size, token_ids_logprobs=[None] * batch_size, positions=positions, ) def test_aiter_matches_argmax_through_sampler(self): batch_size, vocab_size = 64, 129280 logits = torch.randn( batch_size, vocab_size, device=self.device, dtype=torch.bfloat16 ) sampling_info = _make_sampling_info(batch_size, vocab_size, self.device) out_aiter = self._run_sampler(True, logits, sampling_info) out_argmax = self._run_sampler(False, logits, sampling_info) self.assertTrue( torch.equal(out_aiter.cpu(), out_argmax.cpu()), f"Sampler mismatch: aiter={out_aiter[:8].tolist()}, " f"argmax={out_argmax[:8].tolist()}", ) def test_fallback_to_argmax_when_disabled(self): batch_size, vocab_size = 32, 32000 logits = torch.randn( batch_size, vocab_size, device=self.device, dtype=torch.bfloat16 ) sampling_info = _make_sampling_info(batch_size, vocab_size, self.device) out = self._run_sampler(False, logits, sampling_info) expected = torch.argmax(logits, dim=-1) self.assertTrue( torch.equal(out.cpu(), expected.cpu()), "Fallback path should produce torch.argmax results", ) def test_aiter_greedy_with_return_logprob(self): batch_size, vocab_size = 16, 32000 logits = torch.randn( batch_size, vocab_size, device=self.device, dtype=torch.bfloat16 ) sampling_info = _make_sampling_info(batch_size, vocab_size, self.device) out_aiter = self._run_sampler(True, logits, sampling_info, return_logprob=True) out_argmax = self._run_sampler( False, logits, sampling_info, return_logprob=True ) self.assertTrue( torch.equal(out_aiter.cpu(), out_argmax.cpu()), "Token IDs should match with return_logprob=True", ) def test_aiter_output_dtype(self): """Document that the aiter path returns int32 (vs int64 from argmax).""" batch_size, vocab_size = 16, 32000 logits = torch.randn( batch_size, vocab_size, device=self.device, dtype=torch.bfloat16 ) sampling_info = _make_sampling_info(batch_size, vocab_size, self.device) out_aiter = self._run_sampler(True, logits, sampling_info) out_argmax = self._run_sampler(False, logits, sampling_info) self.assertEqual(out_aiter.dtype, torch.int32) self.assertEqual(out_argmax.dtype, torch.int64) def test_various_batch_sizes_through_sampler(self): vocab_size = 129280 for batch_size in [1, 4, 16, 64, 128]: with self.subTest(batch_size=batch_size): logits = torch.randn( batch_size, vocab_size, device=self.device, dtype=torch.bfloat16, ) sampling_info = _make_sampling_info(batch_size, vocab_size, self.device) out_aiter = self._run_sampler(True, logits, sampling_info) out_argmax = self._run_sampler(False, logits, sampling_info) self.assertTrue( torch.equal(out_aiter.cpu(), out_argmax.cpu()), f"Mismatch at batch_size={batch_size}", ) if __name__ == "__main__": unittest.main()