[Feature] Support compile mode for batch invariance on SM80 (#42456)

Signed-off-by: yewentao256 <[email protected]>

.buildkite/test_areas/misc.yaml

@@ -305,6 +305,20 @@ steps:
   - pytest -v -s transformers_utils
   - pytest -v -s config
 
+- label: Batch Invariance (A100)
+  key: batch-invariance-a100
+  timeout_in_minutes: 30
+  device: a100
+  source_file_dependencies:
+    - vllm/v1/attention
+    - vllm/model_executor/layers
+    - tests/v1/determinism/
+  commands:
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pip install pytest-timeout pytest-forked
+    - pytest -v -s v1/determinism/test_batch_invariance.py
+    - VLLM_TEST_MODEL=deepseek-ai/DeepSeek-V2-Lite-Chat pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle[TRITON_MLA]
+
 - label: Batch Invariance (H100)
   key: batch-invariance-h100
   timeout_in_minutes: 30

tests/v1/determinism/test_batch_invariance.py

@@ -11,15 +11,12 @@ from utils import (
     TEST_MODEL,
     _extract_step_logprobs,
     _random_prompt,
-    is_device_capability_below_90,
     skip_unsupported,
 )
 
 import vllm.envs as envs
 from vllm import LLM, SamplingParams
 
-IS_DEVICE_CAPABILITY_BELOW_90 = is_device_capability_below_90()
-
 
 @skip_unsupported
 @pytest.mark.timeout(1000)
@@ -175,7 +172,6 @@ def test_logprobs_bitwise_batch_invariance_bs1_vs_bsN(
         max_model_len=8192,
         dtype="auto",  # not everything is supported
         gpu_memory_utilization=0.9,
-        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
         attention_config={"backend": backend},
     )
 
@@ -388,7 +384,6 @@ def test_simple_generation(backend):
         max_model_len=2048,
         dtype="auto",
         enable_prefix_caching=False,
-        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
         attention_config={"backend": backend},
     )
 
@@ -453,7 +448,6 @@ def test_logprobs_without_batch_invariance_should_fail(
         max_num_seqs=32,
         max_model_len=8192,
         dtype="auto",
-        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
         attention_config={"backend": backend},
     )
 
@@ -673,7 +667,6 @@ def test_decode_logprobs_match_prefill_logprobs(
         max_num_seqs=32,
         max_model_len=8192,
         dtype="auto",
-        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
         attention_config={"backend": backend},
     )
 
@@ -920,7 +913,6 @@ def LLM_with_max_seqs(
         dtype="auto",
         tensor_parallel_size=int(os.getenv("VLLM_TP_SIZE", "1")),
         enable_prefix_caching=False,
-        enforce_eager=IS_DEVICE_CAPABILITY_BELOW_90,
         attention_config=attention_config,
         # Enable for MOE models
         # enable_expert_parallel=True,