[habana_main]enable padding_aware_scheduler for speculative decoding (#1264)

performance number is at https://jira.habana-labs.com/browse/SW-211730

Tested with MTP deepseek R1

with batch_size == 4, see 1.6x improvement w/wo MTP in throughput.
```
Method: no_spec
Acceptance rate:  [0, 0]

Method: mtp
Acceptance rate:  [2180, 1918]
```

with batch_size == 32, see 1.5x improvement
```
Method: no_spec
Acceptance rate:  [0, 0]

Method: mtp
Acceptance rate:  [17408, 15371]
```

with batch_size == 128, see 1.35x improvement
```
Method: no_spec
Acceptance rate:  [0, 0]

Method: mtp
Acceptance rate:  [69988, 61143]
```

with batch_size == 256, see no obvious improvement (which is reasonable
because we hit maximum batch_size)
```
method: no_spec
Acceptance rate:  [0, 0]

Method: mtp
Acceptance rate:  [139286, 122973]
```

Signed-off-by: Chendi Xue <[email protected]>

Author: xuechendi

tests/spec_decode/e2e/test_eagle_correctness.py

@@ -52,9 +52,6 @@
 
         # Main model
         "model_name": MAIN_MODEL,
-
-        # schedule
-        "use_padding_aware_scheduling": False,
     }])
 @pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
 @pytest.mark.parametrize("baseline_llm_kwargs", [{}])

tests/spec_decode/e2e/test_medusa_correctness.py

@@ -55,9 +55,6 @@
 
         # Main model
         "model_name": MAIN_MODEL,
-
-        # schedule
-        "use_padding_aware_scheduling": False,
     }])
 @pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
 @pytest.mark.parametrize("baseline_llm_kwargs", [{}])

tests/spec_decode/e2e/test_mlp_correctness.py

@@ -57,9 +57,6 @@
 
         # Main model
         "model_name": MAIN_MODEL,
-
-        # schedule
-        "use_padding_aware_scheduling": False,
     }])
 @pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
 @pytest.mark.parametrize("baseline_llm_kwargs", [{}])

tests/spec_decode/e2e/test_mtp_correctness.py

@@ -52,9 +52,6 @@
 
         # GPU memory utilization
         "gpu_memory_utilization": 0.85,
-
-        # scheduler
-        "use_padding_aware_scheduling": False,
     }])
 @pytest.mark.parametrize("per_test_common_llm_kwargs", [{}])
 @pytest.mark.parametrize("baseline_llm_kwargs", [{}])

vllm/sequence.py

@@ -1336,7 +1336,14 @@ def prune(self,
         seq_ids = get_all_seq_ids(seq_group_metadata_list)
         if seq_ids != self._seq_ids:
             # Batch contents changed - prune removed sequences.
-            index = [self._seq_ids.index(seq_id) for seq_id in seq_ids]
+            if len(seq_ids) < len(self._seq_ids):
+                index = [self._seq_ids.index(seq_id) for seq_id in seq_ids]
+            else:
+                # This path is added for use_padding_aware_scheduling
+                index = [
+                    self._seq_ids.index(seq_id)
+                    if seq_id in self._seq_ids else 0 for seq_id in seq_ids
+                ]
             self.hidden_states = self.hidden_states[index]
             if self.second_last_token_hidden_states is not None:
                 self.second_last_token_hidden_states = self\