@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6e20924350

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

Add attention sink in attention backends (#22320)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu> Co-authored-by: LiuXiaoxuanPKU <lilyliupku@gmail.com> Co-authored-by: simon-mo <xmo@berkeley.edu> Co-authored-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Hongxia Yang <62075498+hongxiayang@users.noreply.github.com> Co-authored-by: Minseok Lee <47620120+minseokl@users.noreply.github.com> Co-authored-by: Yongye Zhu <zyy1102000@gmail.com>

Woosuk Kwon committed Aug 5, 2025 at 22:37 UTC 6e20924350e3fed375bc63d55166a303b6f0828a

7 files changed +176 -45

vllm/attention/ops/chunked_prefill_paged_decode.py

+27 -6

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

28 query_ptr, # [num_tokens, num_query_heads, head_size]

29 key_cache_ptr, # [num_blks, num_kv_heads, head_size // x, blk_size, x]

30 value_cache_ptr, # [num_blks, num_kv_heads, head_size, blk_size]

31 + sink_ptr, # [num_query_heads]

32 block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]

33 seq_lens_ptr, # [num_seqs]

34 alibi_slopes_ptr, # [num_query_heads]

@@ -95,7 +96,17 @@ def kernel_paged_attention_2d(

97 block_table_offset = seq_idx * block_table_stride

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

99 + if sink_ptr is None:

100 + M = tl.full([num_queries_per_kv_padded],

101 + float("-inf"),

102 + dtype=tl.float32)

103 + else:

104 + M = tl.load(

105 + sink_ptr + query_head_idx,

106 + mask=head_mask,

107 + other=float("-inf"),

108 + ).to(dtype=tl.float32)

109 +

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

111 acc = tl.zeros([num_queries_per_kv_padded, HEAD_SIZE_PADDED],

112 dtype=tl.float32)

@@ -223,6 +234,8 @@ def chunked_prefill_paged_decode(

234 alibi_slopes=None,

235 sliding_window=None,

236 sm_scale=None,

237 + # Optional tensor for sinks

238 + sinks=None,

239 ):

240

241 if sm_scale is None:

@@ -253,6 +266,7 @@ def chunked_prefill_paged_decode(

266 sliding_window=sliding_window,

267 sm_scale=sm_scale,

268 skip_decode=True,

269 + sinks=sinks,

270 )

271

272 block_size = value_cache.shape[3]

@@ -281,11 +295,17 @@ def chunked_prefill_paged_decode(

295 num_queries_per_kv_padded = max(triton.next_power_of_2(num_queries_per_kv),

296 16)

297

284 - use_custom = use_rocm_custom_paged_attention(query.dtype, head_size,

285 - block_size,

286 - num_queries_per_kv,

287 - max_seq_len, sliding_window,

288 - kv_cache_dtype, alibi_slopes)

298 + use_custom = use_rocm_custom_paged_attention(

299 + query.dtype,

300 + head_size,

301 + block_size,

302 + num_queries_per_kv,

303 + max_seq_len,

304 + sliding_window,

305 + kv_cache_dtype,

306 + alibi_slopes,

307 + sinks,

308 + )

309 if use_custom:

310 _PARTITION_SIZE_ROCM = 256

311 max_num_partitions = ((max_seq_len + _PARTITION_SIZE_ROCM - 1) //

@@ -334,6 +354,7 @@ def chunked_prefill_paged_decode(

354 query_ptr=query,

355 key_cache_ptr=key_cache,

356 value_cache_ptr=value_cache,

357 + sink_ptr=sinks,

358 block_tables_ptr=block_table,

359 seq_lens_ptr=seq_lens,

360 alibi_slopes_ptr=alibi_slopes,

vllm/attention/ops/prefix_prefill.py

+15 -3

@@ -38,6 +38,7 @@ def _fwd_kernel(Q,

38 V,

39 K_cache,

40 V_cache,

41 + sink_ptr,

42 B_Loc,

43 sm_scale,

44 k_scale,

@@ -126,7 +127,15 @@ def _fwd_kernel(Q,

127 other=0.0) # [M,D]

128

129 # initialize pointer to m and l

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

130 + if sink_ptr is None:

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

132 + else:

133 + m_i = tl.load(

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

135 + mask=(offs_m < cur_batch_query_len),

136 + other=float("-inf"),

137 + ).to(dtype=tl.float32)

138 +

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

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

141

@@ -732,7 +741,8 @@ def context_attention_fwd(q,

741 alibi_slopes=None,

742 sliding_window=None,

743 sm_scale=None,

735 - skip_decode=False):

744 + skip_decode=False,

745 + sinks=None):

746

747 q_dtype_is_f32 = q.dtype is torch.float32

748

@@ -781,6 +791,7 @@ def context_attention_fwd(q,

791 sliding_window = 0

792

793 if alibi_slopes is not None:

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

795 # need to reduce num. blocks when using fp32

796 # due to increased use of GPU shared memory

797 # if q.dtype is torch.float32:

@@ -843,7 +854,7 @@ def context_attention_fwd(q,

854 max_seq_len = 0 if max_seq_len is None else max_seq_len

855 extra_kargs = {}

856 if current_platform.is_rocm():

846 - extra_kargs = {"kpack": 2, "waves_per_eu": 2}

857 + extra_kargs = {"kpack": 1, "waves_per_eu": 2}

858

859 grid = lambda META: (batch, head,

860 triton.cdiv(max_input_len, META["BLOCK_M"]))

@@ -853,6 +864,7 @@ def context_attention_fwd(q,

864 v,

865 k_cache,

866 v_cache,

867 + sinks,

868 b_loc,

869 sm_scale,

870 k_scale,

vllm/attention/ops/triton_unified_attention.py

+28 -2

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

52 query_ptr, # [num_tokens, num_query_heads, head_size]

53 key_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]

54 value_cache_ptr, # [num_blks, blk_size, num_kv_heads, head_size]

55 + sink_ptr, # [num_query_heads]

56 block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]

57 seq_lens_ptr, # [num_seqs]

58 alibi_slopes_ptr, # [num_query_heads]

@@ -131,7 +132,15 @@ def kernel_unified_attention_2d(

132

133 block_table_offset = seq_idx * block_table_stride

134

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

135 + if sink_ptr is None:

136 + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)

137 + else:

138 + M = tl.load(

139 + sink_ptr + query_offset_1,

140 + mask=query_mask_1,

141 + other=float("-inf"),

142 + ).to(dtype=tl.float32)

143 +

144 L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)

145 acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)

146

@@ -292,6 +301,7 @@ def kernel_unified_attention_3d(

301 query_ptr, # [num_tokens, num_query_heads, head_size]

302 key_cache_ptr, # [num_blks, num_kv_heads, head_size // x, blk_size, x]

303 value_cache_ptr, # [num_blks, num_kv_heads, head_size, blk_size]

304 + sink_ptr, # [num_query_heads]

305 block_tables_ptr, # [num_seqs, max_num_blocks_per_seq]

306 seq_lens_ptr, # [num_seqs]

307 alibi_slopes_ptr, # [num_query_heads]

@@ -383,7 +393,15 @@ def kernel_unified_attention_3d(

393

394 block_table_offset = seq_idx * block_table_stride

395

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

396 + if sink_ptr is None or segm_idx != 0:

397 + M = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32)

398 + else:

399 + M = tl.load(

400 + sink_ptr + query_offset_1,

401 + mask=query_mask_1,

402 + other=float("-inf"),

403 + ).to(dtype=tl.float32)

404 +

405 L = tl.full([BLOCK_M], 1.0, dtype=tl.float32)

406 acc = tl.zeros([BLOCK_M, HEAD_SIZE_PADDED], dtype=tl.float32)

407

@@ -627,6 +645,8 @@ def unified_attention(

645 v_descale,

646 alibi_slopes=None,

647 qq_bias=None,

648 + # Optional tensor for sinks

649 + sinks=None,

650 ):

651 assert causal, "Only causal attention is supported"

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

@@ -635,6 +655,10 @@ def unified_attention(

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

656 "Block size must be at least 32 for fp8"

657

658 + if sinks is not None:

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

660 + "Sinks must be num_query_heads size"

661 +

662 use_alibi_slopes = alibi_slopes is not None

663 use_qq_bias = qq_bias is not None

664

@@ -669,6 +693,7 @@ def unified_attention(

693 query_ptr=q,

694 key_cache_ptr=k,

695 value_cache_ptr=v,

696 + sink_ptr=sinks,

697 block_tables_ptr=block_table,

698 seq_lens_ptr=seqused_k,

699 alibi_slopes_ptr=alibi_slopes,

@@ -741,6 +766,7 @@ def unified_attention(

766 query_ptr=q,

767 key_cache_ptr=k,

768 value_cache_ptr=v,

769 + sink_ptr=sinks,

770 block_tables_ptr=block_table,

771 seq_lens_ptr=seqused_k,

772 alibi_slopes_ptr=alibi_slopes,

vllm/envs.py

+16 -3

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

17 LD_LIBRARY_PATH: Optional[str] = None

18 VLLM_USE_TRITON_FLASH_ATTN: bool = True

19 VLLM_V1_USE_PREFILL_DECODE_ATTENTION: bool = False

20 + VLLM_USE_AITER_UNIFIED_ATTENTION: bool = False

21 VLLM_FLASH_ATTN_VERSION: Optional[int] = None

22 LOCAL_RANK: int = 0

23 CUDA_VISIBLE_DEVICES: Optional[str] = None

@@ -151,6 +152,8 @@ if TYPE_CHECKING:

152 VLLM_LOOPBACK_IP: str = ""

153 VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = False

154 VLLM_ENABLE_RESPONSES_API_STORE: bool = False

155 + VLLM_USE_TRTLLM_CONTEXT_ATTENTION: bool = False

156 + VLLM_USE_TRTLLM_DECODE_ATTENTION: bool = False

157

158

159 def get_default_cache_root():

@@ -326,6 +329,12 @@ environment_variables: dict[str, Callable[[], Any]] = {

329 (os.getenv("VLLM_V1_USE_PREFILL_DECODE_ATTENTION", "False").lower() in

330 ("true", "1")),

331

332 + # Use AITER triton unified attention for V1 attention

333 + "VLLM_USE_AITER_UNIFIED_ATTENTION":

334 + lambda:

335 + (os.getenv("VLLM_USE_AITER_UNIFIED_ATTENTION", "False").lower() in

336 + ("true", "1")),

337 +

338 # Force vllm to use a specific flash-attention version (2 or 3), only valid

339 # when using the flash-attention backend.

340 "VLLM_FLASH_ATTN_VERSION":

@@ -1022,9 +1031,13 @@ environment_variables: dict[str, Callable[[], Any]] = {

1031 "VLLM_USE_CUDNN_PREFILL":

1032 lambda: bool(int(os.getenv("VLLM_USE_CUDNN_PREFILL", "0"))),

1033

1025 - # If set to 1, use the TRTLLM Attention backend in flashinfer.

1026 - "VLLM_USE_TRTLLM_ATTENTION":

1027 - lambda: os.getenv("VLLM_USE_TRTLLM_ATTENTION", None),

1034 + # If set to 1, use the TRTLLM Context Attention backend in flashinfer.

1035 + "VLLM_USE_TRTLLM_CONTEXT_ATTENTION":

1036 + lambda: bool(int(os.getenv("VLLM_USE_TRTLLM_CONTEXT_ATTENTION", "0"))),

1037 +

1038 + # If set to 1, use the TRTLLM Decode Attention backend in flashinfer.

1039 + "VLLM_USE_TRTLLM_DECODE_ATTENTION":

1040 + lambda: bool(int(os.getenv("VLLM_USE_TRTLLM_DECODE_ATTENTION", "0"))),

1041

1042 # Controls garbage collection during CUDA graph capture.

1043 # If set to 0 (default), enables GC freezing to speed up capture time.

vllm/v1/attention/backends/flash_attn.py

+10

@@ -373,6 +373,7 @@ class FlashAttentionImpl(AttentionImpl):

373 logits_soft_cap: Optional[float] = None,

374 attn_type: AttentionType = AttentionType.DECODER,

375 kv_sharing_target_layer_name: Optional[str] = None,

376 + sinks: Optional[torch.Tensor] = None,

377 ) -> None:

378 self.num_heads = num_heads

379 self.head_size = head_size

@@ -410,6 +411,14 @@ class FlashAttentionImpl(AttentionImpl):

411 raise NotImplementedError(

412 "FlashAttention does not support fp8 kv-cache on this device.")

413

414 + self.sinks = sinks

415 + if self.sinks is not None:

416 + assert self.vllm_flash_attn_version == 3, (

417 + "Sinks are only supported in FlashAttention 3")

418 + assert self.sinks.shape[0] == num_heads, (

419 + "Sinks must have the same number of heads as the number of "

420 + "heads in the layer")

421 +

422 def forward(

423 self,

424 layer: torch.nn.Module,

@@ -534,6 +543,7 @@ class FlashAttentionImpl(AttentionImpl):

543 k_descale=layer._k_scale.expand(descale_shape),

544 v_descale=layer._v_scale.expand(descale_shape),

545 num_splits=attn_metadata.max_num_splits,

546 + s_aux=self.sinks,

547 )

548 return output

549

vllm/v1/attention/backends/triton_attn.py

+56 -19

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """Attention layer with PagedAttention and Triton prefix prefill."""

4 from dataclasses import dataclass

5 +from functools import cache

6 from typing import ClassVar, Optional

8 import torch

@@ -13,7 +14,6 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,

14 from vllm.attention.ops.chunked_prefill_paged_decode import (

15 chunked_prefill_paged_decode)

16 from vllm.attention.ops.paged_attn import PagedAttention

16 -from vllm.attention.ops.triton_unified_attention import unified_attention

17 from vllm.config import VllmConfig

18 from vllm.logger import init_logger

19 from vllm.platforms import current_platform

@@ -193,6 +193,15 @@ class TritonAttentionBackend(AttentionBackend):

193 return TritonAttentionMetadataBuilder

194

195

196 +@cache

197 +def use_aiter_unified_attention() -> bool:

198 + """Check if aiter unified attention should be used."""

199 + # VLLM_ROCM_USE_AITER_MHA needs to set to 0 as well as it is set

200 + # to 1 as default

201 + return envs.VLLM_ROCM_USE_AITER \

202 + and envs.VLLM_USE_AITER_UNIFIED_ATTENTION

203 +

204 +

205 class TritonAttentionImpl(AttentionImpl):

206

207 def __init__(

@@ -207,6 +216,7 @@ class TritonAttentionImpl(AttentionImpl):

216 logits_soft_cap: Optional[float] = None,

217 attn_type: AttentionType = AttentionType.DECODER,

218 kv_sharing_target_layer_name: Optional[int] = None,

219 + sinks: Optional[torch.Tensor] = None,

220 ) -> None:

221 self.num_heads = num_heads

222 self.head_size = head_size

@@ -240,6 +250,29 @@ class TritonAttentionImpl(AttentionImpl):

250 self.force_prefill_decode_attn = \

251 envs.VLLM_V1_USE_PREFILL_DECODE_ATTENTION

252

253 + if not self.force_prefill_decode_attn:

254 + # If not using prefill decode attention, we use the Triton

255 + # unified attention implementation.

256 + if use_aiter_unified_attention():

257 + logger.info_once(

258 + "Using aiter unified attention for TritonAttentionImpl")

259 + from aiter.ops.triton.unified_attention import (

260 + unified_attention)

261 + self.unified_attention = unified_attention

262 + else:

263 + logger.info_once(

264 + "Using vllm unified attention for TritonAttentionImpl")

265 + from vllm.attention.ops.triton_unified_attention import (

266 + unified_attention)

267 + self.unified_attention = unified_attention

268 +

269 + self.sinks = sinks

270 + if sinks is not None:

271 + assert sinks.shape[0] == num_heads, (

272 + "Sinks must have the same number of heads as the number of "

273 + f"heads in the layer. Sinks shape: {sinks.shape}, "

274 + f"num_heads: {num_heads}.")

275 +

276 def forward(

277 self,

278 layer: torch.nn.Module,

@@ -342,28 +375,31 @@ class TritonAttentionImpl(AttentionImpl):

375

376 if use_prefill_decode_attn:

377 # Compute attention and update output up to `num_actual_tokens`.

345 - chunked_prefill_paged_decode(query=query[:num_actual_tokens],

346 - key=key[:num_actual_tokens],

347 - value=value[:num_actual_tokens],

348 - output=output[:num_actual_tokens],

349 - kv_cache_dtype=self.kv_cache_dtype,

350 - key_cache=key_cache,

351 - value_cache=value_cache,

352 - block_table=block_table,

353 - query_start_loc=cu_seqlens_q,

354 - seq_lens=seqused_k,

355 - max_seq_len=max_seqlen_k,

356 - max_query_len=max_seqlen_q,

357 - k_scale=layer._k_scale,

358 - v_scale=layer._v_scale,

359 - alibi_slopes=self.alibi_slopes,

360 - sliding_window=self.sliding_window[0],

361 - sm_scale=self.scale)

378 + chunked_prefill_paged_decode(

379 + query=query[:num_actual_tokens],

380 + key=key[:num_actual_tokens],

381 + value=value[:num_actual_tokens],

382 + output=output[:num_actual_tokens],

383 + kv_cache_dtype=self.kv_cache_dtype,

384 + key_cache=key_cache,

385 + value_cache=value_cache,

386 + block_table=block_table,

387 + query_start_loc=cu_seqlens_q,

388 + seq_lens=seqused_k,

389 + max_seq_len=max_seqlen_k,

390 + max_query_len=max_seqlen_q,

391 + k_scale=layer._k_scale,

392 + v_scale=layer._v_scale,

393 + alibi_slopes=self.alibi_slopes,

394 + sliding_window=self.sliding_window[0],

395 + sm_scale=self.scale,

396 + sinks=self.sinks,

397 + )

398

399 else:

400 descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])

401

366 - unified_attention(

402 + self.unified_attention(

403 q=query[:num_actual_tokens],

404 k=key_cache,

405 v=value_cache,

@@ -381,6 +417,7 @@ class TritonAttentionImpl(AttentionImpl):

417 q_descale=None, # Not supported

418 k_descale=layer._k_scale.expand(descale_shape),

419 v_descale=layer._v_scale.expand(descale_shape),

420 + sinks=self.sinks,

421 )

422

423 return output

vllm/v1/attention/backends/utils.py

+24 -12

@@ -254,7 +254,11 @@ def get_kv_cache_layout():

254 # Override with format specified by the user.

255 cache_layout = envs.VLLM_KV_CACHE_LAYOUT

256 if cache_layout is None:

257 - cache_layout = get_kv_connector_cache_layout()

257 + if (envs.VLLM_USE_TRTLLM_CONTEXT_ATTENTION

258 + or envs.VLLM_USE_TRTLLM_DECODE_ATTENTION):

259 + cache_layout = "HND"

260 + else:

261 + cache_layout = get_kv_connector_cache_layout()

262 else:

263 logger.info_once("`VLLM_KV_CACHE_LAYOUT` environment variable " \

264 "detected. Setting KV cache layout to %s.", cache_layout)

@@ -272,7 +276,9 @@ def set_kv_cache_layout(cache_layout: str):

276 class PerLayerParameters:

277 """

278 Currently, FlashInfer backend only support models in which all layers share

275 - the same values for the following hyperparameters.

279 + the same values for the following hyperparameters. Should not be used for

280 + trtllm-gen backend since it supports different values for the following

281 + hyperparameters.

282 """

283

284 window_left: int

@@ -310,7 +316,8 @@ def get_per_layer_parameters(

316 def infer_global_hyperparameters(

317 per_layer_params: dict[str, PerLayerParameters]) -> PerLayerParameters:

318 """

313 - Currently, FlashInfer backend only support models in which all layers share

319 + Currently, FlashInfer backend other than trtllm-gen

320 + only support models in which all layers share

321 the same values for the following hyperparameters:

322 - `window_left`

323 - `logits_soft_cap`

@@ -324,15 +331,20 @@ def infer_global_hyperparameters(

331

332 param_sets = list(per_layer_params.values())

333 global_params = param_sets[0]

327 - for params in param_sets:

328 - if params.window_left != global_params.window_left:

329 - raise ValueError(

330 - "Window left is not the same for all layers. One potential fix "

331 - "is to set disable_sliding_window=True")

332 - assert params == global_params, (

333 - "FlashInfer backend currently only supports models in which all "

334 - "layers share the same values for the following hyperparameters: "

335 - "`window_left`, `logits_soft_cap`, `sm_scale`.")

334 +

335 + # trtllm attention doesn't need global hyper params so disable the check

336 + if (not envs.VLLM_USE_TRTLLM_CONTEXT_ATTENTION

337 + and not envs.VLLM_USE_TRTLLM_DECODE_ATTENTION):

338 + for params in param_sets:

339 + if params.window_left != global_params.window_left:

340 + raise ValueError(

341 + "Window left is not the same for all layers. " \

342 + "One potential fix is to set disable_sliding_window=True")

343 + assert params == global_params, (

344 + "FlashInfer backend currently only supports models in which all"

345 + "layers share the same values "

346 + "for the following hyperparameters:"

347 + "`window_left`, `logits_soft_cap`, `sm_scale`.")

348

349 return global_params

350