07a606aa7e · tduncan/Nvidia-Optimization-Stack-VLLM

vllm/attention/backends/abstract.py

+14

@@ -142,6 +142,17 @@ class AttentionBackend(ABC):

142 def is_sparse(cls) -> bool:

143 return False

144

145 + @classmethod

146 + def supports_attn_type(cls, attn_type: str) -> bool:

147 + """Check if backend supports a given attention type.

148 +

149 + By default, only supports decoder attention.

150 + Backends should override this to support other attention types.

151 + """

152 + from vllm.attention import AttentionType

153 +

154 + return attn_type == AttentionType.DECODER

155 +

156 @classmethod

157 def supports_compute_capability(cls, capability: "DeviceCapability") -> bool:

158 return True

@@ -171,6 +182,7 @@ class AttentionBackend(ABC):

182 has_sink: bool,

183 use_sparse: bool,

184 device_capability: "DeviceCapability",

185 + attn_type: str,

186 ) -> list[str]:

187 invalid_reasons = []

188 if not cls.supports_head_size(head_size):

@@ -195,6 +207,8 @@ class AttentionBackend(ABC):

207 invalid_reasons.append("non-sparse not supported")

208 if not cls.supports_compute_capability(device_capability):

209 invalid_reasons.append("compute capability not supported")

210 + if not cls.supports_attn_type(attn_type):

211 + invalid_reasons.append(f"attention type {attn_type} not supported")

212 combination_reason = cls.supports_combination(

213 head_size,

214 dtype,

vllm/attention/layer.py

+1

@@ -291,6 +291,7 @@ class Attention(nn.Module, AttentionLayerBase):

291 block_size,

292 use_mla=False,

293 has_sink=self.has_sink,

294 + attn_type=attn_type,

295 )

296 else:

297 self.attn_backend = attn_backend

vllm/attention/layers/encoder_only_attention.py

+5 -1

@@ -74,7 +74,11 @@ class EncoderOnlyAttention(Attention):

74 block_size = 16

75

76 underlying_attn_backend = get_attn_backend(

77 - head_size, dtype, kv_cache_dtype, block_size

77 + head_size,

78 + dtype,

79 + kv_cache_dtype,

80 + block_size,

81 + attn_type=AttentionType.ENCODER_ONLY,

82 )

83

84 attn_backend = create_encoder_only_attention_backend(underlying_attn_backend)

vllm/attention/selector.py

+5

@@ -76,6 +76,7 @@ def get_attn_backend(

76 use_mla: bool = False,

77 has_sink: bool = False,

78 use_sparse: bool = False,

79 + attn_type: str | None = None,

80 ) -> type[AttentionBackend]:

81 """Selects which attention backend to use and lazily imports it."""

82

@@ -94,6 +95,7 @@ def get_attn_backend(

95 use_mla=use_mla,

96 has_sink=has_sink,

97 use_sparse=use_sparse,

98 + attn_type=attn_type,

99 )

100

101

@@ -106,6 +108,7 @@ def _cached_get_attn_backend(

108 use_mla: bool = False,

109 has_sink: bool = False,

110 use_sparse: bool = False,

111 + attn_type: str | None = None,

112 ) -> type[AttentionBackend]:

113 # Check whether a particular choice of backend was

114 # previously forced.

@@ -159,6 +162,7 @@ def _cached_get_attn_backend(

162 use_mla,

163 has_sink,

164 use_sparse,

165 + attn_type,

166 )

167 else:

168 attention_cls = current_platform.get_attn_backend_cls(

@@ -170,6 +174,7 @@ def _cached_get_attn_backend(

174 use_mla,

175 has_sink,

176 use_sparse,

177 + attn_type,

178 )

179 if not attention_cls:

180 raise ValueError(

vllm/platforms/cpu.py

+1

@@ -134,6 +134,7 @@ class CpuPlatform(Platform):

134 use_mla: bool,

135 has_sink: bool,

136 use_sparse: bool,

137 + attn_type: str | None = None,

138 ) -> str:

139 from vllm.attention.backends.registry import AttentionBackendEnum

140

vllm/platforms/cuda.py

+10

@@ -298,6 +298,7 @@ class CudaPlatformBase(Platform):

298 has_sink,

299 use_sparse,

300 device_capability,

301 + attn_type,

302 ) -> tuple[

303 list[tuple["AttentionBackendEnum", int]],

304 dict["AttentionBackendEnum", list[str]],

@@ -318,6 +319,7 @@ class CudaPlatformBase(Platform):

319 has_sink,

320 use_sparse,

321 device_capability,

322 + attn_type,

323 )

324 except ImportError:

325 invalid_reasons_i = ["ImportError"]

@@ -339,7 +341,13 @@ class CudaPlatformBase(Platform):

341 use_mla: bool,

342 has_sink: bool,

343 use_sparse: bool,

344 + attn_type: str | None = None,

345 ) -> str:

346 + from vllm.attention import AttentionType

347 +

348 + if attn_type is None:

349 + attn_type = AttentionType.DECODER

350 +

351 device_capability = cls.get_device_capability()

352 assert device_capability is not None

353

@@ -356,6 +364,7 @@ class CudaPlatformBase(Platform):

364 has_sink,

365 use_sparse,

366 device_capability,

367 + attn_type,

368 )

369 except ImportError:

370 invalid_reasons = ["ImportError"]

@@ -379,6 +388,7 @@ class CudaPlatformBase(Platform):

388 has_sink,

389 use_sparse,

390 device_capability,

391 + attn_type,

392 )

393 reasons_str = (

394 "{"

vllm/platforms/interface.py

+1

@@ -222,6 +222,7 @@ class Platform:

222 use_mla: bool,

223 has_sink: bool,

224 use_sparse: bool,

225 + attn_type: str | None = None,

226 ) -> str:

227 """Get the attention backend class of a device."""

228 return ""

vllm/platforms/rocm.py

+1

@@ -216,6 +216,7 @@ class RocmPlatform(Platform):

216 use_mla,

217 has_sink,

218 use_sparse,

219 + attn_type: str | None = None,

220 ) -> str:

221 from vllm._aiter_ops import rocm_aiter_ops

222 from vllm.attention.backends.registry import AttentionBackendEnum

vllm/platforms/tpu.py

+1

@@ -61,6 +61,7 @@ class TpuPlatform(Platform):

61 use_mla: bool,

62 has_sink,

63 use_sparse,

64 + attn_type: str | None = None,

65 ) -> str:

66 from vllm.attention.backends.registry import AttentionBackendEnum

67

vllm/platforms/xpu.py

+1

@@ -51,6 +51,7 @@ class XPUPlatform(Platform):

51 use_mla: bool,

52 has_sink: bool,

53 use_sparse,

54 + attn_type: str | None = None,

55 ) -> str:

56 from vllm.v1.attention.backends.utils import set_kv_cache_layout

57

vllm/v1/attention/backends/cpu_attn.py

+11

@@ -48,6 +48,17 @@ class CPUAttentionBackend(AttentionBackend):

48 def get_name() -> str:

49 return "CPU_ATTN"

50

51 + @classmethod

52 + def supports_attn_type(cls, attn_type: str) -> bool:

53 + """CPU attention supports decoder and encoder-only attention."""

54 + from vllm.attention import AttentionType

55 +

56 + return attn_type in (

57 + AttentionType.DECODER,

58 + AttentionType.ENCODER,

59 + AttentionType.ENCODER_ONLY,

60 + )

61 +

62 @staticmethod

63 def get_impl_cls() -> type["CPUAttentionBackendImpl"]:

64 return CPUAttentionBackendImpl

vllm/v1/attention/backends/flash_attn.py

+12

@@ -66,6 +66,18 @@ class FlashAttentionBackend(AttentionBackend):

66 def get_name() -> str:

67 return "FLASH_ATTN"

68

69 + @classmethod

70 + def supports_attn_type(cls, attn_type: str) -> bool:

71 + """FlashAttention supports all attention types."""

72 + from vllm.attention import AttentionType

73 +

74 + return attn_type in (

75 + AttentionType.DECODER,

76 + AttentionType.ENCODER,

77 + AttentionType.ENCODER_ONLY,

78 + AttentionType.ENCODER_DECODER,

79 + )

80 +

81 @staticmethod

82 def get_impl_cls() -> type["FlashAttentionImpl"]:

83 return FlashAttentionImpl

vllm/v1/attention/backends/flex_attention.py

+7

@@ -84,6 +84,13 @@ class FlexAttentionBackend(AttentionBackend):

84 def get_name() -> str:

85 return "FLEX_ATTENTION"

86

87 + @classmethod

88 + def supports_attn_type(cls, attn_type: str) -> bool:

89 + """FlexAttention supports both decoder and encoder-only attention."""

90 + from vllm.attention import AttentionType

91 +

92 + return attn_type in (AttentionType.DECODER, AttentionType.ENCODER_ONLY)

93 +

94 @staticmethod

95 def get_impl_cls() -> type["FlexAttentionImpl"]:

96 return FlexAttentionImpl

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

+5 -5

@@ -40,14 +40,14 @@ logger = init_logger(__name__)

40 """

41 NOTE: FlashMLA Sparse uses an fp8 cache with the following format

42

43 -In the "FP8 with scale" format, each token's KV cache is 656 Bytes,

43 +In the "FP8 with scale" format, each token's KV cache is 656 Bytes,

44 structured as:

45 -- **First 512 bytes:** The "quantized NoPE" part, containing 512

45 +- **First 512 bytes:** The "quantized NoPE" part, containing 512

46 `float8_e4m3` values.

47 -- **Next 16 bytes:** Scale factors, containing 4 `float32` values.

48 - The first `float32` is the scale for the first 128 `float8_e4m3` values,

47 +- **Next 16 bytes:** Scale factors, containing 4 `float32` values.

48 + The first `float32` is the scale for the first 128 `float8_e4m3` values,

49 the second for the next 128, and so on.

50 -- **Last 128 bytes:** The "RoPE" part, containing 64 `bfloat16` values. This

50 +- **Last 128 bytes:** The "RoPE" part, containing 64 `bfloat16` values. This

51 part is not quantized for accuracy.

52 """

53

[CI Failure] Fix backend selection for encoder-only models (#28534)