@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9898f94abe

Fork 0

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 3, 15:15 UTC

[Attention] Remove deprecated MLA prefill arguments (#42555)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed May 14, 2026 at 13:34 UTC 9898f94abe005f675e0a7f40b0ae891afb5cf992

5 files changed +27 -126

benchmarks/attention_benchmarks/mla_runner.py

+18 -39

@@ -29,6 +29,7 @@ from vllm.config import (

29 VllmConfig,

30 set_current_vllm_config,

31 )

32 +from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum

34 # ============================================================================

35 # VllmConfig Creation

@@ -79,8 +80,8 @@ def create_minimal_vllm_config(

80 index_topk: Optional topk value for sparse MLA backends. If provided,

81 the config will include index_topk for sparse attention.

82 prefill_backend: Prefill backend name (e.g., "fa3", "fa4", "flashinfer",

82 - "cudnn", "trtllm"). Configures the attention config to

83 - force the specified prefill backend.

83 + "trtllm"). Configures the attention config to force

84 + the specified prefill backend.

86 Returns:

87 VllmConfig for benchmarking

@@ -179,27 +180,13 @@ def create_minimal_vllm_config(

180

181 if prefill_backend is not None:

182 prefill_cfg = get_prefill_backend_config(prefill_backend)

182 - if prefill_cfg.get("mla_prefill_backend_enum") is not None:

183 - # Registry-based backends bypass the deprecated boolean flags.

184 - from vllm.v1.attention.backends.mla.prefill import MLAPrefillBackendEnum

185 -

186 - vllm_config.attention_config.mla_prefill_backend = MLAPrefillBackendEnum[

187 - prefill_cfg["mla_prefill_backend_enum"]

188 - ]

189 - else:

190 - if prefill_cfg["flash_attn_version"] is not None:

191 - vllm_config.attention_config.flash_attn_version = prefill_cfg[

192 - "flash_attn_version"

193 - ]

194 - vllm_config.attention_config.disable_flashinfer_prefill = prefill_cfg[

195 - "disable_flashinfer_prefill"

196 - ]

197 - vllm_config.attention_config.use_cudnn_prefill = prefill_cfg[

198 - "use_cudnn_prefill"

183 + vllm_config.attention_config.mla_prefill_backend = prefill_cfg[

184 + "mla_prefill_backend"

185 + ]

186 + if prefill_cfg["flash_attn_version"] is not None:

187 + vllm_config.attention_config.flash_attn_version = prefill_cfg[

188 + "flash_attn_version"

189 ]

200 - vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill = (

201 - prefill_cfg["use_trtllm_ragged_deepseek_prefill"]

202 - )

190

191 return vllm_config

192

@@ -214,34 +201,27 @@ def create_minimal_vllm_config(

201 _PREFILL_BACKEND_CONFIG: dict[str, dict] = {

202 "fa2": {

203 "flash_attn_version": 2,

217 - "disable_flashinfer_prefill": True,

218 - "use_cudnn_prefill": False,

219 - "use_trtllm_ragged_deepseek_prefill": False,

204 + "mla_prefill_backend": MLAPrefillBackendEnum.FLASH_ATTN,

205 },

206 "fa3": {

207 "flash_attn_version": 3,

223 - "disable_flashinfer_prefill": True,

224 - "use_cudnn_prefill": False,

225 - "use_trtllm_ragged_deepseek_prefill": False,

208 + "mla_prefill_backend": MLAPrefillBackendEnum.FLASH_ATTN,

209 },

210 "fa4": {

211 "flash_attn_version": 4,

229 - "disable_flashinfer_prefill": True,

230 - "use_cudnn_prefill": False,

231 - "use_trtllm_ragged_deepseek_prefill": False,

212 + "mla_prefill_backend": MLAPrefillBackendEnum.FLASH_ATTN,

213 },

214 "flashinfer": {

234 - "mla_prefill_backend_enum": "FLASHINFER",

235 - },

236 - "cudnn": {

237 - # cuDNN prefill backend was removed; AttentionConfig raises on use.

238 - "mla_prefill_backend_enum": "FLASHINFER",

215 + "flash_attn_version": None,

216 + "mla_prefill_backend": MLAPrefillBackendEnum.FLASHINFER,

217 },

218 "trtllm": {

241 - "mla_prefill_backend_enum": "TRTLLM_RAGGED",

219 + "flash_attn_version": None,

220 + "mla_prefill_backend": MLAPrefillBackendEnum.TRTLLM_RAGGED,

221 },

222 "tokenspeed": {

244 - "mla_prefill_backend_enum": "TOKENSPEED_MLA",

223 + "flash_attn_version": None,

224 + "mla_prefill_backend": MLAPrefillBackendEnum.TOKENSPEED_MLA,

225 },

226 }

227

@@ -1020,7 +1000,6 @@ def _run_mla_benchmark_batched(

1000 f"version {fa_version}, got "

1001 f"{actual_fa_version} on {actual_class}."

1002 )

1023 -

1003 # Run each benchmark with the shared impl

1004 for config, threshold, num_splits in configs_with_params:

1005 # Set threshold for this benchmark (FlashAttn/FlashMLA only)

tests/engine/test_arg_utils.py

-12

@@ -368,12 +368,8 @@ def test_attention_config():

368 "true",

369 "--attention-config.flash_attn_max_num_splits_for_cuda_graph",

370 "16",

371 - "--attention-config.use_trtllm_ragged_deepseek_prefill",

372 - "true",

371 "--attention-config.use_trtllm_attention",

372 "true",

375 - "--attention-config.disable_flashinfer_prefill",

376 - "true",

373 "--attention-config.disable_flashinfer_q_quantization",

374 "true",

375 ]

@@ -385,9 +381,7 @@ def test_attention_config():

381 assert engine_args.attention_config.flash_attn_version == 3

382 assert engine_args.attention_config.use_prefill_decode_attention is True

383 assert engine_args.attention_config.flash_attn_max_num_splits_for_cuda_graph == 16

388 - assert engine_args.attention_config.use_trtllm_ragged_deepseek_prefill is True

384 assert engine_args.attention_config.use_trtllm_attention is True

390 - assert engine_args.attention_config.disable_flashinfer_prefill is True

385 assert engine_args.attention_config.disable_flashinfer_q_quantization is True

386

387 # set to string form of a dict with all fields

@@ -397,10 +391,7 @@ def test_attention_config():

391 '{"backend": "FLASHINFER", "flash_attn_version": 2, '

392 '"use_prefill_decode_attention": false, '

393 '"flash_attn_max_num_splits_for_cuda_graph": 8, '

400 - '"use_cudnn_prefill": false, '

401 - '"use_trtllm_ragged_deepseek_prefill": false, '

394 '"use_trtllm_attention": false, '

403 - '"disable_flashinfer_prefill": false, '

395 '"disable_flashinfer_q_quantization": false}',

396 ]

397 )

@@ -411,10 +402,7 @@ def test_attention_config():

402 assert engine_args.attention_config.flash_attn_version == 2

403 assert engine_args.attention_config.use_prefill_decode_attention is False

404 assert engine_args.attention_config.flash_attn_max_num_splits_for_cuda_graph == 8

414 - assert engine_args.attention_config.use_cudnn_prefill is False

415 - assert engine_args.attention_config.use_trtllm_ragged_deepseek_prefill is False

405 assert engine_args.attention_config.use_trtllm_attention is False

417 - assert engine_args.attention_config.disable_flashinfer_prefill is False

406 assert engine_args.attention_config.disable_flashinfer_q_quantization is False

407

408 # test --attention-backend flows into VllmConfig.attention_config

tests/v1/attention/test_mla_prefill_selector.py

+6 -21

@@ -269,36 +269,21 @@ class TestMLAPrefillBackendParsing:

269 )

270

271

272 -class TestDeprecatedFlagMigration:

273 - """Tests for _migrate_deprecated_mla_prefill_flags in AttentionConfig."""

272 +class TestMLAPrefillBackendConfig:

273 + """Tests for mla_prefill_backend configuration in AttentionConfig."""

274

275 - def test_no_deprecated_flags_leaves_backend_none(self):

275 + def test_default_backend_is_none(self):

276 config = AttentionConfig()

277 assert config.mla_prefill_backend is None

278

279 - def test_use_trtllm_ragged_migrates_to_trtllm_ragged(self):

280 - config = AttentionConfig(use_trtllm_ragged_deepseek_prefill=True)

281 - assert config.mla_prefill_backend == MLAPrefillBackendEnum.TRTLLM_RAGGED

282 -

283 - def test_disable_flashinfer_prefill_migrates_to_flash_attn(self):

284 - config = AttentionConfig(disable_flashinfer_prefill=True)

285 - assert config.mla_prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN

286 -

287 - def test_explicit_backend_ignores_deprecated_flags(self):

279 + def test_explicit_flash_attn_backend(self):

280 config = AttentionConfig(

281 mla_prefill_backend=MLAPrefillBackendEnum.FLASH_ATTN,

290 - use_cudnn_prefill=True,

282 )

283 assert config.mla_prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN

284

294 - def test_cudnn_raises_error(self):

295 - match = "cuDNN MLA prefill backend has been removed"

296 - with pytest.raises(ValueError, match=match):

297 - AttentionConfig(use_cudnn_prefill=True)

298 -

299 - def test_trtllm_takes_priority_over_disable_flashinfer(self):

285 + def test_explicit_trtllm_ragged_backend(self):

286 config = AttentionConfig(

301 - use_trtllm_ragged_deepseek_prefill=True,

302 - disable_flashinfer_prefill=True,

287 + mla_prefill_backend=MLAPrefillBackendEnum.TRTLLM_RAGGED,

288 )

289 assert config.mla_prefill_backend == MLAPrefillBackendEnum.TRTLLM_RAGGED

vllm/config/attention.py

+1 -52

@@ -6,12 +6,9 @@ from typing import Any, Literal

6 from pydantic import field_validator

8 from vllm.config.utils import config

9 -from vllm.logger import init_logger

9 from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum

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

13 -logger = init_logger(__name__)

14 -

13 @config

14 class AttentionConfig:

@@ -36,27 +33,16 @@ class AttentionConfig:

33 Fixes the split count so grid dimensions are constant across captures,

34 and buffers can be pre-allocated to avoid inflating the memory estimate."""

39 - use_cudnn_prefill: bool = False

40 - """Deprecated: cuDNN prefill backend has been removed."""

41 -

42 - use_trtllm_ragged_deepseek_prefill: bool = False

43 - """Whether to use TRTLLM ragged deepseek prefill."""

44 -

36 use_trtllm_attention: bool | None = None

37 """If set to True/False, use or don't use the TRTLLM attention backend

38 in flashinfer. If None, auto-detect the attention backend in flashinfer."""

49 - disable_flashinfer_prefill: bool | None = None

50 - """Whether to disable flashinfer prefill."""

51 -

40 disable_flashinfer_q_quantization: bool = False

41 """If set, when using fp8 kv, do not quantize Q to fp8."""

43 mla_prefill_backend: MLAPrefillBackendEnum | None = None

44 """MLA prefill backend to use. If None, will be selected automatically.

57 - Valid options: FLASH_ATTN (FA3/FA4), FLASHINFER, TRTLLM_RAGGED.

58 - This option supersedes use_trtllm_ragged_deepseek_prefill

59 - and disable_flashinfer_prefill which are deprecated."""

45 + Valid options: FLASH_ATTN (FA3/FA4), FLASHINFER, TRTLLM_RAGGED."""

47 use_prefill_query_quantization: bool = False

48 """If set, quantize query for attention in prefill."""

@@ -123,40 +109,3 @@ class AttentionConfig:

109 if isinstance(value, str):

110 return MLAPrefillBackendEnum[value.upper()]

111 return value

126 -

127 - def __post_init__(self) -> None:

128 - self._migrate_deprecated_mla_prefill_flags()

129 -

130 - def _migrate_deprecated_mla_prefill_flags(self) -> None:

131 - """Migrate deprecated MLA prefill flags to mla_prefill_backend."""

132 - # If the new option is already set, it takes precedence

133 - if self.mla_prefill_backend is not None:

134 - return

135 -

136 - # Check for deprecated flags and migrate them.

137 - # Only the first flag encountered sets the backend.

138 - if self.use_cudnn_prefill:

139 - raise ValueError(

140 - "The cuDNN MLA prefill backend has been removed. "

141 - "Use --attention-config.mla_prefill_backend=FLASH_ATTN or "

142 - "FLASHINFER or TRTLLM_RAGGED instead."

143 - )

144 -

145 - if self.use_trtllm_ragged_deepseek_prefill:

146 - if self.mla_prefill_backend is None:

147 - self.mla_prefill_backend = MLAPrefillBackendEnum.TRTLLM_RAGGED

148 - logger.warning_once(

149 - "use_trtllm_ragged_deepseek_prefill is deprecated and "

150 - "will be removed in v0.22. Use "

151 - "--attention-config.mla_prefill_backend=TRTLLM_RAGGED "

152 - "instead."

153 - )

154 -

155 - if self.disable_flashinfer_prefill:

156 - if self.mla_prefill_backend is None:

157 - self.mla_prefill_backend = MLAPrefillBackendEnum.FLASH_ATTN

158 - logger.warning_once(

159 - "disable_flashinfer_prefill is deprecated and will be removed "

160 - "in v0.22. Use --attention-config.mla_prefill_backend="

161 - "FLASH_ATTN instead."

162 - )

vllm/model_executor/layers/attention/mla_attention.py

+2 -2

@@ -779,14 +779,14 @@ class MLAAttention(nn.Module, AttentionLayerBase):

779 attn_out,

780 lse,

781 get_dcp_group(),

782 - is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),

782 + is_lse_base_on_e=True,

783 )

784 else:

785 attn_out = cp_lse_ag_out_rs(

786 attn_out,

787 lse,

788 get_dcp_group(),

789 - is_lse_base_on_e=not getattr(self, "_use_fi_prefill", False),

789 + is_lse_base_on_e=True,

790 )

791

792 # v_up projection