066209a045 · tduncan/Nvidia-Optimization-Stack-VLLM

tests/v1/attention/test_mla_backends.py

+1 -1

@@ -61,7 +61,7 @@ for backend in BACKENDS_TO_TEST:

61

62 BACKEND_BLOCK_SIZES = {}

63 for backend in BACKENDS_TO_TEST:

64 - supported_sizes = backend.get_class().supported_kernel_block_sizes

64 + supported_sizes = backend.get_class().get_supported_kernel_block_sizes()

65 if supported_sizes:

66 default_size = supported_sizes[0]

67 block_size = (

tests/v1/worker/test_gpu_model_runner.py

+3 -1

@@ -185,7 +185,9 @@ def _make_mock_backend_for_kernel_block_size(

185 supported_sizes: list[int | MultipleOf],

186 ):

187 class _MockBackend:

188 - supported_kernel_block_sizes = supported_sizes

188 + @staticmethod

189 + def get_supported_kernel_block_sizes():

190 + return supported_sizes

191

192 return _MockBackend()

193

vllm/attention/backends/abstract.py

+7 -3

@@ -46,9 +46,12 @@ class AttentionBackend(ABC):

46 # makes sure the output tensor is allocated inside the cudagraph.

47 accept_output_buffer: bool = False

48 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

49 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(1)]

49 supported_kv_cache_dtypes: ClassVar[list["CacheDType"]] = ["auto"]

50

51 + @staticmethod

52 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

53 + return [MultipleOf(1)]

54 +

55 @staticmethod

56 @abstractmethod

57 def get_name() -> str:

@@ -142,10 +145,11 @@ class AttentionBackend(ABC):

145 if block_size not in valid_sizes:

146 return False

147

145 - if not cls.supported_kernel_block_sizes:

148 + supported_kernel_block_sizes = cls.get_supported_kernel_block_sizes()

149 + if not supported_kernel_block_sizes:

150 return True

151

148 - for supported_size in cls.supported_kernel_block_sizes:

152 + for supported_size in supported_kernel_block_sizes:

153 if isinstance(supported_size, MultipleOf):

154 supported_size = supported_size.base

155 # With hybrid_blocks feature, the framework-level block size

vllm/v1/attention/backends/flash_attn.py

+21 -6

@@ -32,7 +32,7 @@ if is_flash_attn_varlen_func_available():

32 get_scheduler_metadata,

33 reshape_and_cache_flash,

34 )

35 -from vllm.config import VllmConfig, get_layers_from_vllm_config

35 +from vllm.config import VllmConfig, get_current_vllm_config, get_layers_from_vllm_config

36 from vllm.config.cache import CacheDType

37 from vllm.distributed.parallel_state import get_dcp_group

38 from vllm.logger import init_logger

@@ -56,11 +56,26 @@ logger = init_logger(__name__)

56 class FlashAttentionBackend(AttentionBackend):

57 accept_output_buffer: bool = True

58 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

59 - # NOTE(tdoublep): while in principle, FA supports

60 - # MultipleOf(16), these are the block sizes that do not

61 - # suffer from the NaN propagation problem described here:

62 - # https://github.com/Dao-AILab/flash-attention/issues/1974

63 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [16, 32, 64]

59 +

60 + @staticmethod

61 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

62 + vllm_config = get_current_vllm_config()

63 + model_config = vllm_config.model_config

64 + cache_config = vllm_config.cache_config

65 + if (

66 + model_config

67 + and model_config.is_hybrid

68 + and (

69 + cache_config.mamba_ssm_cache_dtype == "float32"

70 + or cache_config.mamba_cache_dtype == "float32"

71 + )

72 + ):

73 + # NOTE(tdoublep): while in principle, FA supports

74 + # MultipleOf(16), these are the block sizes that do not

75 + # suffer from the NaN propagation problem described here:

76 + # https://github.com/Dao-AILab/flash-attention/issues/1974

77 + return [16, 32, 64]

78 + return [MultipleOf(16)]

79

80 @staticmethod

81 def get_name() -> str:

vllm/v1/attention/backends/flashinfer.py

+6 -6

@@ -16,7 +16,6 @@ from flashinfer import (

16 from flashinfer.decode import _get_range_buf, trtllm_batch_decode_with_kv_cache

17 from flashinfer.prefill import trtllm_batch_context_with_kv_cache

18 from flashinfer.utils import FP4Tensor

19 -from typing_extensions import override

19

20 from vllm import envs

21 from vllm.attention.backends.abstract import (

@@ -275,10 +274,6 @@ class BatchDCPPrefillWrapper:

274 class FlashInferBackend(AttentionBackend):

275 accept_output_buffer: bool = True

276 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

278 - # Note: Not sure for all platforms,

279 - # but on Blackwell, only support a page size of

280 - # 16, 32, 64

281 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [16, 32, 64]

277 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

278 "auto",

279 "fp8",

@@ -286,6 +281,12 @@ class FlashInferBackend(AttentionBackend):

281 "fp8_e5m2",

282 ]

283

284 + @staticmethod

285 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

286 + # Note: Not sure for all platforms, but on Blackwell,

287 + # only support a page size of 16, 32, 64.

288 + return [16, 32, 64]

289 +

290 @staticmethod

291 def get_name() -> str:

292 return "FLASHINFER"

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

567 )

568

569 @classmethod

569 - @override

570 def get_cudagraph_support(

571 cls: type["FlashInferMetadataBuilder"],

572 vllm_config: VllmConfig,

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

+4 -1

@@ -36,13 +36,16 @@ class CutlassMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):

36

37 class CutlassMLABackend(MLACommonBackend):

38 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

39 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [128]

39 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

40 "auto",

41 "fp8",

42 "fp8_e4m3",

43 ]

44

45 + @staticmethod

46 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

47 + return [128]

48 +

49 @staticmethod

50 def get_name() -> str:

51 return "CUTLASS_MLA"

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

+4 -1

@@ -41,9 +41,12 @@ logger = init_logger(__name__)

41

42 class FlashAttnMLABackend(MLACommonBackend):

43 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

44 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]

44 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto"]

45

46 + @staticmethod

47 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

48 + return [MultipleOf(16)]

49 +

50 @staticmethod

51 def get_name() -> str:

52 return "FLASH_ATTN_MLA"

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

+4 -1

@@ -35,13 +35,16 @@ class FlashInferMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):

35

36 class FlashInferMLABackend(MLACommonBackend):

37 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

38 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [32, 64]

38 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

39 "auto",

40 "fp8",

41 "fp8_e4m3",

42 ]

43

44 + @staticmethod

45 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

46 + return [32, 64]

47 +

48 @staticmethod

49 def get_name() -> str:

50 return "FLASHINFER_MLA"

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

+4 -1

@@ -39,13 +39,16 @@ logger = init_logger(__name__)

39

40 class FlashMLABackend(MLACommonBackend):

41 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

42 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [64]

42 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

43 "auto",

44 "fp8",

45 "fp8_e4m3",

46 ]

47

48 + @staticmethod

49 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

50 + return [64]

51 +

52 @staticmethod

53 def get_name() -> str:

54 return "FLASHMLA"

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

+4 -1

@@ -55,9 +55,12 @@ structured as:

55 class FlashMLASparseBackend(AttentionBackend):

56 accept_output_buffer: bool = True

57 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]

58 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [64]

58 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = ["auto", "fp8_ds_mla"]

59

60 + @staticmethod

61 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

62 + return [64]

63 +

64 @staticmethod

65 def get_name() -> str:

66 return "FLASHMLA_SPARSE"

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

+3 -3

@@ -24,9 +24,9 @@ logger = init_logger(__name__)

24

25

26 class DeepseekV32IndexerBackend(AttentionBackend):

27 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [

28 - 1 if current_platform.is_rocm() else 64

29 - ]

27 + @staticmethod

28 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

29 + return [1 if current_platform.is_rocm() else 64]

30

31 @classmethod

32 def get_supported_head_sizes(cls) -> list[int]:

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

+3 -1

@@ -21,7 +21,9 @@ from vllm.v1.kv_cache_interface import AttentionSpec

21

22

23 class AiterMLABackend(MLACommonBackend):

24 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [1]

24 + @staticmethod

25 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

26 + return [1]

27

28 @staticmethod

29 def get_name() -> str:

vllm/v1/attention/backends/rocm_aiter_fa.py

+4 -1

@@ -447,7 +447,10 @@ class AiterFlashAttentionMetadataBuilder(

447 class AiterFlashAttentionBackend(AttentionBackend):

448 accept_output_buffer: bool = True

449 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

450 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]

450 +

451 + @staticmethod

452 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

453 + return [MultipleOf(16)]

454

455 @classmethod

456 def get_supported_head_sizes(cls) -> list[int]:

vllm/v1/attention/backends/tree_attn.py

+4 -1

@@ -31,7 +31,10 @@ logger = init_logger(__name__)

31 class TreeAttentionBackend(AttentionBackend):

32 accept_output_buffer: bool = True

33 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

34 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]

34 +

35 + @staticmethod

36 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

37 + return [MultipleOf(16)]

38

39 @classmethod

40 def get_supported_head_sizes(cls) -> list[int]:

vllm/v1/attention/backends/triton_attn.py

+4 -1

@@ -154,7 +154,6 @@ class TritonAttentionBackend(AttentionBackend):

154 torch.bfloat16,

155 torch.float32,

156 ]

157 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]

157 supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

158 "auto",

159 "fp8",

@@ -162,6 +161,10 @@ class TritonAttentionBackend(AttentionBackend):

161 "fp8_e5m2",

162 ]

163

164 + @staticmethod

165 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

166 + return [MultipleOf(16)]

167 +

168 @staticmethod

169 def get_name() -> str:

170 return "TRITON_ATTN"

vllm/v1/attention/backends/xformers.py

+4 -1

@@ -42,7 +42,10 @@ logger = init_logger(__name__)

42 class XFormersAttentionBackend(AttentionBackend):

43 accept_output_buffer: bool = True

44 supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

45 - supported_kernel_block_sizes: ClassVar[list[int | MultipleOf]] = [MultipleOf(16)]

45 +

46 + @staticmethod

47 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

48 + return [MultipleOf(16)]

49

50 @classmethod

51 def get_supported_head_sizes(cls) -> list[int]:

vllm/v1/worker/gpu_model_runner.py

+2 -2

@@ -4618,7 +4618,7 @@ class GPUModelRunner(

4618 """

4619 for backend in backends:

4620 is_supported = False

4621 - for supported_size in backend.supported_kernel_block_sizes:

4621 + for supported_size in backend.get_supported_kernel_block_sizes():

4622 if isinstance(supported_size, int):

4623 if block_size == supported_size:

4624 is_supported = True

@@ -4649,7 +4649,7 @@ class GPUModelRunner(

4649 all_int_supported_sizes = set(

4650 supported_size

4651 for backend in backends

4652 - for supported_size in backend.supported_kernel_block_sizes

4652 + for supported_size in backend.get_supported_kernel_block_sizes()

4653 if isinstance(supported_size, int)

4654 )

4655

[Attention] Refactor FA `block_size` limitations to hybrid models only (#29084)