[ROCm][Bugfix] Fix chunk alignment when using context parallelism with TRITON_MLA (#46114)
Signed-off-by: Micah Williamson <micah.williamson@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Micah Williamson committed
Jun 24, 2026 at 23:07 UTC
710ebaa1897e930bf039f6c6deba75bd0aa75a77
2 files changed
+53
-39
tests/distributed/test_context_parallel.py
+47
-30
@@ -13,13 +13,14 @@ import os
13
from dataclasses import dataclass
14
from typing import Literal, NamedTuple
15
16
+import lm_eval
17
import pytest
18
import torch
19
19
-from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k
20
from tests.utils import RemoteOpenAIServer, create_new_process_for_each_test
21
from vllm.config.model import RunnerOption
22
from vllm.logger import init_logger
23
+from vllm.platforms import current_platform
24
25
from ..models.registry import HF_EXAMPLE_MODELS
26
@@ -35,8 +36,10 @@ CP_TEST_MODELS = [
36
]
37
38
# GSM8K eval configuration
38
-NUM_QUESTIONS = 256 # Fast eval for CI
39
NUM_SHOTS = 5 # Few-shot examples
40
+TASK = "gsm8k"
41
+FILTER = "exact_match,strict-match"
42
+NUM_CONCURRENT = 128
43
# tp accuracy with 2% buffer
44
MIN_ACCURACY = {
45
# .buildkite/lm-eval-harness/configs/DeepSeek-V2-Lite-Chat.yaml
@@ -121,24 +124,34 @@ class CPTestSettings:
124
)
125
126
124
-CP_TEXT_GENERATION_MODELS = {
125
- "deepseek-ai/DeepSeek-V2-Lite-Chat": [
126
- CPTestSettings.detailed(dcp_multipliers=[1]),
127
- CPTestSettings.detailed(
128
- dcp_multipliers=[0.5],
129
- cp_kv_cache_interleave_size=64,
130
- attn_backend="FLASHMLA",
131
- ),
132
- ],
133
- "Qwen/Qwen2.5-1.5B-Instruct": [
134
- CPTestSettings.detailed(
135
- cp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN"
136
- ),
137
- CPTestSettings.detailed(
138
- cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER"
139
- ),
140
- ],
141
-}
127
+if current_platform.is_rocm():
128
+ CP_TEXT_GENERATION_MODELS = {
129
+ "deepseek-ai/DeepSeek-V2-Lite-Chat": [
130
+ CPTestSettings.detailed(dcp_multipliers=[1]),
131
+ ],
132
+ "Qwen/Qwen2.5-1.5B-Instruct": [
133
+ CPTestSettings.detailed(dcp_multipliers=[1]),
134
+ ],
135
+ }
136
+else:
137
+ CP_TEXT_GENERATION_MODELS = {
138
+ "deepseek-ai/DeepSeek-V2-Lite-Chat": [
139
+ CPTestSettings.detailed(dcp_multipliers=[1]),
140
+ CPTestSettings.detailed(
141
+ dcp_multipliers=[0.5],
142
+ cp_kv_cache_interleave_size=64,
143
+ attn_backend="FLASHMLA",
144
+ ),
145
+ ],
146
+ "Qwen/Qwen2.5-1.5B-Instruct": [
147
+ CPTestSettings.detailed(
148
+ cp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN"
149
+ ),
150
+ CPTestSettings.detailed(
151
+ cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER"
152
+ ),
153
+ ],
154
+ }
155
156
157
def _test_cp_gsm8k(
@@ -227,19 +240,23 @@ def _test_cp_gsm8k(
240
server_args,
241
max_wait_seconds=720,
242
) as remote_server:
230
- host = f"http://{remote_server.host}"
231
- port = remote_server.port
232
-
233
- # Run GSM8K evaluation
234
- results = evaluate_gsm8k(
235
- num_questions=NUM_QUESTIONS,
236
- num_shots=NUM_SHOTS,
237
- host=host,
238
- port=port,
243
+ url = f"{remote_server.url_for('v1')}/completions"
244
+
245
+ model_args = (
246
+ f"model={model_id},"
247
+ f"base_url={url},"
248
+ f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False"
249
+ )
250
+
251
+ results = lm_eval.simple_evaluate(
252
+ model="local-completions",
253
+ model_args=model_args,
254
+ tasks=TASK,
255
+ num_fewshot=NUM_SHOTS,
256
)
257
258
# Validate accuracy is reasonable
242
- accuracy = results["accuracy"]
259
+ accuracy = results["results"][TASK][FILTER]
260
min_accuracy = MIN_ACCURACY[model_id]
261
assert accuracy >= min_accuracy, (
262
f"TP+DCP accuracy too low: {accuracy:.3f} < {min_accuracy:.3f}"
vllm/model_executor/layers/attention/mla_attention.py
+6
-9
@@ -1532,9 +1532,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1532
self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size
1533
self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size
1534
1535
- # Don't try to access the runner on AMD
1536
- if self.aot_schedule:
1537
- self.page_size = self.kv_cache_spec.block_size
1535
+ self.page_size = self.kv_cache_spec.block_size
1536
1537
self.chunked_prefill_workspace_size = (
1538
self.determine_chunked_prefill_workspace_size(vllm_config)
@@ -1684,12 +1682,11 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
1682
self.chunked_prefill_workspace_size // num_prefills_with_context_cpu
1683
)
1684
1687
- if self.aot_schedule:
1688
- # align max_context_chunk to page_size by rounding down,
1689
- # currently the `gather_and_maybe_dequant_cache` kernel
1690
- # cannot handle `context_chunk_starts` that are not aligned
1691
- # to page_size
1692
- max_context_chunk = round_down(max_context_chunk, self.page_size)
1685
+ # align max_context_chunk to page_size by rounding down,
1686
+ # currently the `gather_and_maybe_dequant_cache` kernel
1687
+ # cannot handle `context_chunk_starts` that are not aligned
1688
+ # to page_size
1689
+ max_context_chunk = round_down(max_context_chunk, self.page_size)
1690
1691
assert max_context_chunk > 0
1692
num_chunks = cdiv(max_context_len_cpu, max_context_chunk)