diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index e3b8f4d79..bc79cca97 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/python/sglang/srt/layers/sampler.py @@ -46,6 +46,13 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip() if _use_aiter: from aiter import greedy_sample as _aiter_greedy_sample +# The aiter greedy_sample kernel can return an out-of-range token id (== vocab_size, +# e.g. 151666 for MiniCPM-V) for all-NaN / all -inf logit rows on ROCm, which decodes +# to an empty string and breaks downstream consumers. Set this to 1 to fall back to +# torch.argmax (which always returns a valid index). Default off so behavior is +# unchanged elsewhere. +_disable_aiter_greedy_sample = get_bool_env_var("SGLANG_DISABLE_AITER_GREEDY_SAMPLE") + if is_npu(): import torch_npu @@ -110,7 +117,7 @@ class Sampler(nn.Module): logits = self._preprocess_logits(logits, sampling_info) if sampling_info.is_all_greedy: - if _use_aiter: + if _use_aiter and not _disable_aiter_greedy_sample: batch_next_token_ids = torch.empty( logits.shape[0], device=logits.device, dtype=torch.int32 ) diff --git a/test/registered/models/test_vlm_models.py b/test/registered/models/test_vlm_models.py index e6dd05e2f..a0eea75ae 100644 --- a/test/registered/models/test_vlm_models.py +++ b/test/registered/models/test_vlm_models.py @@ -8,7 +8,7 @@ from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci from sglang.test.kits.mmmu_vlm_kit import ( MMMUMultiModelTestBase, ) -from sglang.test.test_utils import is_in_ci +from sglang.test.test_utils import is_in_amd_ci, is_in_ci # VLM (Vision Language Model) tests @@ -41,7 +41,14 @@ class TestVLMModels(MMMUMultiModelTestBase): with tempfile.TemporaryDirectory( prefix=f"test_vlm_mmmu_{model.model.replace('/', '_')}_" ) as temp_dir: - self._run_vlm_mmmu_test(model, temp_dir) + # On AMD CI, the aiter greedy_sample kernel returns an out-of-range + # token id (== vocab_size) for degenerate (all-NaN / all -inf) logit + # rows, producing empty completions that crash the MMMU eval. Disable + # it there so greedy sampling falls back to torch.argmax. + custom_env = None + if is_in_amd_ci(): + custom_env = {"SGLANG_DISABLE_AITER_GREEDY_SAMPLE": "1"} + self._run_vlm_mmmu_test(model, temp_dir, custom_env=custom_env) if __name__ == "__main__":