From 710ebaa1897e930bf039f6c6deba75bd0aa75a77 Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Wed, 24 Jun 2026 23:07:28 -0500 Subject: [PATCH] [ROCm][Bugfix] Fix chunk alignment when using context parallelism with TRITON_MLA (#46114) Signed-off-by: Micah Williamson Co-authored-by: Andreas Karatzas Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- tests/distributed/test_context_parallel.py | 75 ++++++++++++------- .../layers/attention/mla_attention.py | 15 ++-- 2 files changed, 52 insertions(+), 38 deletions(-) diff --git a/tests/distributed/test_context_parallel.py b/tests/distributed/test_context_parallel.py index a2863092177..484d29c5b53 100644 --- a/tests/distributed/test_context_parallel.py +++ b/tests/distributed/test_context_parallel.py @@ -13,13 +13,14 @@ import os from dataclasses import dataclass from typing import Literal, NamedTuple +import lm_eval import pytest import torch -from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k from tests.utils import RemoteOpenAIServer, create_new_process_for_each_test from vllm.config.model import RunnerOption from vllm.logger import init_logger +from vllm.platforms import current_platform from ..models.registry import HF_EXAMPLE_MODELS @@ -35,8 +36,10 @@ CP_TEST_MODELS = [ ] # GSM8K eval configuration -NUM_QUESTIONS = 256 # Fast eval for CI NUM_SHOTS = 5 # Few-shot examples +TASK = "gsm8k" +FILTER = "exact_match,strict-match" +NUM_CONCURRENT = 128 # tp accuracy with 2% buffer MIN_ACCURACY = { # .buildkite/lm-eval-harness/configs/DeepSeek-V2-Lite-Chat.yaml @@ -121,24 +124,34 @@ class CPTestSettings: ) -CP_TEXT_GENERATION_MODELS = { - "deepseek-ai/DeepSeek-V2-Lite-Chat": [ - CPTestSettings.detailed(dcp_multipliers=[1]), - CPTestSettings.detailed( - dcp_multipliers=[0.5], - cp_kv_cache_interleave_size=64, - attn_backend="FLASHMLA", - ), - ], - "Qwen/Qwen2.5-1.5B-Instruct": [ - CPTestSettings.detailed( - cp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN" - ), - CPTestSettings.detailed( - cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER" - ), - ], -} +if current_platform.is_rocm(): + CP_TEXT_GENERATION_MODELS = { + "deepseek-ai/DeepSeek-V2-Lite-Chat": [ + CPTestSettings.detailed(dcp_multipliers=[1]), + ], + "Qwen/Qwen2.5-1.5B-Instruct": [ + CPTestSettings.detailed(dcp_multipliers=[1]), + ], + } +else: + CP_TEXT_GENERATION_MODELS = { + "deepseek-ai/DeepSeek-V2-Lite-Chat": [ + CPTestSettings.detailed(dcp_multipliers=[1]), + CPTestSettings.detailed( + dcp_multipliers=[0.5], + cp_kv_cache_interleave_size=64, + attn_backend="FLASHMLA", + ), + ], + "Qwen/Qwen2.5-1.5B-Instruct": [ + CPTestSettings.detailed( + cp_kv_cache_interleave_size=16, attn_backend="FLASH_ATTN" + ), + CPTestSettings.detailed( + cp_kv_cache_interleave_size=16, attn_backend="FLASHINFER" + ), + ], + } def _test_cp_gsm8k( @@ -227,19 +240,23 @@ def _test_cp_gsm8k( server_args, max_wait_seconds=720, ) as remote_server: - host = f"http://{remote_server.host}" - port = remote_server.port + url = f"{remote_server.url_for('v1')}/completions" - # Run GSM8K evaluation - results = evaluate_gsm8k( - num_questions=NUM_QUESTIONS, - num_shots=NUM_SHOTS, - host=host, - port=port, + model_args = ( + f"model={model_id}," + f"base_url={url}," + f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False" + ) + + results = lm_eval.simple_evaluate( + model="local-completions", + model_args=model_args, + tasks=TASK, + num_fewshot=NUM_SHOTS, ) # Validate accuracy is reasonable - accuracy = results["accuracy"] + accuracy = results["results"][TASK][FILTER] min_accuracy = MIN_ACCURACY[model_id] assert accuracy >= min_accuracy, ( f"TP+DCP accuracy too low: {accuracy:.3f} < {min_accuracy:.3f}" diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 4dd666f0c64..8d9a674319d 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1532,9 +1532,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size - # Don't try to access the runner on AMD - if self.aot_schedule: - self.page_size = self.kv_cache_spec.block_size + self.page_size = self.kv_cache_spec.block_size self.chunked_prefill_workspace_size = ( self.determine_chunked_prefill_workspace_size(vllm_config) @@ -1684,12 +1682,11 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): self.chunked_prefill_workspace_size // num_prefills_with_context_cpu ) - if self.aot_schedule: - # align max_context_chunk to page_size by rounding down, - # currently the `gather_and_maybe_dequant_cache` kernel - # cannot handle `context_chunk_starts` that are not aligned - # to page_size - max_context_chunk = round_down(max_context_chunk, self.page_size) + # align max_context_chunk to page_size by rounding down, + # currently the `gather_and_maybe_dequant_cache` kernel + # cannot handle `context_chunk_starts` that are not aligned + # to page_size + max_context_chunk = round_down(max_context_chunk, self.page_size) assert max_context_chunk > 0 num_chunks = cdiv(max_context_len_cpu, max_context_chunk)