@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5b115bb8a3

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 2, 13:30 UTC

[Attention][AMD] Standardize kv layout to blocks first for AMD (#43660)

Signed-off-by: NickLucche <nlucches@redhat.com>

Nicolò Lucchesi committed May 28, 2026 at 19:28 UTC 5b115bb8a33d72820075450ecefcd292607bfe57

6 files changed +57 -44

vllm/platforms/rocm.py

+7 -3

@@ -362,6 +362,7 @@ def flash_attn_triton_available() -> bool:

362 def _get_backend_priorities(

363 use_mla: bool,

364 use_sparse: bool,

365 + use_kv_connector: bool = False,

366 ) -> list[AttentionBackendEnum]:

367 from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops

368

@@ -380,9 +381,11 @@ def _get_backend_priorities(

381 AttentionBackendEnum.TRITON_MLA,

382 ]

383

383 - backends = [

384 - AttentionBackendEnum.ROCM_ATTN,

385 - ]

384 + backends = []

385 + # ROCM_ATTN uses (2, num_blocks, ...) KV cache layout which is

386 + # incompatible with KV connectors that require blocks-first layout.

387 + if not use_kv_connector:

388 + backends.append(AttentionBackendEnum.ROCM_ATTN)

389 if rocm_aiter_ops.is_mha_enabled():

390 backends.append(AttentionBackendEnum.ROCM_AITER_FA)

391 if is_aiter_found_and_supported():

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

464 backend_priorities = _get_backend_priorities(

465 attn_selector_config.use_mla,

466 attn_selector_config.use_sparse,

467 + attn_selector_config.use_kv_connector,

468 )

469 for priority, backend in enumerate(backend_priorities):

470 try:

vllm/v1/attention/backend.py

@@ -240,6 +240,10 @@ class AttentionBackend(ABC):

240 def supports_batch_invariance(cls) -> bool:

241 return False

242

243 + @classmethod

244 + def supports_kv_connector(cls) -> bool:

245 + return True

246 +

247 @classmethod

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

249 """Check if backend supports a given attention type.

@@ -283,6 +287,7 @@ class AttentionBackend(ABC):

287 attn_type: str,

288 use_non_causal: bool = False,

289 use_batch_invariant: bool = False,

290 + use_kv_connector: bool = False,

291 ) -> list[str]:

292 invalid_reasons = []

293 if not cls.supports_head_size(head_size):

@@ -319,6 +324,8 @@ class AttentionBackend(ABC):

324 invalid_reasons.append("non-causal attention not supported")

325 if use_batch_invariant and not cls.supports_batch_invariance():

326 invalid_reasons.append("batch invariance not supported")

327 + if use_kv_connector and not cls.supports_kv_connector():

328 + invalid_reasons.append("KV connector not supported")

329 combination_reason = cls.supports_combination(

330 head_size,

331 dtype,

vllm/v1/attention/backends/rocm_aiter_fa.py

+23 -35

@@ -225,6 +225,8 @@ if current_platform.is_rocm():

225 x,

226 k_stride0,

227 v_stride0,

228 + k_cache_block_stride,

229 + v_cache_block_stride,

230 block_size,

231 head_size,

232 num_kv_heads,

@@ -241,15 +243,17 @@ if current_platform.is_rocm():

243 return

244 block_id = slot_id // block_size

245 block_offset = slot_id % block_size

244 - dst_offset = (

245 - block_id * num_kv_heads * head_size * block_size

246 - + head_id * head_size * block_size

246 + k_dst_offset = (

247 + block_id * k_cache_block_stride + head_id * head_size * block_size

248 )

249 dst_k_shuffle_offset = (

249 - dst_offset + offset // x * block_size * x + block_offset * x + offset % x

250 + k_dst_offset + offset // x * block_size * x + block_offset * x + offset % x

251 + )

252 + v_dst_offset = (

253 + block_id * v_cache_block_stride + head_id * head_size * block_size

254 )

255 dst_v_shuffle_offset = (

252 - dst_offset

256 + v_dst_offset

257 + block_offset // x * head_size * x

258 + offset * x

259 + block_offset % x

@@ -280,18 +284,6 @@ if current_platform.is_rocm():

284 _, num_kv_heads, head_size = key.shape

285 num_blocks, block_size, _, _ = key_cache.shape

286 x = 16 // key_cache.element_size()

283 - k_cache_template = torch.empty(

284 - [num_blocks, num_kv_heads, head_size // x, block_size, x],

285 - dtype=key_cache.dtype,

286 - device="meta",

287 - )

288 - v_cache_template = torch.empty(

289 - [num_blocks, num_kv_heads, block_size // x, head_size, x],

290 - dtype=value_cache.dtype,

291 - device="meta",

292 - )

293 - new_key_cache = key_cache.view_as(k_cache_template)

294 - new_value_cache = value_cache.view_as(v_cache_template)

287 QUANT = False

288 if is_quantized_kv_cache(kv_cache_dtype):

289 QUANT = True

@@ -302,14 +294,16 @@ if current_platform.is_rocm():

294 reshape_and_cache_shuffle_kernel[grid](

295 key,

296 value,

305 - new_key_cache,

306 - new_value_cache,

297 + key_cache,

298 + value_cache,

299 slot_mapping,

300 k_scales,

301 v_scales,

302 x,

303 key.stride(0),

304 value.stride(0),

305 + key_cache.stride(0),

306 + value_cache.stride(0),

307 block_size,

308 head_size,

309 num_kv_heads,

@@ -485,7 +479,7 @@ class AiterFlashAttentionMetadataBuilder(

479 kv_cache_shape = self.vllm_config.compilation_config.static_forward_context[

480 first_layer_name

481 ].kv_cache.shape

488 - num_blocks = kv_cache_shape[1]

482 + num_blocks = kv_cache_shape[0]

483 self.scale = torch.ones(

484 [num_blocks, self.num_heads_kv, self.block_size],

485 dtype=torch.float32,

@@ -765,7 +759,7 @@ class AiterFlashAttentionBackend(AttentionBackend):

759 ) -> tuple[int, ...]:

760 if block_size % 16 != 0:

761 raise ValueError("Block size must be a multiple of 16.")

768 - return (2, num_blocks, block_size, num_kv_heads, head_size)

762 + return (num_blocks, 2, block_size, num_kv_heads, head_size)

763

764 @classmethod

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

@@ -1025,7 +1019,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

1019 key: shape = [num_tokens, num_kv_heads, head_size]

1020 value: shape = [num_tokens, num_kv_heads, head_size]

1021 kv_cache: shape =

1028 - [2, num_blocks, block_size, num_kv_heads, head_size]

1022 + [num_blocks, 2, block_size, num_kv_heads, head_size]

1023 attn_metadata: Metadata for attention.

1024 Returns:

1025 shape = [num_tokens, num_heads * head_size]

@@ -1053,7 +1047,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

1047 # Whenever making a change in this method, please benchmark the

1048 # performance to make sure it does not introduce any overhead.

1049 num_actual_tokens = attn_metadata.num_actual_tokens

1056 - key_cache, value_cache = kv_cache.unbind(0)

1050 + key_cache, value_cache = kv_cache.unbind(1)

1051

1052 if is_quantized_kv_cache(self.kv_cache_dtype):

1053 key_cache = key_cache.view(current_platform.fp8_dtype())

@@ -1278,18 +1272,12 @@ class AiterFlashAttentionImpl(AttentionImpl):

1272 max_logits = torch.empty_like(exp_sums)

1273 num_blocks, block_size, num_kv_heads, _ = key_cache.shape

1274 x = 16 // key_cache.element_size()

1281 - k_cache_template = torch.empty(

1282 - [num_blocks, num_kv_heads, head_size // x, block_size, x],

1283 - dtype=key_cache.dtype,

1284 - device="meta",

1275 + new_key_cache = key_cache.reshape(

1276 + num_blocks, num_kv_heads, head_size // x, block_size, x

1277 )

1286 - v_cache_template = torch.empty(

1287 - [num_blocks, num_kv_heads, block_size // x, head_size, x],

1288 - dtype=value_cache.dtype,

1289 - device="meta",

1278 + new_value_cache = value_cache.reshape(

1279 + num_blocks, num_kv_heads, block_size // x, head_size, x

1280 )

1291 - new_key_cache = key_cache.view_as(k_cache_template)

1292 - new_value_cache = value_cache.view_as(v_cache_template)

1281 k_qscale = (

1282 layer._k_scale

1283 if attn_metadata.k_scale is None

@@ -1378,7 +1366,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

1366 kv_cache: torch.Tensor,

1367 slot_mapping: torch.Tensor,

1368 ):

1381 - key_cache, value_cache = kv_cache.unbind(0)

1369 + key_cache, value_cache = kv_cache.unbind(1)

1370

1371 # key and value may be None in the case of cross attention. They are

1372 # calculated once based on the output from the encoder and then cached

@@ -1446,7 +1434,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

1434 kv_cache: torch.Tensor,

1435 layer_slot_mapping: torch.Tensor,

1436 ):

1449 - key_cache, value_cache = kv_cache.unbind(0)

1437 + key_cache, value_cache = kv_cache.unbind(1)

1438 flash_layout = True

1439

1440 is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

+5 -5

@@ -74,7 +74,7 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):

74 ) -> tuple[int, ...]:

75 if block_size % 16 != 0:

76 raise ValueError("Block size must be a multiple of 16.")

77 - return (2, num_blocks, block_size, num_kv_heads, head_size)

77 + return (num_blocks, 2, block_size, num_kv_heads, head_size)

79 @staticmethod

80 def use_cascade_attention(*args, **kwargs) -> bool:

@@ -153,7 +153,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

153 key: shape = [num_tokens, num_kv_heads, head_size]

154 value: shape = [num_tokens, num_kv_heads, head_size]

155 kv_cache: shape =

156 - [2, num_blocks, block_size, num_kv_heads, head_size]

156 + [num_blocks, 2, block_size, num_kv_heads, head_size]

157 attn_metadata: Metadata for attention.

158 Returns:

159 shape = [num_tokens, num_heads * head_size]

@@ -194,7 +194,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

194 layer,

195 )

196

197 - key_cache, value_cache = kv_cache.unbind(0)

197 + key_cache, value_cache = kv_cache.unbind(1)

198

199 softmax_scale = self.scale

200 if is_quantized_kv_cache(self.kv_cache_dtype):

@@ -243,7 +243,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

243 # For encoder attention,

244 # we use direct Q, K, V tensors without caching

245 return

246 - key_cache, value_cache = kv_cache.unbind(0)

246 + key_cache, value_cache = kv_cache.unbind(1)

247

248 # Reshape the input keys and values and store them in the cache.

249 ops.reshape_and_cache_flash(

@@ -276,7 +276,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):

276 # For encoder attention,

277 # we use direct Q, K, V tensors without caching

278 return

279 - key_cache, value_cache = kv_cache.unbind(0)

279 + key_cache, value_cache = kv_cache.unbind(1)

280 flash_layout = True

281

282 is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)

vllm/v1/attention/backends/rocm_attn.py

@@ -208,6 +208,12 @@ class RocmAttentionBackend(AttentionBackend):

208 def supports_non_causal(cls) -> bool:

209 return True

210

211 + @classmethod

212 + def supports_kv_connector(cls) -> bool:

213 + # ROCM_ATTN uses (2, num_blocks, ...) KV cache layout which is

214 + # incompatible with KV connectors that require blocks-first layout.

215 + return False

216 +

217 forward_includes_kv_cache_update: bool = False

218

219 @staticmethod

vllm/v1/attention/selector.py

+9 -1

@@ -31,6 +31,7 @@ class AttentionSelectorConfig(NamedTuple):

31 attn_type: str = AttentionType.DECODER

32 use_non_causal: bool = False

33 use_batch_invariant: bool = False

34 + use_kv_connector: bool = False

36 def __repr__(self):

37 return (

@@ -45,7 +46,8 @@ class AttentionSelectorConfig(NamedTuple):

46 f"use_per_head_quant_scales={self.use_per_head_quant_scales}, "

47 f"attn_type={self.attn_type}, "

48 f"use_non_causal={self.use_non_causal}, "

48 - f"use_batch_invariant={self.use_batch_invariant})"

49 + f"use_batch_invariant={self.use_batch_invariant}, "

50 + f"use_kv_connector={self.use_kv_connector})"

51 )

@@ -80,6 +82,11 @@ def get_attn_backend(

82 else:

83 block_size = None

85 + kv_transfer_config = vllm_config.kv_transfer_config

86 + use_kv_connector = (

87 + kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance

88 + )

89 +

90 attn_selector_config = AttentionSelectorConfig(

91 head_size=head_size,

92 dtype=dtype,

@@ -93,6 +100,7 @@ def get_attn_backend(

100 attn_type=attn_type or AttentionType.DECODER,

101 use_non_causal=vllm_config.attention_config.use_non_causal,

102 use_batch_invariant=envs.VLLM_BATCH_INVARIANT,

103 + use_kv_connector=use_kv_connector,

104 )

105

106 return _cached_get_attn_backend(