@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 01a583fea4

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 14, 05:30 UTC

[Kernel] Decouple Tile Size from Block Size in Triton Unified Attention Kernel (#21197)

Signed-off-by: Jan van Lunteren <jvl@zurich.ibm.com>

jvlunteren committed Sep 18, 2025 at 16:27 UTC 01a583fea40571986ffe277549e5bb441d409768

2 files changed +70 -55

tests/kernels/attention/test_triton_unified_attention.py

-3

@@ -102,9 +102,6 @@ def test_triton_unified_attn(

102 ) -> None:

103 torch.set_default_device("cuda")

104

105 - if q_dtype is not None and q_dtype.itemsize < 2 and block_size < 32:

106 - pytest.skip("block size must be at least 32 for fp8")

107 -

105 current_platform.seed_everything(0)

106 num_seqs = len(seq_lens)

107 query_lens = [x[0] for x in seq_lens]

vllm/attention/ops/triton_unified_attention.py

+70 -52

@@ -73,6 +73,7 @@ def kernel_unified_attention_2d(

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

74 qq_bias_stride_0: tl.int64, # int

75 BLOCK_SIZE: tl.constexpr, # int

76 + TILE_SIZE: tl.constexpr, # int must be power of 2

77 HEAD_SIZE: tl.constexpr, # int

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

79 USE_ALIBI_SLOPES: tl.constexpr, # bool

@@ -118,6 +119,7 @@ def kernel_unified_attention_2d(

119

120 offs_m = tl.arange(0, BLOCK_M)

121 offs_d = tl.arange(0, HEAD_SIZE_PADDED)

122 + offs_t = tl.arange(0, TILE_SIZE)

123 query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv

124

125 query_offset_0 = cur_batch_in_all_start_index + query_pos

@@ -177,31 +179,32 @@ def kernel_unified_attention_2d(

179 # actual sequence length

180 max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)

181

180 - # calculate the number of tiles (blocks) that need to be processed to

181 - # cover the longest sequence prefix (due to causal masking, blocks beyond

182 + # calculate the number of tiles that need to be processed to

183 + # cover the longest sequence prefix (due to causal masking, tiles beyond

184 # this prefix can be skipped)

183 - num_blocks = cdiv_fn(max_seq_prefix_len, BLOCK_SIZE)

185 + num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE)

186

187 # iterate through tiles

186 - for j in range(0, num_blocks):

188 + for j in range(0, num_tiles):

189 + seq_offset = j * TILE_SIZE + offs_t

190 + tile_mask = seq_offset < max_seq_prefix_len

191

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

192 + physical_block_idx = tl.load(block_tables_ptr + block_table_offset +

193 + seq_offset // BLOCK_SIZE).to(tl.int64)

194

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

191 -

192 - v_offset = (physical_block_idx * stride_v_cache_0 +

195 + v_offset = (physical_block_idx[:, None] * stride_v_cache_0 +

196 kv_head_idx * stride_v_cache_2 +

197 offs_d[None, :] * stride_v_cache_3 +

195 - offs_n[:, None] * stride_v_cache_1)

198 + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1)

199

197 - k_offset = (physical_block_idx * stride_k_cache_0 +

200 + k_offset = (physical_block_idx[None, :] * stride_k_cache_0 +

201 kv_head_idx * stride_k_cache_2 +

202 offs_d[:, None] * stride_k_cache_3 +

200 - offs_n[None, :] * stride_k_cache_1)

203 + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1)

204

202 - # K : (HEAD_SIZE, BLOCK_SIZE)

205 + # K : (HEAD_SIZE, TILE_SIZE)

206 K_load = tl.load(key_cache_ptr + k_offset,

204 - mask=dim_mask[:, None],

207 + mask=dim_mask[:, None] & tile_mask[None, :],

208 other=0.0)

209

210 if K_load.dtype.is_fp8():

@@ -212,9 +215,9 @@ def kernel_unified_attention_2d(

215 else:

216 K = K_load

217

215 - # V : (BLOCK_SIZE, HEAD_SIZE)

218 + # V : (TILE_SIZE, HEAD_SIZE)

219 V_load = tl.load(value_cache_ptr + v_offset,

217 - mask=dim_mask[None, :],

220 + mask=dim_mask[None, :] & tile_mask[:, None],

221 other=0.0)

222

223 if V_load.dtype.is_fp8():

@@ -225,12 +228,10 @@ def kernel_unified_attention_2d(

228 else:

229 V = V_load

230

228 - seq_offset = j * BLOCK_SIZE + offs_n

229 -

231 seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1

232

232 - # S : (BLOCK_M, BLOCK_SIZE)

233 - S = tl.zeros(shape=(BLOCK_M, BLOCK_SIZE), dtype=tl.float32)

233 + # S : (BLOCK_M, TILE_SIZE)

234 + S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)

235

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

237

@@ -262,11 +263,12 @@ def kernel_unified_attention_2d(

263 # compute running maximum

264 # m_j : (BLOCK_M,)

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

266 +

267 # For sliding window there's a chance the max is -inf due to masking of

268 # the entire row. In this case we need to set m_j 0 to avoid NaN

269 m_j = tl.where(m_j > float("-inf"), m_j, 0.0)

270

269 - # P : (BLOCK_M, BLOCK_SIZE)

271 + # P : (BLOCK_M, TILE_SIZE)

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

273

274 # l_j : (BLOCK_M,)

@@ -327,6 +329,7 @@ def kernel_unified_attention_3d(

329 query_stride_1: tl.int64, # int, should be equal to head_size

330 qq_bias_stride_0: tl.int64, # int

331 BLOCK_SIZE: tl.constexpr, # int

332 + TILE_SIZE: tl.constexpr, # int, must be power of 2

333 HEAD_SIZE: tl.constexpr, # int

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

335 USE_ALIBI_SLOPES: tl.constexpr, # bool

@@ -374,20 +377,19 @@ def kernel_unified_attention_3d(

377

378 # number of segments for this particular sequence

379 num_segments = NUM_SEGMENTS_PER_SEQ

377 - blocks_per_segment = cdiv_fn(seq_len, num_segments * BLOCK_SIZE)

380 + tiles_per_segment = cdiv_fn(seq_len, num_segments * TILE_SIZE)

381

379 - if segm_idx * blocks_per_segment * BLOCK_SIZE >= seq_len:

382 + if segm_idx * tiles_per_segment * TILE_SIZE >= seq_len:

383 return

384

385 offs_m = tl.arange(0, BLOCK_M)

386 offs_d = tl.arange(0, HEAD_SIZE_PADDED)

384 -

387 + offs_t = tl.arange(0, TILE_SIZE)

388 query_pos = q_block_local_idx * BLOCK_Q + offs_m // num_queries_per_kv

389

390 query_offset_0 = cur_batch_in_all_start_index + query_pos

391 query_offset_1 = kv_head_idx * num_queries_per_kv + \

392 offs_m % num_queries_per_kv

390 -

393 query_offset = (query_offset_0[:, None] * query_stride_0 +

394 query_offset_1[:, None] * query_stride_1 + offs_d[None, :])

395

@@ -433,30 +435,44 @@ def kernel_unified_attention_3d(

435 qq_bias_row_ptrs = (qq_bias_ptr + query_pos[:, None] * qq_bias_stride_0

436 ) # shape: [BLOCK_M]

437

436 - num_blocks = cdiv_fn(seq_len, BLOCK_SIZE)

438 + # compute the length of the longest sequence prefix spanned by any

439 + # query token in the current q_block (q_block_local_idx)

440 + max_seq_prefix_len = context_len + q_block_local_idx * BLOCK_Q + (

441 + BLOCK_M - 1) // num_queries_per_kv + 1

442 +

443 + # adjust for potential padding in the last q_block by considering the

444 + # actual sequence length

445 + max_seq_prefix_len = tl.minimum(max_seq_prefix_len, seq_len)

446 +

447 + # calculate the number of tiles that need to be processed to

448 + # cover the longest sequence prefix (due to causal masking, tiles beyond

449 + # this prefix can be skipped)

450 + num_tiles = cdiv_fn(max_seq_prefix_len, TILE_SIZE)

451

452 # iterate through tiles within current segment

453 for j in range(

440 - segm_idx * blocks_per_segment,

441 - min((segm_idx + 1) * blocks_per_segment, num_blocks),

454 + segm_idx * tiles_per_segment,

455 + min((segm_idx + 1) * tiles_per_segment, num_tiles),

456 ):

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

457 + seq_offset = j * TILE_SIZE + offs_t

458 + tile_mask = seq_offset < max_seq_prefix_len

459

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

460 + physical_block_idx = tl.load(block_tables_ptr + block_table_offset +

461 + seq_offset // BLOCK_SIZE).to(tl.int64)

462

447 - v_offset = (physical_block_idx * stride_v_cache_0 +

463 + v_offset = (physical_block_idx[:, None] * stride_v_cache_0 +

464 kv_head_idx * stride_v_cache_2 +

465 offs_d[None, :] * stride_v_cache_3 +

450 - offs_n[:, None] * stride_v_cache_1)

466 + (seq_offset % BLOCK_SIZE)[:, None] * stride_v_cache_1)

467

452 - k_offset = (physical_block_idx * stride_k_cache_0 +

468 + k_offset = (physical_block_idx[None, :] * stride_k_cache_0 +

469 kv_head_idx * stride_k_cache_2 +

470 offs_d[:, None] * stride_k_cache_3 +

455 - offs_n[None, :] * stride_k_cache_1)

471 + (seq_offset % BLOCK_SIZE)[None, :] * stride_k_cache_1)

472

457 - # K : (HEAD_SIZE, BLOCK_SIZE)

473 + # K : (HEAD_SIZE, TILE_SIZE)

474 K_load = tl.load(key_cache_ptr + k_offset,

459 - mask=dim_mask[:, None],

475 + mask=dim_mask[:, None] & tile_mask[None, :],

476 other=0.0)

477

478 if K_load.dtype.is_fp8():

@@ -467,9 +483,9 @@ def kernel_unified_attention_3d(

483 else:

484 K = K_load

485

470 - # V : (BLOCK_SIZE, HEAD_SIZE)

486 + # V : (TILE_SIZE, HEAD_SIZE)

487 V_load = tl.load(value_cache_ptr + v_offset,

472 - mask=dim_mask[None, :],

488 + mask=dim_mask[None, :] & tile_mask[:, None],

489 other=0.0)

490

491 if V_load.dtype.is_fp8():

@@ -480,13 +496,10 @@ def kernel_unified_attention_3d(

496 else:

497 V = V_load

498

483 - seq_offset = j * BLOCK_SIZE + offs_n

484 -

499 seq_mask = seq_offset[None, :] < context_len + query_pos[:, None] + 1

500

487 - # S : (BLOCK_M, BLOCK_SIZE)

488 - S = tl.zeros(shape=(BLOCK_M, BLOCK_SIZE), dtype=tl.float32)

489 -

501 + # S : (BLOCK_M, TILE_SIZE)

502 + S = tl.zeros(shape=(BLOCK_M, TILE_SIZE), dtype=tl.float32)

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

504

505 if USE_SOFTCAP:

@@ -517,11 +530,12 @@ def kernel_unified_attention_3d(

530 # compute running maximum

531 # m_j : (BLOCK_M,)

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

533 +

534 # For sliding window there's a chance the max is -inf due to masking of

535 # the entire row. In this case we need to set m_j 0 to avoid NaN

536 m_j = tl.where(m_j > float("-inf"), m_j, 0.0)

537

524 - # P : (BLOCK_M, BLOCK_SIZE,)

538 + # P : (BLOCK_M, TILE_SIZE,)

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

540

541 # l_j : (BLOCK_M,)

@@ -573,7 +587,7 @@ def reduce_segments(

587 output_stride_0: tl.int64, # int

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

589 block_table_stride: tl.int64, # int

576 - BLOCK_SIZE: tl.constexpr, # int

590 + TILE_SIZE: tl.constexpr, # int

591 HEAD_SIZE: tl.constexpr, # int, must be power of 2

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

593 query_start_len_ptr, # [num_seqs+1]

@@ -594,10 +608,10 @@ def reduce_segments(

608

609 # number of segments for this particular sequence

610 num_segments = NUM_SEGMENTS_PER_SEQ

597 - blocks_per_segment = cdiv_fn(seq_len, num_segments * BLOCK_SIZE)

611 + tiles_per_segment = cdiv_fn(seq_len, num_segments * TILE_SIZE)

612

613 # create masks for subsequent loads

600 - act_num_segments = cdiv_fn(seq_len, blocks_per_segment * BLOCK_SIZE)

614 + act_num_segments = cdiv_fn(seq_len, tiles_per_segment * TILE_SIZE)

615 segm_mask = tl.arange(0, NUM_SEGMENTS_PER_SEQ) < tl.full(

616 [NUM_SEGMENTS_PER_SEQ], act_num_segments, dtype=tl.int32)

617 dim_mask = tl.where(tl.arange(0, HEAD_SIZE_PADDED) < HEAD_SIZE, 1,

@@ -671,13 +685,10 @@ def unified_attention(

685 # Optional tensor for sinks

686 sinks=None,

687 ):

688 +

689 assert causal, "Only causal attention is supported"

690 assert q_descale is None, "Q scales not supported"

691

677 - block_size = v.shape[1]

678 - assert q.element_size() >= 2 or block_size >= 32, \

679 - "Block size must be at least 32 for fp8"

680 -

692 if sinks is not None:

693 assert sinks.shape[0] == q.shape[1], \

694 "Sinks must be num_query_heads size"

@@ -707,6 +718,12 @@ def unified_attention(

718 # = floor(q.shape[0] / BLOCK_Q) + num_seqs

719 total_num_q_blocks = q.shape[0] // BLOCK_Q + num_seqs

720

721 + # Assigning default tile sizes for prefill and decode.

722 + # Note: each tile size must be at least 32 for "fp8" (q.element_size() == 1)

723 + # and at least 16 for all other data types.

724 + TILE_SIZE_PREFILL = 32

725 + TILE_SIZE_DECODE = 16 if q.element_size() >= 2 else 32

726 +

727 # if batch contains a prefill

728 if max_seqlen_q > 1 or total_num_q_blocks * num_kv_heads > 128:

729 kernel_unified_attention_2d[(

@@ -736,6 +753,7 @@ def unified_attention(

753 output_stride_1=out.stride(1),

754 qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,

755 BLOCK_SIZE=block_size,

756 + TILE_SIZE=TILE_SIZE_PREFILL,

757 HEAD_SIZE=head_size,

758 HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),

759 USE_ALIBI_SLOPES=use_alibi_slopes,

@@ -809,6 +827,7 @@ def unified_attention(

827 query_stride_1=q.stride(1),

828 qq_bias_stride_0=qq_bias.stride(0) if use_qq_bias else 0,

829 BLOCK_SIZE=block_size,

830 + TILE_SIZE=TILE_SIZE_DECODE,

831 HEAD_SIZE=head_size,

832 HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),

833 USE_ALIBI_SLOPES=use_alibi_slopes,

@@ -830,7 +849,6 @@ def unified_attention(

849 BLOCK_M=BLOCK_M,

850 NUM_SEGMENTS_PER_SEQ=NUM_SEGMENTS,

851 )

833 -

852 reduce_segments[(q.shape[0], num_query_heads)](

853 output_ptr=out,

854 segm_output_ptr=segm_output,

@@ -844,7 +862,7 @@ def unified_attention(

862 output_stride_0=out.stride(0),

863 output_stride_1=out.stride(1),

864 block_table_stride=block_table.stride(0),

847 - BLOCK_SIZE=block_size,

865 + TILE_SIZE=TILE_SIZE_DECODE,

866 HEAD_SIZE=head_size,

867 HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),

868 query_start_len_ptr=cu_seqlens_q,