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 11, 13:45 UTC
[Kernel] Support CUDA Graphs in 3D Triton Attention Kernel (#28306)
Signed-off-by: Jan van Lunteren <jvl@zurich.ibm.com> Signed-off-by: jvlunteren <161835099+jvlunteren@users.noreply.github.com> Co-authored-by: Thomas Parnell <tom.parnell@gmail.com> Co-authored-by: Thomas Parnell <tpa@zurich.ibm.com>
jvlunteren committed
Dec 12, 2025 at 16:55 UTC
9c0ee995a81fbd87b397c956ca56fc94f784966e
3 files changed
+140
-40
tests/kernels/attention/test_triton_unified_attention.py
+27
@@ -7,6 +7,7 @@ import torch
7
8
from vllm.attention.ops.triton_unified_attention import unified_attention
9
from vllm.platforms import current_platform
10
+from vllm.utils.math_utils import next_power_of_2
11
12
NUM_HEADS = [(4, 4), (8, 2)]
13
HEAD_SIZES = [128, 256]
@@ -22,6 +23,10 @@ QDTYPES = (
23
# one value small enough to test the schema op check
24
NUM_BLOCKS = [32768, 2048]
25
26
+# 0: use 2D kernel for decode
27
+# 8: use 3D kernel for decode
28
+SEQ_THRESHOLD_3D_VALUES = [0, 8]
29
+
30
31
def ref_paged_attn(
32
query: torch.Tensor,
@@ -92,6 +97,7 @@ def ref_paged_attn(
97
@pytest.mark.parametrize("soft_cap", [None, 50.0])
98
@pytest.mark.parametrize("num_blocks", NUM_BLOCKS)
99
@pytest.mark.parametrize("q_dtype", QDTYPES)
100
+@pytest.mark.parametrize("seq_threshold_3D", SEQ_THRESHOLD_3D_VALUES)
101
@torch.inference_mode()
102
def test_triton_unified_attn(
103
seq_lens: list[tuple[int, int]],
@@ -103,6 +109,7 @@ def test_triton_unified_attn(
109
soft_cap: float | None,
110
num_blocks: int,
111
q_dtype: torch.dtype | None,
112
+ seq_threshold_3D: int,
113
) -> None:
114
torch.set_default_device("cuda")
115
@@ -152,6 +159,21 @@ def test_triton_unified_attn(
159
k_descale = torch.rand(scale_shape, dtype=torch.float32)
160
v_descale = torch.rand(scale_shape, dtype=torch.float32)
161
162
+ num_par_softmax_segments = 16
163
+ head_size_padded = next_power_of_2(head_size)
164
+ softmax_segm_output = torch.empty(
165
+ (seq_threshold_3D, num_query_heads, num_par_softmax_segments, head_size_padded),
166
+ dtype=torch.float32,
167
+ )
168
+ softmax_segm_max = torch.empty(
169
+ (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
170
+ dtype=torch.float32,
171
+ )
172
+ softmax_segm_expsum = torch.empty(
173
+ (seq_threshold_3D, num_query_heads, num_par_softmax_segments),
174
+ dtype=torch.float32,
175
+ )
176
+
177
unified_attention(
178
q=maybe_quantized_query,
179
k=maybe_quantized_key_cache,
@@ -169,6 +191,11 @@ def test_triton_unified_attn(
191
q_descale=q_descale,
192
k_descale=k_descale,
193
v_descale=v_descale,
194
+ seq_threshold_3D=seq_threshold_3D,
195
+ num_par_softmax_segments=num_par_softmax_segments,
196
+ softmax_segm_output=softmax_segm_output,
197
+ softmax_segm_max=softmax_segm_max,
198
+ softmax_segm_expsum=softmax_segm_expsum,
199
)
200
201
ref_output = ref_paged_attn(
vllm/attention/ops/triton_unified_attention.py
+30
-39
@@ -355,7 +355,7 @@ def kernel_unified_attention_2d(
355
@triton.jit
356
def kernel_unified_attention_3d(
357
segm_output_ptr,
358
- # [num_tokens, num_query_heads, num_segments, head_size]
358
+ # [num_tokens, num_query_heads, num_segments, head_size_padded]
359
segm_max_ptr, # [num_tokens, num_query_heads, num_segments]
360
segm_expsum_ptr, # [num_tokens, num_query_heads, num_segments]
361
query_ptr, # [num_tokens, num_query_heads, head_size]
@@ -749,6 +749,11 @@ def unified_attention(
749
q_descale,
750
k_descale,
751
v_descale,
752
+ seq_threshold_3D=None,
753
+ num_par_softmax_segments=None,
754
+ softmax_segm_output=None,
755
+ softmax_segm_max=None,
756
+ softmax_segm_expsum=None,
757
alibi_slopes=None,
758
output_scale=None,
759
qq_bias=None,
@@ -793,8 +798,19 @@ def unified_attention(
798
TILE_SIZE_PREFILL = 32
799
TILE_SIZE_DECODE = 16 if q.element_size() >= 2 else 32
800
796
- # if batch contains a prefill
797
- if max_seqlen_q > 1 or total_num_q_blocks * num_kv_heads > 128:
801
+ # Launch the 2D kernel if
802
+ # 1. No intermediate tiled softmax buffers for the 3D kernel have been allocated, or
803
+ # 2. The batch includes at least one prefill request, or
804
+ # 3. The number of sequences exceeds the configured threshold
805
+ if (
806
+ seq_threshold_3D is None
807
+ or num_par_softmax_segments is None
808
+ or softmax_segm_output is None
809
+ or softmax_segm_max is None
810
+ or softmax_segm_expsum is None
811
+ or max_seqlen_q > 1
812
+ or num_seqs > seq_threshold_3D
813
+ ):
814
kernel_unified_attention_2d[
815
(
816
total_num_q_blocks,
@@ -847,37 +863,12 @@ def unified_attention(
863
USE_FP8=output_scale is not None,
864
)
865
else:
850
- # for initial version, NUM_SEGMENTS = 16 is chosen as a default
851
- # value that showed good performance in tests
852
- NUM_SEGMENTS = 16
853
-
854
- segm_output = torch.empty(
855
- q.shape[0],
856
- num_query_heads,
857
- NUM_SEGMENTS,
858
- triton.next_power_of_2(head_size),
859
- dtype=torch.float32,
860
- device=q.device,
861
- )
862
- segm_max = torch.empty(
863
- q.shape[0],
864
- num_query_heads,
865
- NUM_SEGMENTS,
866
- dtype=torch.float32,
867
- device=q.device,
868
- )
869
- segm_expsum = torch.empty(
870
- q.shape[0],
871
- num_query_heads,
872
- NUM_SEGMENTS,
873
- dtype=torch.float32,
874
- device=q.device,
875
- )
876
-
877
- kernel_unified_attention_3d[(total_num_q_blocks, num_kv_heads, NUM_SEGMENTS)](
878
- segm_output_ptr=segm_output,
879
- segm_max_ptr=segm_max,
880
- segm_expsum_ptr=segm_expsum,
866
+ kernel_unified_attention_3d[
867
+ (total_num_q_blocks, num_kv_heads, num_par_softmax_segments)
868
+ ](
869
+ segm_output_ptr=softmax_segm_output,
870
+ segm_max_ptr=softmax_segm_max,
871
+ segm_expsum_ptr=softmax_segm_expsum,
872
query_ptr=q,
873
key_cache_ptr=k,
874
value_cache_ptr=v,
@@ -917,13 +908,13 @@ def unified_attention(
908
BLOCK_Q=BLOCK_Q,
909
num_seqs=num_seqs,
910
BLOCK_M=BLOCK_M,
920
- NUM_SEGMENTS_PER_SEQ=NUM_SEGMENTS,
911
+ NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,
912
)
913
reduce_segments[(q.shape[0], num_query_heads)](
914
output_ptr=out,
924
- segm_output_ptr=segm_output,
925
- segm_max_ptr=segm_max,
926
- segm_expsum_ptr=segm_expsum,
915
+ segm_output_ptr=softmax_segm_output,
916
+ segm_max_ptr=softmax_segm_max,
917
+ segm_expsum_ptr=softmax_segm_expsum,
918
seq_lens_ptr=seqused_k,
919
num_seqs=num_seqs,
920
num_query_heads=num_query_heads,
@@ -936,6 +927,6 @@ def unified_attention(
927
HEAD_SIZE_PADDED=triton.next_power_of_2(head_size),
928
query_start_len_ptr=cu_seqlens_q,
929
BLOCK_Q=BLOCK_Q,
939
- NUM_SEGMENTS_PER_SEQ=NUM_SEGMENTS,
930
+ NUM_SEGMENTS_PER_SEQ=num_par_softmax_segments,
931
USE_FP8=output_scale is not None,
932
)
vllm/v1/attention/backends/triton_attn.py
+83
-1
@@ -17,7 +17,7 @@ from vllm.attention.ops.triton_reshape_and_cache_flash import (
17
triton_reshape_and_cache_flash,
18
)
19
from vllm.attention.ops.triton_unified_attention import unified_attention
20
-from vllm.config import VllmConfig
20
+from vllm.config import CUDAGraphMode, VllmConfig
21
from vllm.config.cache import CacheDType
22
from vllm.logger import init_logger
23
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -26,6 +26,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
26
)
27
from vllm.platforms import current_platform
28
from vllm.platforms.interface import DeviceCapability
29
+from vllm.utils.math_utils import next_power_of_2
30
from vllm.v1.attention.backends.utils import (
31
AttentionCGSupport,
32
AttentionMetadataBuilder,
@@ -36,6 +37,11 @@ from vllm.v1.kv_cache_interface import AttentionSpec
37
logger = init_logger(__name__)
38
39
40
+# constants
41
+MIN_LAUNCH_GRID_SIZE_2D = 128 # Minimum launch grid size of 2D kernel
42
+NUM_PAR_SOFTMAX_SEGMENTS = 16 # Number of parallel tiled softmax segments
43
+
44
+
45
@dataclass
46
class TritonAttentionMetadata:
47
# NOTE(sang): Definition of context_len, query_len, and seq_len.
@@ -54,6 +60,12 @@ class TritonAttentionMetadata:
60
block_table: torch.Tensor
61
slot_mapping: torch.Tensor
62
63
+ seq_threshold_3D: int
64
+ num_par_softmax_segments: int
65
+ softmax_segm_output: torch.Tensor
66
+ softmax_segm_max: torch.Tensor
67
+ softmax_segm_expsum: torch.Tensor
68
+
69
# For cascade attention.
70
use_cascade: bool
71
common_prefix_len: int
@@ -87,6 +99,60 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
99
self.num_heads_kv = model_config.get_num_kv_heads(vllm_config.parallel_config)
100
self.headdim = model_config.get_head_size()
101
102
+ # Check if CUDA Graphs are enabled for decode
103
+ self.decode_cudagraph_enabled = (
104
+ self.vllm_config.compilation_config.cudagraph_mode
105
+ in (
106
+ CUDAGraphMode.FULL_AND_PIECEWISE,
107
+ CUDAGraphMode.FULL_DECODE_ONLY,
108
+ CUDAGraphMode.FULL,
109
+ )
110
+ )
111
+
112
+ # The launch grid for the 2D kernel is defined as (num_q_blocks, num_heads_kv).
113
+ # A lower bound for num_q_blocks is the number of sequences.
114
+ # To ensure the minimum launch grid size is achieved, the number of sequences
115
+ # must be at least equal to the threshold below.
116
+ # If this threshold is not reached (i.e., the batch size is not large enough),
117
+ # the 3D kernel will be selected instead.
118
+ self.seq_threshold_3D = MIN_LAUNCH_GRID_SIZE_2D // self.num_heads_kv
119
+
120
+ # Modify the threshold if needed.
121
+ if self.decode_cudagraph_enabled:
122
+ capture_sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes
123
+ assert capture_sizes, "CUDA Graphs enabled but no capture sizes specified."
124
+
125
+ # Select the CUDA Graph capture size closest to self.seq_threshold_3D
126
+ # as threshold. This ensures that each captured graph covers the
127
+ # correct execution path.
128
+ self.seq_threshold_3D = min(
129
+ capture_sizes,
130
+ key=lambda x: abs(x - self.seq_threshold_3D),
131
+ )
132
+
133
+ self.num_par_softmax_segments = NUM_PAR_SOFTMAX_SEGMENTS
134
+ headdim_padded = next_power_of_2(self.headdim)
135
+ self.softmax_segm_output = torch.empty(
136
+ (
137
+ self.seq_threshold_3D,
138
+ self.num_heads_q,
139
+ self.num_par_softmax_segments,
140
+ headdim_padded,
141
+ ),
142
+ dtype=torch.float32,
143
+ device=device,
144
+ )
145
+ self.softmax_segm_max = torch.empty(
146
+ (self.seq_threshold_3D, self.num_heads_q, self.num_par_softmax_segments),
147
+ dtype=torch.float32,
148
+ device=device,
149
+ )
150
+ self.softmax_segm_expsum = torch.empty(
151
+ (self.seq_threshold_3D, self.num_heads_q, self.num_par_softmax_segments),
152
+ dtype=torch.float32,
153
+ device=device,
154
+ )
155
+
156
def build_for_cudagraph_capture(
157
self, common_attn_metadata: CommonAttentionMetadata
158
) -> TritonAttentionMetadata:
@@ -143,6 +209,11 @@ class TritonAttentionMetadataBuilder(AttentionMetadataBuilder[TritonAttentionMet
209
prefix_kv_lens=prefix_kv_lens,
210
suffix_kv_lens=suffix_kv_lens,
211
prefix_scheduler_metadata=prefix_scheduler_metadata,
212
+ seq_threshold_3D=self.seq_threshold_3D,
213
+ num_par_softmax_segments=self.num_par_softmax_segments,
214
+ softmax_segm_output=self.softmax_segm_output,
215
+ softmax_segm_max=self.softmax_segm_max,
216
+ softmax_segm_expsum=self.softmax_segm_expsum,
217
)
218
return attn_metadata
219
@@ -349,6 +420,12 @@ class TritonAttentionImpl(AttentionImpl):
420
max_seqlen_k = attn_metadata.max_seq_len
421
block_table = attn_metadata.block_table
422
423
+ seq_threshold_3D = attn_metadata.seq_threshold_3D
424
+ num_par_softmax_segments = attn_metadata.num_par_softmax_segments
425
+ softmax_segm_output = attn_metadata.softmax_segm_output
426
+ softmax_segm_max = attn_metadata.softmax_segm_max
427
+ softmax_segm_expsum = attn_metadata.softmax_segm_expsum
428
+
429
descale_shape = (cu_seqlens_q.shape[0] - 1, key_cache.shape[2])
430
431
unified_attention(
@@ -369,6 +446,11 @@ class TritonAttentionImpl(AttentionImpl):
446
q_descale=None, # Not supported
447
k_descale=layer._k_scale.expand(descale_shape),
448
v_descale=layer._v_scale.expand(descale_shape),
449
+ seq_threshold_3D=seq_threshold_3D,
450
+ num_par_softmax_segments=num_par_softmax_segments,
451
+ softmax_segm_output=softmax_segm_output,
452
+ softmax_segm_max=softmax_segm_max,
453
+ softmax_segm_expsum=softmax_segm_expsum,
454
sinks=self.sinks,
455
output_scale=output_scale,
456
)