From 0da66a6e468653cd8294d5fe16831bb3d7367284 Mon Sep 17 00:00:00 2001
From: "Kevin H. Luu" <khluu000@gmail.com>
Date: Tue, 13 Oct 2026 16:14:18 -0700
Subject: [PATCH] [CI] Shard batch invariance, Kimi prefix cache, EP eval and
 generate APIs automatically (#61474)

Signed-off-by: Kevin Luu <51931015+khluu@users.noreply.github.com>
Co-authored-by: Codex <noreply@openai.com>
---
 .buildkite/test_areas/entrypoints.yaml  | 1 +
 .buildkite/test_areas/lm_eval.yaml      | 1 +
 .buildkite/test_areas/misc.yaml         | 2 ++
 .buildkite/test_areas/models_basic.yaml | 6 +++---
 4 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml
index bc7fa701b2..d125642f99 100644
--- a/.buildkite/test_areas/entrypoints.yaml
+++ b/.buildkite/test_areas/entrypoints.yaml
@@ -142,6 +142,7 @@ steps:
 - label: ":nvidia: (H200 MIG 35GB) Entrypoints Integration (API Server Generate)"
   device: h200_35gb
   key: entrypoints-integration-api-server-generate
+  automatic_shard: true
   timeout_in_minutes: 50
   working_dir: "/vllm-workspace/tests"
   source_file_dependencies:
diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml
index 96e5d4fc5e..430a6d32b9 100644
--- a/.buildkite/test_areas/lm_eval.yaml
+++ b/.buildkite/test_areas/lm_eval.yaml
@@ -384,6 +384,7 @@ steps:
 
 - label: ":nvidia: (B200) LM Eval Large Models EP"
   key: lm-eval-large-models-ep-2xb200
+  automatic_shard: true
   timeout_in_minutes: 60
   device: b200-k8s
   optional: true
diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml
index bf650fd756..198def78e0 100644
--- a/.buildkite/test_areas/misc.yaml
+++ b/.buildkite/test_areas/misc.yaml
@@ -733,6 +733,7 @@ steps:
 
 - label: ":nvidia: (B200) Batch Invariance"
   key: batch-invariance-b200
+  automatic_shard: true
   timeout_in_minutes: 45
   device: b200-k8s
   source_file_dependencies:
@@ -789,6 +790,7 @@ steps:
 
 - label: ":nvidia: (H100) Batch Invariance (Distributed)"
   key: batch-invariance-distributed-4xh100
+  automatic_shard: true
   timeout_in_minutes: 45
   device: h100
   optional: true
diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml
index 045e04533a..6845c2ebfb 100644
--- a/.buildkite/test_areas/models_basic.yaml
+++ b/.buildkite/test_areas/models_basic.yaml
@@ -158,8 +158,9 @@ steps:
       # cuda.bindings.driver), so it is deliberately left out here.
       - pytest -v -s models/kimi_k3
 
-- label: ":nvidia: (4xB200) Kimi K3 Prefix Cache Shard %N"
+- label: ":nvidia: (4xB200) Kimi K3 Prefix Cache"
   key: kimi-k3-prefix-cache-4xb200
+  automatic_shard: true
   timeout_in_minutes: 75
   device: b200-k8s
   num_devices: 4
@@ -177,8 +178,7 @@ steps:
   - .buildkite/scripts/install-kv-connectors.sh
   commands:
   - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
-  - pytest -v -s models/kimi_k3/test_prefix_cache.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
-  parallelism: 2
+  - pytest -v -s models/kimi_k3/test_prefix_cache.py
 
 - label: ":nvidia: (H200) GLM5Next"
   key: glm5next-unit-tests