@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 130d6c9514

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 10, 02:45 UTC

[ROCm][Perf] Enable shuffle kv cache layout and assembly paged attention kernel for `AiterFlashAttentionBackend` (#29887)

Signed-off-by: ganyi <ygan@amd.com>

Pleaplusone committed Jan 15, 2026 at 23:29 UTC 130d6c9514856cb5a152329f0382d60ff6e8d97e

3 files changed +311 -84

vllm/_aiter_ops.py

@@ -833,6 +833,7 @@ class rocm_aiter_ops:

833 _FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE

834 _MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA

835 _MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA

836 + _SHUFFLE_KV_CACHE_ENABLED = envs.VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT

837 _TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION

838 # TODO: Consolidate under _LINEAR_ENABLED

839 _FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM

@@ -859,6 +860,7 @@ class rocm_aiter_ops:

860 cls._FMOE_ENABLED = envs.VLLM_ROCM_USE_AITER_MOE

861 cls._MLA_ENABLED = envs.VLLM_ROCM_USE_AITER_MLA

862 cls._MHA_ENABLED = envs.VLLM_ROCM_USE_AITER_MHA

863 + cls._SHUFFLE_KV_CACHE_ENABLED = envs.VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT

864 cls._TRITON_UNIFIED_ATTN_ENABLED = envs.VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION

865 cls._FP8BMM_ENABLED = envs.VLLM_ROCM_USE_AITER_FP8BMM

866 cls._FP4_GEMM_DYNAMIC_QUANT_ASM = envs.VLLM_ROCM_USE_AITER_FP4_ASM_GEMM

@@ -906,6 +908,11 @@ class rocm_aiter_ops:

908 def is_mha_enabled(cls) -> bool:

909 return cls._AITER_ENABLED and cls._MHA_ENABLED

910

911 + @classmethod

912 + @if_aiter_supported

913 + def is_shuffle_kv_cache_enabled(cls) -> bool:

914 + return cls._AITER_ENABLED and cls._SHUFFLE_KV_CACHE_ENABLED

915 +

916 @classmethod

917 @if_aiter_supported

918 def is_triton_unified_attn_enabled(cls) -> bool:

vllm/envs.py

@@ -128,6 +128,7 @@ if TYPE_CHECKING:

128 VLLM_ROCM_FP8_PADDING: bool = True

129 VLLM_ROCM_MOE_PADDING: bool = True

130 VLLM_ROCM_CUSTOM_PAGED_ATTN: bool = True

131 + VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT: bool = False

132 VLLM_ENABLE_V1_MULTIPROCESSING: bool = True

133 VLLM_LOG_BATCHSIZE_INTERVAL: float = -1

134 VLLM_DISABLE_COMPILE_CACHE: bool = False

@@ -1018,6 +1019,10 @@ environment_variables: dict[str, Callable[[], Any]] = {

1019 "VLLM_ROCM_CUSTOM_PAGED_ATTN": lambda: (

1020 os.getenv("VLLM_ROCM_CUSTOM_PAGED_ATTN", "True").lower() in ("true", "1")

1021 ),

1022 + # Whether to use the shuffled kv cache layout

1023 + "VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT": lambda: (

1024 + os.getenv("VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT", "False").lower() in ("true", "1")

1025 + ),

1026 # Custom quick allreduce kernel for MI3* cards

1027 # Choice of quantization level: FP, INT8, INT6, INT4 or NONE

1028 # Recommended for large models to get allreduce

vllm/v1/attention/backends/rocm_aiter_fa.py

+299 -84

@@ -7,6 +7,7 @@ from typing import ClassVar

8 import torch

10 +from vllm._aiter_ops import rocm_aiter_ops

11 from vllm.attention.layer import Attention

12 from vllm.config import VllmConfig, get_layers_from_vllm_config

13 from vllm.logger import init_logger

@@ -30,7 +31,6 @@ from vllm.v1.kv_cache_interface import AttentionSpec

32 _PARTITION_SIZE_ROCM = 256

33 _CP_TOKENS_PER_ITER_ROCM = 32 * 1024

33 -

34 if current_platform.is_rocm():

35 import aiter

@@ -52,7 +52,7 @@ if current_platform.is_rocm():

52 cu_seqlens_kv_ptr, # [num_batches + 1]

53 token_to_batch_ptr, # [max_cum_tokens]

54 seq_start_ptr, # [num_batches]

55 - k_scale_ptr,

55 + k_scale_ptr, # [1] / [num_blocks, num_kv_heads, page_size]

56 v_scale_ptr,

57 num_heads,

58 head_size,

@@ -64,13 +64,15 @@ if current_platform.is_rocm():

64 BLOCK_SIZE: tl.constexpr,

65 ):

66 token_id = tl.program_id(0)

67 + head_id = tl.program_id(1)

68 col_offsets = tl.arange(0, BLOCK_SIZE)

68 - if DEQUANT:

69 - k_scale = tl.load(k_scale_ptr)

70 - v_scale = tl.load(v_scale_ptr)

72 - key_ptr_offset = key_ptr + token_id * head_size * num_heads

73 - value_ptr_offset = value_ptr + token_id * head_size * num_heads

70 + key_ptr_offset = (

71 + key_ptr + token_id * head_size * num_heads + head_id * head_size

72 + )

73 + value_ptr_offset = (

74 + value_ptr + token_id * head_size * num_heads + head_id * head_size

75 + )

76 batch_idx = tl.load(token_to_batch_ptr + token_id)

77 batch_start = tl.load(seq_start_ptr + batch_idx)

78 token_start = tl.load(cu_seqlens_kv_ptr + batch_idx)

@@ -89,24 +91,54 @@ if current_platform.is_rocm():

91 key_cache_ptr

92 + block_id * num_heads * head_size * PAGE_SIZE

93 + slot_id * num_heads * head_size

94 + + head_id * head_size

95 )

96 value_cache_ptr_offset = (

97 value_cache_ptr

98 + block_id * num_heads * head_size * PAGE_SIZE

99 + slot_id * num_heads * head_size

100 + + head_id * head_size

101 )

98 -

99 - for i in tl.range(0, head_size * num_heads, BLOCK_SIZE):

100 - mask = (col_offsets + i) < head_size * num_heads

101 - k_reg = tl.load(key_cache_ptr_offset + col_offsets + i, mask=mask)

102 - v_reg = tl.load(value_cache_ptr_offset + col_offsets + i, mask=mask)

103 - if DEQUANT:

104 - k_dtype = k_reg.dtype

105 - v_dtype = v_reg.dtype

106 - k_reg = (k_reg.to(tl.float32) * k_scale).to(k_dtype)

107 - v_reg = (v_reg.to(tl.float32) * v_scale).to(v_dtype)

108 - tl.store(key_ptr_offset + col_offsets + i, k_reg, mask=mask)

109 - tl.store(value_ptr_offset + col_offsets + i, v_reg, mask=mask)

102 + k_reg = tl.load(key_cache_ptr_offset + col_offsets)

103 + v_reg = tl.load(value_cache_ptr_offset + col_offsets)

104 + if DEQUANT:

105 + k_scale = tl.load(k_scale_ptr)

106 + v_scale = tl.load(v_scale_ptr)

107 + k_dtype = k_reg.dtype

108 + v_dtype = v_reg.dtype

109 + k_reg = (k_reg.to(tl.float32) * k_scale).to(k_dtype)

110 + v_reg = (v_reg.to(tl.float32) * v_scale).to(v_dtype)

111 + tl.store(key_ptr_offset + col_offsets, k_reg)

112 + tl.store(value_ptr_offset + col_offsets, v_reg)

113 +

114 + elif CACHE_FORMAT == "SHUFFLE":

115 + # for kv cache layout as

116 + # K: [num_blocks, num_head, head_dim // x, page_size, x]

117 + # V: [num_blocks, num_head, page_size // x, head_dim, x]

118 + key_cache_ptr_offset = (

119 + key_cache_ptr

120 + + block_id * num_heads * head_size * PAGE_SIZE

121 + + head_id * head_size * PAGE_SIZE

122 + + slot_id * x

123 + )

124 + value_cache_ptr_offset = (

125 + value_cache_ptr

126 + + block_id * num_heads * head_size * PAGE_SIZE

127 + + head_id * head_size * PAGE_SIZE

128 + + (slot_id // x) * head_size * x

129 + + slot_id % x

130 + )

131 + k_reg_offset = col_offsets // x * PAGE_SIZE * x + col_offsets % x

132 + v_reg_offset = col_offsets * x

133 + k_reg = tl.load(key_cache_ptr_offset + k_reg_offset)

134 + v_reg = tl.load(value_cache_ptr_offset + v_reg_offset)

135 + if DEQUANT:

136 + k_scale = 1.0

137 + v_scale = 1.0

138 + k_reg = k_reg.to(tl.float32) * k_scale

139 + v_reg = v_reg.to(tl.float32) * v_scale

140 + tl.store(key_ptr_offset + col_offsets, k_reg)

141 + tl.store(value_ptr_offset + col_offsets, v_reg)

142

143 def cp_mha_gather_cache(

144 key_cache: torch.Tensor,

@@ -123,17 +155,14 @@ if current_platform.is_rocm():

155 kv_cache_layout: str,

156 total_tokens: int,

157 ):

126 - assert kv_cache_layout in ["v0", "NHD", "HND"], (

127 - "kv_cache_layout only support v0, NHD, HND"

158 + assert kv_cache_layout in ["NHD", "SHUFFLE"], (

159 + "kv_cache_layout only support NHD, SHUFFLE"

160 )

161 head_dim = key.shape[2]

130 - x = 0

162 + x = 16 // key_cache.element_size()

163 # assert dequant is True, "Currently, we only support "\

164 # "gather cache with dequant"

165 # For k cache layout: [num_blocks, num_heads, page_size, head_dim]

134 - assert kv_cache_layout == "NHD", (

135 - "ROCM_AITER_FA_BACKEND Only support NHD kv cache layout for now"

136 - )

166 assert head_dim == key_cache.shape[3], (

167 "We assume your kv cache layout is [num_blocks, "

168 "page_size, num_heads, head_dim], but got otherwise"

@@ -141,7 +170,7 @@ if current_platform.is_rocm():

170 page_size = key_cache.shape[1]

171 num_heads = key_cache.shape[2]

172

144 - grid = lambda meta: (total_tokens,)

173 + grid = lambda meta: (total_tokens, num_heads)

174 cp_mha_gather_cache_kernel[grid](

175 key_cache,

176 value_cache,

@@ -163,6 +192,112 @@ if current_platform.is_rocm():

192 BLOCK_SIZE=head_dim,

193 )

194

195 + @triton.jit

196 + def reshape_and_cache_shuffle_kernel(

197 + key_ptr, # [num_tokens, num_kv_heads, head_size]

198 + value_ptr, # [num_tokens, num_kv_heads, head_size]

199 + key_cache_ptr, # [num_blocks, num_kv_heads, head_size // x, block_size, x]

200 + value_cache_ptr, # [num_blocks, num_kv_heads, block_size // x, head_size, x]

201 + slot_mapping_ptr, # [num_tokens]

202 + k_scale_ptr, # [num_blocks, num_kv_heads, block_size]

203 + v_scale_ptr, # [num_blocks, num_kv_heads, block_size]

204 + x,

205 + k_stride0,

206 + v_stride0,

207 + block_size,

208 + head_size,

209 + num_kv_heads,

210 + BLOCK_SIZE: tl.constexpr,

211 + QUANT: tl.constexpr,

212 + IS_FNUZ: tl.constexpr,

213 + ):

214 + tid = tl.program_id(0)

215 + head_id = tl.program_id(1)

216 + offset = tl.arange(0, BLOCK_SIZE)

217 + src_offset_k = tid * k_stride0 + head_id * head_size

218 + src_offset_v = tid * v_stride0 + head_id * head_size

219 + slot_id = tl.load(slot_mapping_ptr + tid)

220 + if slot_id < 0:

221 + return

222 + block_id = slot_id // block_size

223 + block_offset = slot_id % block_size

224 + dst_offset = (

225 + block_id * num_kv_heads * head_size * block_size

226 + + head_id * head_size * block_size

227 + )

228 + dst_k_shuffle_offset = (

229 + dst_offset + offset // x * block_size * x + block_offset * x + offset % x

230 + )

231 + dst_v_shuffle_offset = (

232 + dst_offset

233 + + block_offset // x * head_size * x

234 + + offset * x

235 + + block_offset % x

236 + )

237 + k_val = tl.load(key_ptr + src_offset_k + offset)

238 + v_val = tl.load(value_ptr + src_offset_v + offset)

239 + if QUANT:

240 + k_scale = 1.0

241 + v_scale = 1.0

242 + k_dtype = key_cache_ptr.type.element_ty

243 + v_dtype = value_cache_ptr.type.element_ty

244 + k_val = (k_val.to(tl.float32) / k_scale).to(k_dtype)

245 + v_val = (v_val.to(tl.float32) / v_scale).to(v_dtype)

246 + tl.store(key_cache_ptr + dst_k_shuffle_offset, k_val)

247 + tl.store(value_cache_ptr + dst_v_shuffle_offset, v_val)

248 +

249 + def reshape_and_cache_shuffle_triton(

250 + key: torch.Tensor,

251 + value: torch.Tensor,

252 + key_cache: torch.Tensor,

253 + value_cache: torch.Tensor,

254 + slot_mapping: torch.Tensor,

255 + kv_cache_dtype: str,

256 + k_scales: torch.Tensor,

257 + v_scales: torch.Tensor,

258 + ):

259 + num_tokens = slot_mapping.shape[0]

260 + _, num_kv_heads, head_size = key.shape

261 + num_blocks, block_size, _, _ = key_cache.shape

262 + x = 16 // key_cache.element_size()

263 + k_cache_template = torch.empty(

264 + [num_blocks, num_kv_heads, head_size // x, block_size, x],

265 + dtype=key_cache.dtype,

266 + device="meta",

267 + )

268 + v_cache_template = torch.empty(

269 + [num_blocks, num_kv_heads, block_size // x, head_size, x],

270 + dtype=value_cache.dtype,

271 + device="meta",

272 + )

273 + new_key_cache = key_cache.view_as(k_cache_template)

274 + new_value_cache = value_cache.view_as(v_cache_template)

275 + QUANT = False

276 + if kv_cache_dtype.startswith("fp8"):

277 + QUANT = True

278 + grid = (

279 + num_tokens,

280 + num_kv_heads,

281 + )

282 + reshape_and_cache_shuffle_kernel[grid](

283 + key,

284 + value,

285 + new_key_cache,

286 + new_value_cache,

287 + slot_mapping,

288 + k_scales,

289 + v_scales,

290 + x,

291 + key.stride(0),

292 + value.stride(0),

293 + block_size,

294 + head_size,

295 + num_kv_heads,

296 + BLOCK_SIZE=head_size,

297 + QUANT=QUANT,

298 + IS_FNUZ=current_platform.fp8_dtype() == torch.float8_e4m3fnuz,

299 + )

300 +

301

302 logger = init_logger(__name__)

303

@@ -253,6 +388,11 @@ class AiterFlashAttentionMetadata:

388 common_prefix_len: int

389 total_tokens: int

390

391 + # Only for fp8 shuffle layout kv cache, we allocate kv_scale for each layer

392 + # since we might integrate per token quant for kv cache in the future.

393 + k_scale: dict[str, torch.Tensor] | None

394 + v_scale: dict[str, torch.Tensor] | None

395 +

396

397 class AiterFlashAttentionMetadataBuilder(

398 AttentionMetadataBuilder[AiterFlashAttentionMetadata]

@@ -303,6 +443,7 @@ class AiterFlashAttentionMetadataBuilder(

443 dtype=self.model_config.dtype,

444 device=device,

445 )

446 + self.scale = torch.tensor([1.0], dtype=torch.float, device=self.device)

447

448 def build_for_cudagraph_capture(

449 self, common_attn_metadata: CommonAttentionMetadata

@@ -325,7 +466,27 @@ class AiterFlashAttentionMetadataBuilder(

466 common_attn_metadata,

467 decode_threshold=self.reorder_batch_threshold,

468 )

328 -

469 + # Allocate scales for fp8 shuffle kv cache with shuffle_kv_cache enabled

470 + if (

471 + rocm_aiter_ops.is_shuffle_kv_cache_enabled()

472 + and self.scale.numel() == 1

473 + and self.vllm_config.cache_config.cache_dtype.startswith("fp8")

474 + ):

475 + layers = get_layers_from_vllm_config(self.vllm_config, Attention)

476 + first_layer_name = [k for k in layers][0]

477 + kv_cache_shape = (

478 + self.vllm_config.compilation_config.static_forward_context[

479 + first_layer_name

480 + ]

481 + .kv_cache[0]

482 + .shape

483 + )

484 + num_blocks = kv_cache_shape[1]

485 + self.scale = torch.ones(

486 + [num_blocks, self.num_heads_kv, self.block_size],

487 + dtype=torch.float32,

488 + device=self.device,

489 + )

490 (

491 num_decodes,

492 num_extends,

@@ -507,6 +668,8 @@ class AiterFlashAttentionMetadataBuilder(

668 use_cascade=use_cascade,

669 common_prefix_len=common_prefix_len,

670 total_tokens=self.total_tokens,

671 + k_scale=self.scale,

672 + v_scale=self.scale,

673 )

674 return attn_metadata

675

@@ -548,7 +711,6 @@ class AiterFlashAttentionBackend(AttentionBackend):

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

712 if block_size % 16 != 0:

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

551 -

714 return (2, num_blocks, block_size, num_kv_heads, head_size)

715

716

@@ -630,7 +792,7 @@ class AiterFlashAttentionImpl(AttentionImpl):

792 cu_seqlens_kv=swa_cu_seqlens,

793 token_to_batch=swa_token_to_batch,

794 seq_starts=swa_seq_starts,

633 - dequant=False,

795 + dequant=self.kv_cache_dtype.startswith("fp8"),

796 kv_cache_layout="NHD",

797 total_tokens=swa_total_tokens,

798 )

@@ -668,8 +830,8 @@ class AiterFlashAttentionImpl(AttentionImpl):

830 min_seqlen_q: int,

831 block_table: torch.Tensor,

832 slot_mapping: torch.Tensor,

671 - k_scale: float,

672 - v_scale: float,

833 + k_scale: torch.Tensor,

834 + v_scale: torch.Tensor,

835 ):

836 if self.sliding_window[0] != -1:

837 self.extend_for_sliding_window(

@@ -725,8 +887,10 @@ class AiterFlashAttentionImpl(AttentionImpl):

887 cu_seqlens_kv=cu_seqlens_kv[chunk_idx],

888 token_to_batch=token_to_batch[chunk_idx],

889 seq_starts=chunk_starts[chunk_idx],

728 - dequant=False,

729 - kv_cache_layout="NHD",

890 + dequant=self.kv_cache_dtype.startswith("fp8"),

891 + kv_cache_layout="SHUFFLE"

892 + if rocm_aiter_ops.is_shuffle_kv_cache_enabled()

893 + else "NHD",

894 total_tokens=total_token_per_batch[chunk_idx],

895 )

896

@@ -823,6 +987,9 @@ class AiterFlashAttentionImpl(AttentionImpl):

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

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

989 # in KV cache.

990 + if self.kv_cache_dtype.startswith("fp8"):

991 + key_cache = key_cache.view(current_platform.fp8_dtype())

992 + value_cache = value_cache.view(current_platform.fp8_dtype())

993 if (

994 self.kv_sharing_target_layer_name is None

995 and key is not None

@@ -835,21 +1002,31 @@ class AiterFlashAttentionImpl(AttentionImpl):

1002 # key[:num_actual_tokens] and value[:num_actual_tokens] because

1003 # the reshape_and_cache_flash op uses the slot_mapping's shape

1004 # to determine the number of actual tokens.

838 -

839 - torch.ops._C_cache_ops.reshape_and_cache_flash(

840 - key,

841 - value,

842 - key_cache,

843 - value_cache,

844 - attn_metadata.slot_mapping,

845 - self.kv_cache_dtype,

846 - layer._k_scale,

847 - layer._v_scale,

848 - )

849 -

850 - if self.kv_cache_dtype.startswith("fp8"):

851 - key_cache = key_cache.view(current_platform.fp8_dtype())

852 - value_cache = value_cache.view(current_platform.fp8_dtype())

1005 + if rocm_aiter_ops.is_shuffle_kv_cache_enabled():

1006 + # We may calculate per token quant scale in

1007 + # reshape_and_cache_shuffle_triton which might differ from

1008 + # vllm's style when shuffle layout is used.

1009 + reshape_and_cache_shuffle_triton(

1010 + key,

1011 + value,

1012 + key_cache,

1013 + value_cache,

1014 + attn_metadata.slot_mapping,

1015 + self.kv_cache_dtype,

1016 + attn_metadata.k_scale,

1017 + attn_metadata.v_scale,

1018 + )

1019 + else:

1020 + torch.ops._C_cache_ops.reshape_and_cache_flash(

1021 + key,

1022 + value,

1023 + key_cache,

1024 + value_cache,

1025 + attn_metadata.slot_mapping,

1026 + self.kv_cache_dtype,

1027 + layer._k_scale,

1028 + layer._v_scale,

1029 + )

1030

1031 # decode:extend:prefill

1032 query = query[:num_actual_tokens]

@@ -902,6 +1079,11 @@ class AiterFlashAttentionImpl(AttentionImpl):

1079 extend_keys = key[extend_tokens_slice]

1080 extend_values = value[extend_tokens_slice]

1081 extend_outputs = output[extend_tokens_slice]

1082 + k_scale = layer._k_scale

1083 + v_scale = layer._v_scale

1084 + if rocm_aiter_ops.is_shuffle_kv_cache_enabled():

1085 + k_scale = attn_metadata.k_scale

1086 + v_scale = attn_metadata.v_scale

1087 self.extend_forward(

1088 attn_metadata=attn_metadata,

1089 query=extend_querys,

@@ -920,14 +1102,17 @@ class AiterFlashAttentionImpl(AttentionImpl):

1102 slot_mapping=attn_metadata.slot_mapping[

1103 num_decodes : num_decodes + num_extends

1104 ],

923 - k_scale=layer._k_scale,

924 - v_scale=layer._v_scale,

1105 + k_scale=k_scale,

1106 + v_scale=v_scale,

1107 )

1108

1109 # calculate for decodes

1110 if num_decodes > 0:

1111 assert attn_metadata.decode_metadata is not None

1112 if self.sliding_window[0] != -1:

1113 + assert not rocm_aiter_ops.is_shuffle_kv_cache_enabled(), (

1114 + "Sliding window with shuffle layout is not supported yet."

1115 + )

1116 from aiter.ops.triton.unified_attention import (

1117 unified_attention,

1118 )

@@ -957,41 +1142,71 @@ class AiterFlashAttentionImpl(AttentionImpl):

1142 )

1143 return

1144 assert attn_metadata.decode_metadata is not None

960 - _, num_heads, head_size = query.shape

961 - nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8

962 - num_seqs = attn_metadata.seq_lens.shape[0]

963 - max_num_partitions = (

964 - attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1

965 - ) // _PARTITION_SIZE_ROCM

966 -

967 - workspace_buffer = torch.empty(

968 - (num_seqs * num_heads * max_num_partitions * head_size)

969 - * nbytes_per_qo_elem

970 - + 2 * (num_seqs * num_heads * max_num_partitions) * 4,

971 - dtype=torch.uint8,

972 - device=output.device,

973 - )

1145

975 - torch.ops.aiter.paged_attention_v1(

976 - output[:num_decode_tokens],

977 - workspace_buffer,

978 - query[:num_decode_tokens],

979 - key_cache,

980 - value_cache,

981 - self.scale,

982 - attn_metadata.block_table[:num_decodes],

983 - attn_metadata.query_start_loc[:num_decodes],

984 - attn_metadata.seq_lens[:num_decodes],

985 - attn_metadata.max_seq_len,

986 - self.alibi_slopes,

987 - self.kv_cache_dtype,

988 - "NHD",

989 - self.logits_soft_cap,

990 - layer._k_scale,

991 - layer._v_scale,

992 - None,

993 - _PARTITION_SIZE_ROCM,

994 - )

1146 + if rocm_aiter_ops.is_shuffle_kv_cache_enabled():

1147 + num_blocks, block_size, num_kv_heads, head_size = key_cache.shape

1148 + x = 16 // key_cache.element_size()

1149 + k_cache_template = torch.empty(

1150 + [num_blocks, num_kv_heads, head_size // x, block_size, x],

1151 + dtype=key_cache.dtype,

1152 + device="meta",

1153 + )

1154 + v_cache_template = torch.empty(

1155 + [num_blocks, num_kv_heads, block_size // x, head_size, x],

1156 + dtype=value_cache.dtype,

1157 + device="meta",

1158 + )

1159 + new_key_cache = key_cache.view_as(k_cache_template)

1160 + new_value_cache = value_cache.view_as(v_cache_template)

1161 + aiter.pa_fwd_asm(

1162 + Q=query[:num_decode_tokens],

1163 + K=new_key_cache,

1164 + V=new_value_cache,

1165 + block_tables=attn_metadata.block_table[:num_decodes],

1166 + context_lens=attn_metadata.seq_lens[:num_decodes],

1167 + block_tables_stride0=attn_metadata.block_table[

1168 + :num_decodes

1169 + ].stride(0),

1170 + K_QScale=attn_metadata.k_scale,

1171 + V_QScale=attn_metadata.v_scale,

1172 + out_=output[:num_decode_tokens],

1173 + )

1174 + else:

1175 + _, num_heads, head_size = query.shape

1176 + nbytes_per_qo_elem = torch.finfo(query.dtype).bits // 8

1177 + num_seqs = attn_metadata.seq_lens.shape[0]

1178 + max_num_partitions = (

1179 + attn_metadata.max_seq_len + _PARTITION_SIZE_ROCM - 1

1180 + ) // _PARTITION_SIZE_ROCM

1181 +

1182 + workspace_buffer = torch.empty(

1183 + (num_seqs * num_heads * max_num_partitions * head_size)

1184 + * nbytes_per_qo_elem

1185 + + 2 * (num_seqs * num_heads * max_num_partitions) * 4,

1186 + dtype=torch.uint8,

1187 + device=output.device,

1188 + )

1189 +

1190 + torch.ops.aiter.paged_attention_v1(

1191 + output[:num_decode_tokens],

1192 + workspace_buffer,

1193 + query[:num_decode_tokens],

1194 + key_cache,

1195 + value_cache,

1196 + self.scale,

1197 + attn_metadata.block_table[:num_decodes],

1198 + attn_metadata.query_start_loc[:num_decodes],

1199 + attn_metadata.seq_lens[:num_decodes],

1200 + attn_metadata.max_seq_len,

1201 + self.alibi_slopes,

1202 + self.kv_cache_dtype,

1203 + "NHD",

1204 + self.logits_soft_cap,

1205 + layer._k_scale,

1206 + layer._v_scale,

1207 + None,

1208 + _PARTITION_SIZE_ROCM,

1209 + )

1210 else:

1211 raise NotImplementedError(

1212 "Cascade attention is not implemented for ROCM AITER"