@tduncan / Nvidia-Optimization-Stack-VLLM / commits / db3814a4f2

[Attention] Require explicit DCP support from attention implementations (#55780)

Signed-off-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: OpenAI Codex <codex@openai.com>

Andreas Karatzas committed Sep 8, 2026 at 16:04 UTC db3814a4f215e666098ffb74a5c21bf46c60bae9

15 files changed +199 -19

tests/kernels/attention/test_attention_selector.py

+94 -16

@@ -10,6 +10,7 @@ import torch

10 from vllm.config import (

11 AttentionConfig,

12 CacheConfig,

13 + ParallelConfig,

14 VllmConfig,

15 set_current_vllm_config,

16 )

@@ -29,7 +30,11 @@ else:

31 from vllm.v1.attention.backend import AttentionType

32 from vllm.v1.attention.backends.registry import AttentionBackendEnum

32 -from vllm.v1.attention.selector import _cached_get_attn_backend, get_attn_backend

33 +from vllm.v1.attention.selector import (

34 + AttentionSelectorConfig,

35 + _cached_get_attn_backend,

36 + get_attn_backend,

37 +)

40 @pytest.fixture(autouse=True)

@@ -77,33 +82,54 @@ def generate_params():

82 else DEVICE_REGULAR_ATTN_BACKENDS[device]

83 )

84 for name in backends:

80 - block_sizes = DEVICE_MLA_BLOCK_SIZES[device] if use_mla else [16]

85 + block_sizes = (

86 + [128]

87 + if name == "CUTLASS_MLA"

88 + else DEVICE_MLA_BLOCK_SIZES[device]

89 + if use_mla

90 + else [16]

91 + )

92 for block_size in block_sizes:

82 - params.append(

83 - pytest.param(

84 - device,

85 - name,

86 - use_mla,

87 - block_size,

88 - id=f"{device}_{name}_mla_{str(use_mla)[0]}_blks{block_size}",

93 + for use_dcp in [False, True] if device != "cpu" else [False]:

94 + params.append(

95 + pytest.param(

96 + device,

97 + name,

98 + use_mla,

99 + block_size,

100 + use_dcp,

101 + id=(

102 + f"{device}_{name}_mla_{str(use_mla)[0]}"

103 + f"_blks{block_size}_dcp{use_dcp}"

104 + ),

105 + )

106 )

90 - )

107 return params

108

109

94 -@pytest.mark.parametrize("device, name, use_mla, block_size", generate_params())

110 +@pytest.mark.parametrize(

111 + "device, name, use_mla, block_size, use_dcp", generate_params()

112 +)

113 def test_backend_selection(

114 device: str,

115 name: str,

116 use_mla: bool,

117 block_size: int,

118 + use_dcp: bool,

119 ):

101 - """Test attention backend selection with valid device-backend pairs."""

120 + """Supported GPU backends remain selectable when DCP is enabled."""

121 # Create AttentionConfig with the specified backend

122 attention_config = AttentionConfig(backend=AttentionBackendEnum[name])

123 cache_config = CacheConfig(block_size=block_size)

124 vllm_config = VllmConfig(

106 - attention_config=attention_config, cache_config=cache_config

125 + attention_config=attention_config,

126 + cache_config=cache_config,

127 + parallel_config=ParallelConfig(

128 + # Skip GPU-count autodetection: selection does not launch workers.

129 + distributed_executor_backend="mp" if use_dcp else None,

130 + tensor_parallel_size=2 if use_dcp else 1,

131 + decode_context_parallel_size=2 if use_dcp else 1,

132 + ),

133 )

134

135 with set_current_vllm_config(vllm_config):

@@ -135,9 +161,15 @@ def test_backend_selection(

161 expected = name

162 assert backend.get_name() == expected

163 else:

138 - backend = get_attn_backend(32, torch.float16, None, use_mla=use_mla)

139 - expected = "ROCM_ATTN"

140 - assert backend.get_name() == expected

164 + if use_dcp:

165 + with pytest.raises(ValueError, match="DCP not supported"):

166 + get_attn_backend(32, torch.float16, None, use_mla=use_mla)

167 + else:

168 + backend = get_attn_backend(

169 + 32, torch.float16, None, use_mla=use_mla

170 + )

171 + expected = "ROCM_ATTN"

172 + assert backend.get_name() == expected

173

174 elif device == "cuda":

175 if CudaPlatform is None:

@@ -230,6 +262,52 @@ def test_backend_selection(

262 assert backend.get_name() == expected

263

264

265 +@pytest.mark.skipif(

266 + not current_platform.is_cuda_alike(), reason="GPU attention backends"

267 +)

268 +@pytest.mark.parametrize(

269 + "backend_name",

270 + [

271 + "FLASH_ATTN",

272 + "FLASH_ATTN_DIFFKV",

273 + "FLASHINFER",

274 + "TRITON_MLA",

275 + "ROCM_AITER_MLA",

276 + "ROCM_AITER_TRITON_MLA",

277 + "CUTLASS_MLA",

278 + "FLASHMLA",

279 + "FLASHINFER_MLA",

280 + "TOKENSPEED_MLA",

281 + "FLASH_ATTN_MLA",

282 + "FLASHMLA_SPARSE",

283 + "FLASHINFER_MLA_SPARSE",

284 + ],

285 +)

286 +def test_supported_backend_preserves_dcp_eligibility(backend_name, default_vllm_config):

287 + """DCP opt-ins preserve eligibility and unrelated configuration restrictions."""

288 + if backend_name.startswith("ROCM_") and not current_platform.is_rocm():

289 + pytest.skip("ROCm-specific backend")

290 + try:

291 + backend_cls = AttentionBackendEnum[backend_name].get_class()

292 + except ImportError as error:

293 + pytest.skip(f"Optional backend dependency unavailable: {error}")

294 +

295 + config = AttentionSelectorConfig(

296 + head_size=576 if backend_cls.is_mla() else 128,

297 + dtype=torch.bfloat16,

298 + kv_cache_dtype="auto",

299 + block_size=backend_cls.get_preferred_block_size(64),

300 + use_mla=backend_cls.is_mla(),

301 + use_sparse=backend_cls.is_sparse(),

302 + )

303 + kwargs = dict(

304 + device_capability=current_platform.get_device_capability(), **config._asdict()

305 + )

306 + invalid_without_dcp = backend_cls.validate_configuration(**kwargs)

307 + kwargs["use_dcp"] = True

308 + assert backend_cls.validate_configuration(**kwargs) == invalid_without_dcp

309 +

310 +

311 @pytest.mark.parametrize("device", ["cpu", "cuda", "hip"])

312 def test_fp32_fallback(device: str):

313 """Test attention backend selection with fp32."""

tests/test_attention_backend_registry.py

+21

@@ -1,5 +1,8 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import torch

4 +

5 +from vllm.platforms.interface import DeviceCapability

6 from vllm.v1.attention.backend import (

7 AttentionBackend,

8 AttentionImpl,

@@ -9,6 +12,7 @@ from vllm.v1.attention.backends.registry import (

12 MambaAttentionBackendEnum,

13 register_backend,

14 )

15 +from vllm.v1.attention.selector import AttentionSelectorConfig

18 class CustomAttentionImpl(AttentionImpl):

@@ -67,6 +71,23 @@ class CustomMambaAttentionBackend(AttentionBackend):

71 return None

74 +def test_custom_attention_backend_requires_dcp_opt_in():

75 + """A new implementation remains usable without implicitly opting into DCP."""

76 + config = AttentionSelectorConfig(

77 + head_size=64,

78 + dtype=torch.float16,

79 + kv_cache_dtype="auto",

80 + block_size=16,

81 + )

82 + kwargs = dict(device_capability=DeviceCapability(9, 0), **config._asdict())

83 + assert CustomAttentionBackend.validate_configuration(**kwargs) == []

84 +

85 + kwargs["use_dcp"] = True

86 + assert CustomAttentionBackend.validate_configuration(**kwargs) == [

87 + "DCP not supported"

88 + ]

89 +

90 +

91 def test_custom_is_not_alias_of_any_backend():

92 # Get all members of AttentionBackendEnum

93 all_backends = list(AttentionBackendEnum)

tests/v1/attention/test_rocm_attention_backends_selection.py

+72 -2

@@ -74,6 +74,7 @@ def test_aiter_unified_attention_capture_preserves_query_start_locations():

74 assert torch.equal(actual.query_start_loc, expected_query_start_loc)

77 +@pytest.mark.parametrize("use_dcp", [False, True])

78 @pytest.mark.parametrize(

79 "env_vars, selected_backend, expected_backend_path",

80 [

@@ -143,11 +144,12 @@ def test_standard_attention_backend_selection(

144 env_vars,

145 selected_backend,

146 expected_backend_path,

147 + use_dcp,

148 mock_vllm_config,

149 mock_get_cdna_version,

150 monkeypatch,

151 ):

150 - """Test standard attention backend selection with various configurations."""

152 + """Standard ROCm backends remain selectable without DCP and reject DCP."""

153 # Set environment variables

154 for key, value in env_vars.items():

155 monkeypatch.setenv(key, value)

@@ -176,8 +178,14 @@ def test_standard_attention_backend_selection(

178 use_mla=False,

179 has_sink=False,

180 use_sparse=False,

181 + use_dcp=use_dcp,

182 )

183

184 + if use_dcp:

185 + with pytest.raises(ValueError, match="DCP not supported"):

186 + RocmPlatform.get_attn_backend_cls(backend_enum, attn_selector_config)

187 + return

188 +

189 backend_path = RocmPlatform.get_attn_backend_cls(

190 selected_backend=backend_enum, attn_selector_config=attn_selector_config

191 )

@@ -185,6 +193,7 @@ def test_standard_attention_backend_selection(

193 assert backend_path == expected_backend_path

194

195

196 +@pytest.mark.parametrize("use_dcp", [False, True])

197 @pytest.mark.parametrize(

198 "env_vars, selected_backend, block_size, expected_backend_path, should_raise",

199 [

@@ -261,10 +270,11 @@ def test_mla_backend_selection(

270 block_size,

271 expected_backend_path,

272 should_raise,

273 + use_dcp,

274 mock_vllm_config,

275 monkeypatch,

276 ):

267 - """Test MLA backend selection with various configurations."""

277 + """Dense MLA remains selectable with DCP for valid block sizes."""

278 # Set environment variables

279 for key, value in env_vars.items():

280 monkeypatch.setenv(key, value)

@@ -303,6 +313,7 @@ def test_mla_backend_selection(

313 use_mla=True,

314 has_sink=False,

315 use_sparse=False,

316 + use_dcp=use_dcp,

317 )

318 attn_selector_config = AttentionSelectorConfig(

319 head_size=128,

@@ -312,6 +323,7 @@ def test_mla_backend_selection(

323 use_mla=True,

324 has_sink=False,

325 use_sparse=False,

326 + use_dcp=use_dcp,

327 )

328 backend_path = RocmPlatform.get_attn_backend_cls(

329 selected_backend=backend_enum,

@@ -327,6 +339,7 @@ def test_mla_backend_selection(

339 use_mla=True,

340 has_sink=False,

341 use_sparse=False,

342 + use_dcp=use_dcp,

343 )

344

345 backend_path = RocmPlatform.get_attn_backend_cls(

@@ -336,6 +349,63 @@ def test_mla_backend_selection(

349 assert backend_path == expected_backend_path

350

351

352 +@pytest.mark.parametrize("use_dcp", [False, True])

353 +@pytest.mark.parametrize(

354 + "selected_backend", [None, AttentionBackendEnum.ROCM_AITER_MLA_SPARSE]

355 +)

356 +def test_sparse_mla_backend_rejects_dcp(selected_backend, use_dcp):

357 + """Sparse MLA remains selectable without DCP and fails early with DCP."""

358 + from vllm.platforms.rocm import RocmPlatform

359 +

360 + selector_config = AttentionSelectorConfig(

361 + head_size=576,

362 + dtype=torch.bfloat16,

363 + kv_cache_dtype="auto",

364 + block_size=16,

365 + use_mla=True,

366 + use_sparse=True,

367 + use_dcp=use_dcp,

368 + )

369 + if use_dcp:

370 + with pytest.raises(ValueError, match="DCP not supported"):

371 + RocmPlatform.get_attn_backend_cls(selected_backend, selector_config)

372 + else:

373 + assert RocmPlatform.get_attn_backend_cls(selected_backend, selector_config) == (

374 + AttentionBackendEnum.ROCM_AITER_MLA_SPARSE.get_path()

375 + )

376 +

377 +

378 +@pytest.mark.parametrize("use_dcp", [False, True])

379 +@pytest.mark.parametrize(

380 + "selected_backend, head_size, kv_cache_dtype",

381 + [

382 + (AttentionBackendEnum.TRITON_ATTN_DIFFKV, 192, "bfloat16"),

383 + # A 128-dimensional head uses 118 packed bytes, or 59 fp16 elements.

384 + (AttentionBackendEnum.TURBOQUANT, 59, "turboquant_k3v4_nc"),

385 + ],

386 +)

387 +def test_specialized_attention_backends_reject_dcp(

388 + selected_backend, head_size, kv_cache_dtype, use_dcp

389 +):

390 + """Valid DiffKV and compressed-cache configurations must reject DCP."""

391 + from vllm.platforms.rocm import RocmPlatform

392 +

393 + selector_config = AttentionSelectorConfig(

394 + head_size=head_size,

395 + dtype=torch.bfloat16,

396 + kv_cache_dtype=kv_cache_dtype,

397 + block_size=16,

398 + use_dcp=use_dcp,

399 + )

400 + if use_dcp:

401 + with pytest.raises(ValueError, match="DCP not supported"):

402 + RocmPlatform.get_attn_backend_cls(selected_backend, selector_config)

403 + else:

404 + assert RocmPlatform.get_attn_backend_cls(selected_backend, selector_config) == (

405 + selected_backend.get_path()

406 + )

407 +

408 +

409 def test_aiter_fa_requires_mi3xx(mock_vllm_config):

410 """Test that ROCM_AITER_FA requires CDNA3+ architecture."""

411 from vllm.platforms.rocm import RocmPlatform

vllm/v1/attention/backend.py

+1 -1

@@ -819,7 +819,7 @@ class AttentionImplBase(ABC, Generic[T]):

819 # Whether the attention impl supports Prefill Context Parallelism.

820 supports_pcp: bool = False

821 # Whether the attention impl supports Decode Context Parallelism.

822 - supports_dcp: bool = True

822 + supports_dcp: bool = False

823 # Whether the attention impl(or ops) supports MTP

824 # when cp_kv_cache_interleave_size > 1

825 supports_mtp_with_cp_non_trivial_interleave_size: bool = False

vllm/v1/attention/backends/flash_attn.py

@@ -832,6 +832,7 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad

832

833 class FlashAttentionImpl(AttentionImpl):

834 can_return_lse_for_decode: bool = True

835 + supports_dcp: bool = True

836

837 def __init__(

838 self,

vllm/v1/attention/backends/flashinfer.py

@@ -1779,6 +1779,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):

1779

1780 class FlashInferImpl(AttentionImpl):

1781 can_return_lse_for_decode: bool = True

1782 + supports_dcp: bool = True

1783

1784 def __init__(

1785 self,

vllm/v1/attention/backends/mla/cutlass_mla.py

@@ -103,6 +103,7 @@ MAX_HEADS = 128

103

104 class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):

105 can_return_lse_for_decode: bool = True

106 + supports_dcp: bool = True

107

108 def __init__(

109 self,

vllm/v1/attention/backends/mla/flashattn_mla.py

@@ -255,6 +255,7 @@ class FlashAttnMLAMetadataBuilder(MLACommonMetadataBuilder[FlashAttnMLAMetadata]

255

256 class FlashAttnMLAImpl(MLACommonImpl[FlashAttnMLAMetadata]):

257 can_return_lse_for_decode: bool = True

258 + supports_dcp: bool = True

259

260 def __init__(

261 self,

vllm/v1/attention/backends/mla/flashinfer_mla.py

@@ -210,6 +210,7 @@ class FlashInferMLABackend(MLACommonBackend):

210

211 class FlashInferMLAImpl(MLACommonImpl[MLACommonMetadata]):

212 can_return_lse_for_decode: bool = True

213 + supports_dcp: bool = True

214 # DCP is the only path that consumes LSE. It uses monolithic CuTeDSL,

215 # whose public LSE contract is natural-log.

216 lse_base_on_e: bool = True

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -319,6 +319,7 @@ class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):

319 """

320

321 can_return_lse_for_decode: bool = True

322 + supports_dcp: bool = True

323 lse_base_on_e: bool = False

324

325 def __init__(

vllm/v1/attention/backends/mla/flashmla.py

@@ -207,6 +207,7 @@ class FlashMLAMetadataBuilder(MLACommonMetadataBuilder[FlashMLAMetadata]):

207

208 class FlashMLAImpl(MLACommonImpl[FlashMLAMetadata]):

209 can_return_lse_for_decode: bool = True

210 + supports_dcp: bool = True

211

212 def __init__(

213 self,

vllm/v1/attention/backends/mla/flashmla_sparse.py

@@ -582,6 +582,7 @@ class FlashMLASparseMetadataBuilder(

582

583 class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

584 can_return_lse_for_decode: bool = True

585 + supports_dcp: bool = True

586

587 @staticmethod

588 def _compute_fp8_decode_padded_heads(num_heads: int) -> int:

vllm/v1/attention/backends/mla/rocm_aiter_mla.py

@@ -1421,6 +1421,7 @@ class AiterMLAHelper:

1421 class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]):

1422 # DCP decode paths return natural-log softmax LSE for the cross-rank merge.

1423 can_return_lse_for_decode: bool = True

1424 + supports_dcp: bool = True

1425 # Measured on gfx950: aiter.mla.mla_decode_fwd(return_lse=True) matches

1426 # logsumexp to fp32 exactly, and merge_mla_segments_triton converts AITER's

1427 # base-2 segment statistics with LOGE2. Stated rather than inherited because

vllm/v1/attention/backends/mla/tokenspeed_mla.py

@@ -154,6 +154,7 @@ class TokenspeedMLABackend(MLACommonBackend):

154

155 class TokenspeedMLAImpl(MLACommonImpl[MLACommonMetadata]):

156 can_return_lse_for_decode: bool = True

157 + supports_dcp: bool = True

158 # tokenspeed_mla_decode returns LSE in log2 units; its own DCP test merges

159 # partial outputs with exp2(lse).

160 lse_base_on_e: bool = False

vllm/v1/attention/backends/mla/triton_mla.py

@@ -181,6 +181,7 @@ class TritonMLABackend(MLACommonBackend):

181

182 class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

183 can_return_lse_for_decode: bool = True

184 + supports_dcp: bool = True

185

186 def __init__(

187 self,