[AMD] Guard aiter greedy_sample OOB token id (fixes VLM MMMU CI) (#27247)

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
                 )

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__":