@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1a19e9cd87

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

[Bugfix][ROCm]Fix Qwen3-Next-80B-A3B-Thinking inference and optimize non-standard block size (544) support under rocm_atten (#31380)

Signed-off-by: vllmellm <vllm.ellm@embeddedllm.com>

vllmellm committed Jan 9, 2026 at 12:28 UTC 1a19e9cd87b6bb091fc9c530e6446601db52b58e

5 files changed +283 -84

tests/kernels/attention/test_prefix_prefill.py

+33 -2

@@ -112,6 +112,7 @@ def test_contexted_kv_attention(

112 kv_cache_dtype: str,

113 device: str,

114 op: Callable,

115 + block_size: int = 32,

116 ) -> None:

117 if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):

118 pytest.skip(

@@ -138,7 +139,6 @@ def test_contexted_kv_attention(

139 MAX_CTX_LEN = 1024

140 BS = 10

141 cache_size = 640

141 - block_size = 32

142 max_block_per_request = 64

143 query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]

144 # ensure one sequence in batch is a decode

@@ -333,6 +333,7 @@ def test_contexted_kv_attention_alibi(

333 kv_cache_dtype: str,

334 device: str,

335 op: Callable,

336 + block_size: int = 32,

337 ) -> None:

338 if "fp8" in kv_cache_dtype and not current_platform.has_device_capability(89):

339 pytest.skip(

@@ -385,7 +386,6 @@ def test_contexted_kv_attention_alibi(

386 MAX_CTX_LEN = 1024

387 BS = 10

388 cache_size = 640

388 - block_size = 32

389 max_block_per_request = 64

390 query_lens = [random.randint(16, MAX_SEQ_LEN) for _ in range(BS)]

391 ctx_lens = [random.randint(16, MAX_CTX_LEN) for _ in range(BS)]

@@ -637,3 +637,34 @@ def test_contexted_kv_attention_alibi_f32(

637 test_contexted_kv_attention_alibi(

638 num_heads, num_queries_per_kv, head_size, dtype, kv_cache_dtype, device, op

639 )

640 +

641 +

642 +@pytest.mark.parametrize("head_size", [128])

643 +@pytest.mark.parametrize("dtype", DTYPES)

644 +@pytest.mark.parametrize("device", CUDA_DEVICES)

645 +@pytest.mark.parametrize("op", OPS)

646 +@torch.inference_mode()

647 +def test_qwen3_nonstandard_block_size(

648 + head_size: int,

649 + dtype: torch.dtype,

650 + device: str,

651 + op: Callable,

652 +) -> None:

653 + """

654 + A separate test function specifically added

655 + for Qwen3-Next-80B (Block Size 544).

656 + """

657 + if not current_platform.is_rocm():

658 + pytest.skip("544 block size optimization is only for ROCm.")

659 +

660 + test_contexted_kv_attention(

661 + num_heads=64,

662 + num_queries_per_kv=1,

663 + head_size=head_size,

664 + block_size=544,

665 + sliding_window=0,

666 + dtype=dtype,

667 + kv_cache_dtype="auto",

668 + device=device,

669 + op=op,

670 + )

vllm/attention/ops/chunked_prefill_paged_decode.py

+83 -28

@@ -46,6 +46,7 @@ def kernel_paged_attention_2d(

46 output_stride_0: tl.int64, # int

47 output_stride_1: tl.int64, # int, should be equal to head_size

48 BLOCK_SIZE: tl.constexpr, # int

49 + PHYSICAL_BLOCK_SIZE: tl.constexpr, # int

50 HEAD_SIZE: tl.constexpr, # int

51 HEAD_SIZE_PADDED: tl.constexpr, # int, must be power of 2

52 USE_ALIBI_SLOPES: tl.constexpr, # bool

@@ -104,14 +105,15 @@ def kernel_paged_attention_2d(

105

106 if not USE_SINKS:

107 M = tl.full([num_queries_per_kv_padded], float("-inf"), dtype=tl.float32)

108 + L = tl.zeros([num_queries_per_kv_padded], dtype=tl.float32)

109 else:

110 M = tl.load(

111 sink_ptr + query_head_idx,

112 mask=head_mask,

113 other=float("-inf"),

114 ).to(dtype=tl.float32)

115 + L = tl.where(float("-inf") < M, 1.0, 0.0)

116

114 - L = tl.full([num_queries_per_kv_padded], 1.0, dtype=tl.float32)

117 acc = tl.zeros([num_queries_per_kv_padded, HEAD_SIZE_PADDED], dtype=tl.float32)

118

119 # sequence len for this particular sequence

@@ -125,30 +127,45 @@ def kernel_paged_attention_2d(

127

128 num_blocks = cdiv_fn(seq_len, BLOCK_SIZE)

129

130 + offs_n = tl.arange(0, BLOCK_SIZE)

131 + offs_d = tl.arange(0, HEAD_SIZE_PADDED)

132 # iterate through tiles

133 for j in range(0, num_blocks):

130 - physical_block_idx = tl.load(block_tables_ptr + block_table_offset + j)

131 -

132 - offs_n = tl.arange(0, BLOCK_SIZE)

133 - offs_d = tl.arange(0, HEAD_SIZE_PADDED)

134 -

135 - v_offset = (

136 - physical_block_idx * stride_v_cache_0

137 - + kv_head_idx * stride_v_cache_1

138 - + offs_d[None, :] * stride_v_cache_2

139 - + offs_n[:, None] * stride_v_cache_3

140 - )

141 -

134 + start_n = j * BLOCK_SIZE

135 + # Calculate the logical location within a non-standard physical block,

136 + # such as 544 in Qwen/Qwen3-Next-80B-A3B-Thinking.

137 + # Supports non-contiguous mapping

138 + # from logical blocks to physical blocks

139 + abs_token_idx = start_n + offs_n

140 + l_block_idx = abs_token_idx // PHYSICAL_BLOCK_SIZE

141 + # Vectorized loading of physical block IDs

142 + p_block_idx = tl.load(block_tables_ptr + block_table_offset + l_block_idx)

143 + internal_offsets = abs_token_idx % PHYSICAL_BLOCK_SIZE

144 +

145 + # 5D addressing logic of K

146 k_offset = (

143 - physical_block_idx * stride_k_cache_0

147 + p_block_idx[None, :] * stride_k_cache_0

148 + kv_head_idx * stride_k_cache_1

149 + (offs_d[:, None] // x) * stride_k_cache_2

146 - + offs_n[None, :] * stride_k_cache_3

150 + + internal_offsets[None, :] * stride_k_cache_3

151 + (offs_d[:, None] % x) * stride_k_cache_4

152 )

153

154 + # 4D addressing logic of V (Slot is innermost)

155 + v_offset = (

156 + p_block_idx[:, None] * stride_v_cache_0

157 + + kv_head_idx * stride_v_cache_1

158 + + offs_d[None, :] * stride_v_cache_2

159 + + internal_offsets[:, None] * stride_v_cache_3

160 + )

161 +

162 # K : (HEAD_SIZE, BLOCK_SIZE)

151 - K_load = tl.load(key_cache_ptr + k_offset, mask=dim_mask[:, None], other=0.0)

163 + K_load = tl.load(

164 + key_cache_ptr + k_offset,

165 + mask=dim_mask[:, None],

166 + other=0.0,

167 + eviction_policy="evict_last",

168 + )

169

170 if K_load.dtype.is_fp8():

171 K = (K_load.to(tl.float32) * tl.load(k_scale)).to(Q.dtype)

@@ -156,7 +173,12 @@ def kernel_paged_attention_2d(

173 K = K_load

174

175 # V : (BLOCK_SIZE, HEAD_SIZE)

159 - V_load = tl.load(value_cache_ptr + v_offset, mask=dim_mask[None, :], other=0.0)

176 + V_load = tl.load(

177 + value_cache_ptr + v_offset,

178 + mask=dim_mask[None, :],

179 + other=0.0,

180 + eviction_policy="evict_last",

181 + )

182

183 if V_load.dtype.is_fp8():

184 V = (V_load.to(tl.float32) * tl.load(v_scale)).to(Q.dtype)

@@ -167,9 +189,9 @@ def kernel_paged_attention_2d(

189 boundary = tl.full([BLOCK_SIZE], seq_len, dtype=tl.int32)

190 seq_mask = seq_offset[None, :] < boundary

191

170 - # S : (num_queries_per_kv, BLOCK_SIZE,)

171 - S = tl.where(head_mask[:, None] & seq_mask, 0.0, float("-inf")).to(tl.float32)

172 - S += scale * tl.dot(Q, K)

192 + # First calculate the dot, then apply the mask.

193 + qk = scale * tl.dot(Q, K)

194 + S = tl.where(head_mask[:, None] & seq_mask, qk, float("-inf"))

195

196 context_len = seq_len - 1

197

@@ -184,13 +206,15 @@ def kernel_paged_attention_2d(

206 m_j = tl.maximum(M, tl.max(S, axis=1))

207

208 # P : (num_queries_per_kv, BLOCK_SIZE,)

187 - P = tl.exp(S - m_j[:, None])

209 + p = tl.exp(S - m_j[:, None])

210 + p = tl.where(m_j[:, None] == float("-inf"), 0.0, p)

211

212 # l_j : (num_queries_per_kv,)

190 - l_j = tl.sum(P, axis=1)

213 + l_j = tl.sum(p, axis=1)

214

215 # alpha : (num_queries_per_kv, )

216 alpha = tl.exp(M - m_j)

217 + alpha = tl.where(float("-inf") == M, 0.0, alpha)

218

219 # acc : (num_queries_per_kv, BLOCK_SIZE,)

220 acc = acc * alpha[:, None]

@@ -200,10 +224,10 @@ def kernel_paged_attention_2d(

224 M = m_j

225

226 # acc : (num_queries_per_kv, BLOCK_SIZE,)

203 - acc += tl.dot(P.to(V.dtype), V)

227 + acc += tl.dot(p.to(V.dtype), V)

228

229 # epilogue

206 - acc = acc / L[:, None]

230 + acc = acc / (L[:, None] + 1e-10)

231 if USE_FP8:

232 acc = acc * tl.load(out_scale_inv)

233 acc = tl.clamp(acc, FP8_MIN, FP8_MAX)

@@ -241,9 +265,10 @@ def chunked_prefill_paged_decode(

265 output_scale=None,

266 # Optional tensor for sinks

267 sinks=None,

268 + is_block_table_ptr: bool = False,

269 ):

270 if sm_scale is None:

246 - sm_scale = 1.0 / (query.shape[1] ** 0.5)

271 + sm_scale = 1.0 / (query.shape[2] ** 0.5)

272

273 use_alibi_slopes = alibi_slopes is not None

274

@@ -315,6 +340,16 @@ def chunked_prefill_paged_decode(

340 alibi_slopes,

341 sinks,

342 )

343 + # Triton is only forced when encountering a non-standard block

344 + # like Qwen3 with a size of 544.

345 + # 1. Check if block_size is a power of 2 (16, 32, 64...)

346 + # 2. If it's a power of 2, we trust the vLLM's native use_custom decision.

347 + # 3. If it's not a power of 2 (such as Qwen3's 544),

348 + # then our Triton path is forced.

349 + is_pow2 = block_size > 0 and (block_size & (block_size - 1) == 0)

350 + if not is_pow2:

351 + use_custom = False

352 +

353 if use_custom:

354 _PARTITION_SIZE_ROCM = 256

355 max_num_partitions = (

@@ -356,6 +391,25 @@ def chunked_prefill_paged_decode(

391 fp8_out_scale=output_scale,

392 )

393 else:

394 + real_block_size = value_cache.shape[3]

395 + # The standard model directly uses the original block_size.

396 + # Non-standard 544 uses 32 to accommodate integer division logic.

397 + TRITON_BLOCK_SIZE = block_size if is_pow2 else 32

398 + if is_block_table_ptr:

399 + # Using the physical base address of tensors

400 + kv_element_size = key_cache.element_size()

401 + block_byte_stride = key_cache.stride(0) * kv_element_size

402 + # Get the starting physical address of the KV Cache

403 + base_addr = key_cache.data_ptr()

404 +

405 + # Normalization: Directly calculate the block offset

406 + # of the pointer relative to the base address

407 + processed_block_table = ((block_table - base_addr) // block_byte_stride).to(

408 + torch.int32

409 + )

410 + else:

411 + processed_block_table = block_table.to(torch.int32)

412 +

413 kernel_paged_attention_2d[

414 (

415 num_seqs,

@@ -367,7 +421,7 @@ def chunked_prefill_paged_decode(

421 key_cache_ptr=key_cache,

422 value_cache_ptr=value_cache,

423 sink_ptr=sinks,

370 - block_tables_ptr=block_table,

424 + block_tables_ptr=processed_block_table,

425 seq_lens_ptr=seq_lens,

426 alibi_slopes_ptr=alibi_slopes,

427 scale=sm_scale,

@@ -377,12 +431,13 @@ def chunked_prefill_paged_decode(

431 num_query_heads=num_query_heads,

432 num_queries_per_kv=num_queries_per_kv,

433 num_queries_per_kv_padded=num_queries_per_kv_padded,

380 - block_table_stride=block_table.stride(0),

434 + block_table_stride=processed_block_table.stride(0),

435 query_stride_0=query.stride(0),

436 query_stride_1=query.stride(1),

437 output_stride_0=output.stride(0),

438 output_stride_1=output.stride(1),

385 - BLOCK_SIZE=block_size,

439 + BLOCK_SIZE=TRITON_BLOCK_SIZE,

440 + PHYSICAL_BLOCK_SIZE=real_block_size,

441 HEAD_SIZE=head_size,

442 HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),

443 USE_ALIBI_SLOPES=use_alibi_slopes,

vllm/attention/ops/prefix_prefill.py

+80 -32

@@ -79,6 +79,7 @@ def _fwd_kernel(

79 BLOCK_DMODEL: tl.constexpr,

80 BLOCK_DMODEL_PADDED: tl.constexpr,

81 BLOCK_SIZE: tl.constexpr,

82 + PHYSICAL_BLOCK_SIZE: tl.constexpr,

83 BLOCK_N: tl.constexpr,

84 SLIDING_WINDOW: tl.constexpr,

85 num_unroll_cache: tl.constexpr,

@@ -139,42 +140,52 @@ def _fwd_kernel(

140 # initialize pointer to m and l

141 if not USE_SINKS:

142 m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)

143 + l_i = tl.zeros([BLOCK_M], dtype=tl.float32)

144 else:

145 m_i = tl.load(

146 sink_ptr + tl.full([BLOCK_M], cur_head, dtype=tl.int64),

147 mask=(offs_m < cur_batch_query_len),

148 other=float("-inf"),

149 ).to(dtype=tl.float32)

150 + l_i = tl.where(m_i > float("-inf"), 1.0, 0.0)

151

149 - l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32)

152 acc = tl.zeros([BLOCK_M, BLOCK_DMODEL_PADDED], dtype=tl.float32) # [M,D]

153

154 # compute query against context (no causal mask here)

155 for start_n in tl.range(

156 0, cur_batch_ctx_len, BLOCK_SIZE, loop_unroll_factor=num_unroll_cache

157 ):

156 - start_n = tl.multiple_of(start_n, BLOCK_SIZE)

157 - # -- compute qk ----

158 + # Under a block size of 544 (Qwen/Qwen3-Next-80B-A3B-Thinking),

159 + # replace one physical block every 17 32-Tile blocks

160 + # Calculate the logical block index of each of the 32 tokens

161 + # in the current Tile (handling cross-block cases).

162 + token_indices = start_n + offs_bs_n

163 + bn_logical_indices = token_indices // PHYSICAL_BLOCK_SIZE

164 +

165 + # 2. Vectorized loading of physical block IDs from B_Loc

166 bn = tl.load(

159 - B_Loc

160 - + cur_batch * stride_b_loc_b

161 - + (start_n // BLOCK_SIZE) * stride_b_loc_s

167 + B_Loc + cur_batch * stride_b_loc_b + bn_logical_indices * stride_b_loc_s

168 ).to(tl.int64)

163 - # [D,BLOCK_SIZE]

169 +

170 + # 3. Calculate the exact offset of

171 + # each token within its physical block.

172 + internal_offsets = token_indices % PHYSICAL_BLOCK_SIZE

173 +

174 + # Addressing of K (5D)

175 off_k = (

176 bn[None, :] * stride_k_cache_bs

177 + cur_kv_head * stride_k_cache_h

178 + (offs_d[:, None] // x) * stride_k_cache_d

168 - + ((start_n + offs_bs_n[None, :]) % BLOCK_SIZE) * stride_k_cache_bl

179 + + internal_offsets[None, :] * stride_k_cache_bl

180 + (offs_d[:, None] % x) * stride_k_cache_x

181 )

182

172 - # [BLOCK_SIZE,D]

183 + # Addressing of V (4D)

184 off_v = (

185 bn[:, None] * stride_v_cache_bs

186 + cur_kv_head * stride_v_cache_h

187 + offs_d[None, :] * stride_v_cache_d

177 - + offs_bs_n[:, None] * stride_v_cache_bl

188 + + internal_offsets[:, None] * stride_v_cache_bl

189 )

190

191 if (

@@ -195,12 +206,12 @@ def _fwd_kernel(

206 else:

207 k = k_load

208

198 - qk = tl.zeros([BLOCK_M, BLOCK_SIZE], dtype=tl.float32) # [M,N]

199 - qk = tl.dot(q, k, acc=qk, input_precision=IN_PRECISION)

209 + # qk = tl.zeros([BLOCK_M, BLOCK_SIZE], dtype=tl.float32) # [M,N]

210 + qk = sm_scale * tl.dot(q, k, input_precision=IN_PRECISION)

211 qk = tl.where(

212 (start_n + offs_bs_n[None, :]) < cur_batch_ctx_len, qk, float("-inf")

213 )

203 - qk *= sm_scale

214 + # qk *= sm_scale

215 if SLIDING_WINDOW > 0:

216 # (cur_batch_ctx_len + offs_m[:, None]) are the positions of

217 # Q entries in sequence

@@ -217,14 +228,16 @@ def _fwd_kernel(

228 (cur_batch_ctx_len + offs_m[:, None]) - (start_n + offs_bs_n[None, :])

229 < SLIDING_WINDOW,

230 qk,

220 - -10000,

231 + float("-inf"),

232 )

233

234 # compute running maximum

235 m_ij = tl.maximum(m_i, tl.max(qk, axis=1))

236 p = tl.exp(qk - m_ij[:, None])

237 + p = tl.where(m_ij[:, None] == float("-inf"), 0.0, p)

238 l_ij = tl.sum(p, axis=1)

239 alpha = tl.exp(m_i - m_ij)

240 + alpha = tl.where(m_i == float("-inf"), 0.0, alpha)

241 acc = acc * alpha[:, None]

242

243 # update acc

@@ -293,14 +306,17 @@ def _fwd_kernel(

306 qk = tl.where(

307 offs_m[:, None] - (start_n + offs_n[None, :]) < SLIDING_WINDOW,

308 qk,

296 - -10000,

309 + float("-inf"),

310 )

311

312 # compute running maximum

313 m_ij = tl.maximum(m_i, tl.max(qk, axis=1))

314 p = tl.exp(qk - m_ij[:, None])

315 + p = tl.where(m_ij[:, None] == float("-inf"), 0.0, p)

316 l_ij = tl.sum(p, axis=1)

317 alpha = tl.exp(m_i - m_ij)

318 + # To prevent NaN from appearing in the first round

319 + alpha = tl.where(m_i == float("-inf"), 0.0, alpha)

320 acc = acc * alpha[:, None]

321

322 # update acc

@@ -317,7 +333,7 @@ def _fwd_kernel(

333 l_i = l_i * alpha + l_ij

334 m_i = m_ij

335

320 - acc = acc / l_i[:, None]

336 + acc = acc / (l_i[:, None] + 1e-10)

337

338 # initialize pointers to output

339 off_o = (

@@ -637,6 +653,7 @@ def context_attention_fwd(

653 skip_decode=False,

654 fp8_out_scale=None,

655 sinks=None,

656 + is_block_table_ptr: bool = False,

657 ):

658 q_dtype_is_f32 = q.dtype is torch.float32

659

@@ -689,6 +706,19 @@ def context_attention_fwd(

706 if sliding_window is None or sliding_window <= 0:

707 sliding_window = 0

708

709 + if is_block_table_ptr:

710 + kv_element_size = k_cache.element_size()

711 + block_byte_stride = k_cache.stride(0) * kv_element_size

712 + # The physical starting point of the obtained KV Cache Pool

713 + base_addr = k_cache.data_ptr()

714 +

715 + mask = b_loc > 0

716 + processed_b_loc = torch.where(

717 + mask, (b_loc - base_addr) // block_byte_stride, b_loc

718 + ).to(torch.int32)

719 + else:

720 + processed_b_loc = b_loc.to(torch.int32)

721 +

722 if alibi_slopes is not None:

723 assert sinks is None, "Sinks arg is not supported with alibi"

724 assert fp8_out_scale is None, "FP8 output not supported with alibi"

@@ -752,7 +782,24 @@ def context_attention_fwd(

782 max_seq_len = 0 if max_seq_len is None else max_seq_len

783 extra_kargs = {}

784 if current_platform.is_rocm():

755 - extra_kargs = {"kpack": 1, "waves_per_eu": 2}

785 + extra_kargs = {}

786 +

787 + real_block_size = v_cache.shape[3]

788 + is_pow2 = real_block_size > 0 and (real_block_size & (real_block_size - 1) == 0)

789 + # For standard models involving powers of 2,

790 + # follow the original logic (Llama 128/64)

791 + # For non-standard models (Qwen3-next block_size 544), set to 32.

792 + if is_pow2:

793 + BLOCK_M = 128

794 + BLOCK_N = 64

795 + else:

796 + BLOCK_M = 32

797 + BLOCK_N = 32

798 +

799 + # TRITON_BLOCK_SIZE is kept at 32 to ensure

800 + # correct alignment logic when the kernel handles

801 + # non-standard sizes (such as 544).

802 + TRITON_BLOCK_SIZE = 32

803

804 grid_fn = lambda META: (batch, head, triton.cdiv(max_input_len, META["BLOCK_M"]))

805 _fwd_kernel[grid_fn](

@@ -762,7 +809,7 @@ def context_attention_fwd(

809 k_cache,

810 v_cache,

811 sinks,

765 - b_loc,

812 + processed_b_loc,

813 sm_scale,

814 k_scale,

815 v_scale,

@@ -771,8 +818,8 @@ def context_attention_fwd(

818 b_seq_len,

819 k_cache.shape[4],

820 o,

774 - b_loc.stride(0),

775 - b_loc.stride(1),

821 + processed_b_loc.stride(0),

822 + processed_b_loc.stride(1),

823 q.stride(0),

824 q.stride(1),

825 q.stride(2),

@@ -785,16 +832,17 @@ def context_attention_fwd(

832 o.stride(0),

833 o.stride(1),

834 o.stride(2),

788 - k_cache.stride(0),

789 - k_cache.stride(1),

790 - k_cache.stride(2),

791 - k_cache.stride(3),

792 - k_cache.stride(4), # [num_blocks, num_kv_heads, head_size/x, block_size, x]

793 - v_cache.stride(0),

794 - v_cache.stride(1),

795 - v_cache.stride(2),

796 - v_cache.stride(3), # [num_blocks, num_kv_heads, head_size, block_size]

797 - BLOCK_SIZE=v_cache.shape[3],

835 + stride_k_cache_bs=k_cache.stride(0),

836 + stride_k_cache_h=k_cache.stride(1),

837 + stride_k_cache_d=k_cache.stride(2),

838 + stride_k_cache_bl=k_cache.stride(3),

839 + stride_k_cache_x=k_cache.stride(4),

840 + stride_v_cache_bs=v_cache.stride(0),

841 + stride_v_cache_h=v_cache.stride(1),

842 + stride_v_cache_d=v_cache.stride(2),

843 + stride_v_cache_bl=v_cache.stride(3),

844 + BLOCK_SIZE=TRITON_BLOCK_SIZE,

845 + PHYSICAL_BLOCK_SIZE=real_block_size,

846 num_queries_per_kv=num_queries_per_kv,

847 IN_PRECISION=IN_PRECISION,

848 BLOCK_DMODEL=Lk,

@@ -802,8 +850,8 @@ def context_attention_fwd(

850 SLIDING_WINDOW=sliding_window,

851 SKIP_DECODE=skip_decode,

852 USE_FP8=fp8_out_scale is not None,

805 - BLOCK_M=128,

806 - BLOCK_N=64,

853 + BLOCK_M=BLOCK_M,

854 + BLOCK_N=BLOCK_N,

855 num_unroll_cache=4,

856 num_unroll_request=1,

857 num_warps=4,

vllm/attention/ops/triton_reshape_and_cache_flash.py

+52 -12

@@ -20,10 +20,15 @@ def reshape_and_cache_kernel_flash(

20 key_stride: tl.int64,

21 value_stride: tl.int64,

22 block_stride: tl.int64,

23 + head_stride: tl.int64,

24 + dim_stride_k: tl.int64,

25 + dim_stride_v: tl.int64,

26 page_stride: tl.int64,

27 num_heads: tl.constexpr,

28 head_size: tl.constexpr,

29 block_size: tl.constexpr,

30 + x: tl.constexpr,

31 + USE_HEAD_MAJOR_LAYOUT: tl.constexpr,

32 # FP8 flags

33 FP8_KV_CACHE: tl.constexpr,

34 # tune parameters

@@ -35,17 +40,38 @@ def reshape_and_cache_kernel_flash(

40 # Padding token that should be ignored.

41 return

38 - tile_i = tl.program_id(axis=1)

39 - tile_offs = tl.arange(0, TILE_SIZE)

40 - tile_pos = tile_i * TILE_SIZE + tile_offs

41 -

43 block_idx = slot_idx // block_size

44 block_offset = slot_idx % block_size

46 + tile_i = tl.program_id(axis=1)

47 + tile_offs = tl.arange(0, TILE_SIZE)

48 + tile_pos = tile_i * TILE_SIZE + tile_offs

49 src_key_idx = token_idx * key_stride

50 src_value_idx = token_idx * value_stride

48 - tgt_idx = block_idx * block_stride + block_offset * page_stride

52 + if USE_HEAD_MAJOR_LAYOUT:

53 + # Decompose the tile index back into head and dim coordinates.

54 + cur_head = tile_pos // head_size

55 + cur_dim = tile_pos % head_size

56 + # Value addressing (4D): [Block, Head, Dim, Slot]

57 + tgt_idx_v = (

58 + block_idx * block_stride

59 + + cur_head * head_stride

60 + + cur_dim * dim_stride_v

61 + + block_offset * 1

62 + )

63 + # Key addressing (5D): [Block, Head, Dim//8, Slot, 8]

64 + tgt_idx_k = (

65 + block_idx * block_stride

66 + + cur_head * head_stride

67 + + (cur_dim // x) * dim_stride_k

68 + + block_offset * x

69 + + (cur_dim % x)

70 + )

71 + else:

72 + tgt_base = block_idx * block_stride + block_offset * page_stride

73 + tgt_idx_k = tgt_base + tile_pos

74 + tgt_idx_v = tgt_base + tile_pos

76 # [TILE_SIZE]

77 key_load = tl.load(

@@ -73,12 +99,12 @@ def reshape_and_cache_kernel_flash(

99 value_tile = value_load

100

101 tl.store(

76 - key_cache_ptr + tgt_idx + tile_pos,

102 + key_cache_ptr + tgt_idx_k,

103 key_tile,

104 mask=tile_pos < (num_heads * head_size),

105 )

106 tl.store(

81 - value_cache_ptr + tgt_idx + tile_pos,

107 + value_cache_ptr + tgt_idx_v,

108 value_tile,

109 mask=tile_pos < (num_heads * head_size),

110 )

@@ -99,17 +125,26 @@ def triton_reshape_and_cache_flash(

125 ):

126 num_heads = key.shape[1]

127 head_size = key.shape[2]

102 - block_size = key_cache.shape[1]

103 - n = num_heads * head_size

128

129 + use_head_major_layout = key_cache.ndim == 5

130 + if use_head_major_layout:

131 + block_size = key_cache.shape[3]

132 + x = key_cache.shape[4]

133 + head_stride = key_cache.stride(1)

134 + dim_stride_k = key_cache.stride(2)

135 + dim_stride_v = value_cache.stride(2)

136 + else:

137 + block_size = key_cache.shape[1]

138 + x = 1

139 + dim_stride_k = 0

140 + dim_stride_v = 0

141 + head_stride = key_cache.stride()[2]

142 + n = num_heads * head_size

143 key_stride = key.stride()[0]

144 value_stride = value.stride()[0]

145 block_stride = key_cache.stride()[0]

146 page_stride = key_cache.stride()[1]

147

110 - head_stride = key_cache.stride()[2]

111 - assert head_stride == head_size, "only continous heads are supported"

112 -

148 assert kv_cache_dtype == "auto" or kv_cache_dtype.startswith("fp8"), (

149 f"unsupported kv_cache_dtype (str), got {kv_cache_dtype}."

150 )

@@ -171,10 +206,15 @@ def triton_reshape_and_cache_flash(

206 key_stride=key_stride,

207 value_stride=value_stride,

208 block_stride=block_stride,

209 + head_stride=head_stride,

210 + dim_stride_k=dim_stride_k,

211 + dim_stride_v=dim_stride_v,

212 page_stride=page_stride,

213 num_heads=num_heads,

214 head_size=head_size,

215 block_size=block_size,

216 + x=x,

217 + USE_HEAD_MAJOR_LAYOUT=use_head_major_layout,

218 # FP8 flags

219 FP8_KV_CACHE=FP8_KV_CACHE,

220 # autotune parameters

vllm/v1/attention/backends/rocm_attn.py

+35 -10

@@ -15,6 +15,9 @@ from vllm.attention.backends.abstract import (

15 )

16 from vllm.attention.ops.chunked_prefill_paged_decode import chunked_prefill_paged_decode

17 from vllm.attention.ops.paged_attn import PagedAttention

18 +from vllm.attention.ops.triton_reshape_and_cache_flash import (

19 + triton_reshape_and_cache_flash,

20 +)

21 from vllm.config import VllmConfig

22 from vllm.logger import init_logger

23 from vllm.model_executor.layers.quantization.utils.quant_utils import (

@@ -321,16 +324,38 @@ class RocmAttentionImpl(AttentionImpl):

324 if self.kv_sharing_target_layer_name is None:

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

326 # Skip this if sharing KV cache with an earlier attention layer.

324 - PagedAttention.write_to_paged_cache(

325 - key,

326 - value,

327 - key_cache,

328 - value_cache,

329 - attn_metadata.slot_mapping,

330 - self.kv_cache_dtype,

331 - layer._k_scale,

332 - layer._v_scale,

333 - )

327 +

328 + # Get the actual block_size from value_cache

329 + # value_cache shape: [num_blocks, num_heads, head_size, block_size]

330 + block_size = value_cache.shape[3]

331 + # Determine if it is a power of 2

332 + is_pow2 = block_size > 0 and (block_size & (block_size - 1) == 0)

333 +

334 + if is_pow2:

335 + # Normal 16, 32, 64, etc., use vLLM native HIP C++ logic

336 + PagedAttention.write_to_paged_cache(

337 + key,

338 + value,

339 + key_cache,

340 + value_cache,

341 + attn_metadata.slot_mapping,

342 + self.kv_cache_dtype,

343 + layer._k_scale,

344 + layer._v_scale,

345 + )

346 + else:

347 + # Case B: Non-standard blocks (e.g., 544 in Qwen3),

348 + # force using our modified Triton logic

349 + triton_reshape_and_cache_flash(

350 + key,

351 + value,

352 + key_cache,

353 + value_cache,

354 + attn_metadata.slot_mapping,

355 + self.kv_cache_dtype,

356 + layer._k_scale,

357 + layer._v_scale,

358 + )

359

360 if self.kv_cache_dtype.startswith("fp8"):

361 key_cache = key_cache.view(self.fp8_dtype)