[AMD][CI] CI: fix AMD 2-GPU multimodal-gen partition-count abort

.github/workflows/pr-test-amd-rocm720.yml

@@ -813,7 +813,7 @@ jobs:
       fail-fast: false
       max-parallel: 1
       matrix:
-        part: [0, 1, 2]  # 3 partitions: 2 parametrized + 1 standalone (single_test_file/test_disagg_server.py)
+        part: [0, 1, 2]  # run_suite.py load-balances the suite's parametrized cases and standalone files over these 3 partitions
     runs-on: ${{ format('linux-{0}-2gpu-sglang', inputs.runner_arch || 'mi300') }}
     steps:
       - name: Checkout code

.github/workflows/pr-test-amd.yml

@@ -803,7 +803,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        part: [0, 1, 2]  # 3 partitions: 2 parametrized + 1 standalone (single_test_file/test_disagg_server.py)
+        part: [0, 1, 2]  # run_suite.py load-balances the suite's parametrized cases and standalone files over these 3 partitions
     runs-on: ${{ format('linux-{0}-2gpu-sglang', inputs.runner_arch || 'mi300') }}
     steps:
       - name: Checkout code

python/sglang/multimodal_gen/test/partitioning.py

@@ -30,3 +30,17 @@ def partition_items_by_lpt(
         partition_sums[min_idx] += item.est_time
 
     return partitions
+
+
+def assign_partition(
+    items: list[PartitionItem], partition_id: int, num_partitions: int
+) -> list[PartitionItem]:
+    """Return the LPT slice of ``items`` owned by ``partition_id``.
+
+    The LPT pass is deterministic, so shards that each call this with the same
+    item list cover the list exactly once between them.
+    """
+    partitions = partition_items_by_lpt(items, num_partitions)
+    if partition_id < 0 or partition_id >= len(partitions):
+        return []
+    return partitions[partition_id]