A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 13:15 UTC
[CI/Test Fix] Fix CP tests on Blackwell (#28404)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Signed-off-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Lucas Wilkinson committed
Nov 10, 2025 at 20:36 UTC
39029d519276fddbe0c36440e0eefcdda069b969
2 files changed
+12
-1
tests/distributed/test_context_parallel.py
+12
@@ -14,6 +14,7 @@ from dataclasses import dataclass
14
from typing import Literal, NamedTuple
15
16
import pytest
17
+import torch
18
19
from vllm.config.model import RunnerOption
20
from vllm.logger import init_logger
@@ -254,6 +255,17 @@ def test_cp_generation(
255
test_options: CPTestOptions,
256
num_gpus_available,
257
):
258
+ if (
259
+ model_id == "deepseek-ai/DeepSeek-V2-Lite-Chat"
260
+ and torch.cuda.get_device_capability() < (9, 0)
261
+ ):
262
+ pytest.skip(reason="MLA+DCP requires compute capability of 9.0 or higher")
263
+ if (
264
+ model_id == "bigcode/gpt_bigcode-santacoder"
265
+ and torch.cuda.get_device_capability() != (9, 0)
266
+ ):
267
+ pytest.skip(reason="GQA+DCP currently requires compute capability of 9.0")
268
+
269
_compare_cp_with_tp(
270
model_id,
271
parallel_setup,
vllm/attention/ops/common.py
-1
@@ -195,7 +195,6 @@ def cp_lse_ag_out_rs(
195
cp_attn_lse = cp_attn_lse.contiguous()
196
lses = cp_group.all_gather(cp_attn_lse, dim=0).view_as(lses)
197
out, lse = correct_attn_out(cp_attn_out, lses, cp_group.rank_in_group, ctx)
198
- assert out.is_contiguous()
198
out = cp_group.reduce_scatter(out, dim=1)
199
200
if return_lse: