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 13, 15:45 UTC
[Kernel] LoRA triton kernels support PDL (#27402)
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Jee Jee Li committed
Nov 7, 2025 at 16:05 UTC
21b82f4ea2f12ab2c3d74f9156b50616b892ea7d
5 files changed
+68
-17
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+22
-7
@@ -6,6 +6,8 @@ import torch
6
from vllm.triton_utils import tl, triton
7
from vllm.utils.torch_utils import direct_register_custom_op
8
9
+from .utils import supports_pdl
10
+
11
_LORA_PTR_DICT: dict[tuple[int, ...], torch.tensor] = {}
12
13
@@ -82,6 +84,8 @@ def _fused_moe_lora_kernel(
84
BLOCK_SIZE_K: tl.constexpr,
85
GROUP_SIZE_M: tl.constexpr,
86
SPLIT_K: tl.constexpr,
87
+ USE_GDC: tl.constexpr,
88
+ IS_PRIMARY: tl.constexpr,
89
):
90
pid = tl.program_id(axis=0)
91
slice_id = tl.program_id(axis=1)
@@ -110,13 +114,11 @@ def _fused_moe_lora_kernel(
114
num_tokens_post_padded = tl.load(num_tokens_post_padded_ptr + lora_id)
115
if pid_m * BLOCK_SIZE_M >= num_tokens_post_padded:
116
return
113
-
117
# get the expert_id to process curr shard
118
ind = lora_id * stride_el + pid_m
119
expert_id = tl.load(expert_ids_ptr + ind, ind < max_loras * stride_el, -1)
120
if expert_id == -1:
121
return
119
-
122
# get a_ptr,b_ptr,c_ptr
123
cur_a_ptr = a_ptr + (slice_id % num_slice_a) * slice_a_size
124
cur_b_ptr = tl.load(b_ptr + slice_id).to(tl.pointer_type(c_ptr.dtype.element_ty))
@@ -149,12 +151,17 @@ def _fused_moe_lora_kernel(
151
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
152
for k in range(0, grid_k):
153
k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K)
154
+ # pre-fetch lora weight
155
+ b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
156
+ # GDC wait waits for ALL programs in the the prior kernel to complete
157
+ # before continuing.
158
+ if USE_GDC and not IS_PRIMARY:
159
+ tl.extra.cuda.gdc_wait()
160
a = tl.load(
161
a_ptrs,
162
mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),
163
other=0.0,
164
)
157
- b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
165
accumulator += tl.dot(a, b)
166
# Advance the ptrs to the next K block.
167
a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
@@ -163,12 +170,15 @@ def _fused_moe_lora_kernel(
170
if MUL_ROUTED_WEIGHT:
171
moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0)
172
accumulator = accumulator * moe_weight[:, None]
166
-
173
+ if USE_GDC and IS_PRIMARY:
174
+ # GDC launch dependents hints the runtime system to launch dependent kernels.
175
+ tl.extra.cuda.gdc_launch_dependents()
176
accumulator = accumulator.to(c_ptr.dtype.element_ty)
177
# Write back the block of the output
178
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
179
c_ptrs = cur_c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
180
c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
181
+
182
if SPLIT_K == 1:
183
tl.store(c_ptrs, accumulator, mask=c_mask)
184
else:
@@ -209,7 +219,7 @@ def _fused_moe_lora_shrink(
219
mul_routed_weight: bool = False,
220
) -> None:
221
w1_lora_a_stacked = lora_a_stacked[0]
212
-
222
+ use_gdc = supports_pdl(qcurr_hidden_states.device)
223
shrink_config = {
224
"BLOCK_SIZE_M": block_size_m,
225
"BLOCK_SIZE_N": block_size_n,
@@ -218,6 +228,8 @@ def _fused_moe_lora_shrink(
228
"num_warps": num_warps,
229
"num_stages": num_stages,
230
"SPLIT_K": split_k,
231
+ "USE_GDC": use_gdc,
232
+ "launch_pdl": use_gdc, # triton kernel metadata
233
}
234
235
b_ptr = _get_ptr(lora_a_stacked, device)
@@ -229,7 +241,6 @@ def _fused_moe_lora_shrink(
241
len(lora_a_stacked),
242
lora_a_stacked[0].shape[0],
243
)
232
-
244
_fused_moe_lora_kernel[grid](
245
qcurr_hidden_states,
246
b_ptr,
@@ -261,6 +272,7 @@ def _fused_moe_lora_shrink(
272
num_slice_c=num_slices,
273
top_k=1 if mul_routed_weight else top_k_num,
274
MUL_ROUTED_WEIGHT=False,
275
+ IS_PRIMARY=True,
276
**shrink_config,
277
)
278
@@ -314,7 +326,7 @@ def _fused_moe_lora_expand(
326
dtype=output.dtype,
327
device=device,
328
)
317
-
329
+ use_gdc = supports_pdl(a_intermediate_cache1.device)
330
expand_config = {
331
"BLOCK_SIZE_M": block_size_m,
332
"BLOCK_SIZE_N": block_size_n,
@@ -323,6 +335,8 @@ def _fused_moe_lora_expand(
335
"num_warps": num_warps,
336
"num_stages": num_stages,
337
"SPLIT_K": split_k, # Set split_k = 1 for expand calls
338
+ "USE_GDC": use_gdc,
339
+ "launch_pdl": use_gdc, # triton kernel metadata
340
}
341
342
grid = lambda META: (
@@ -361,6 +375,7 @@ def _fused_moe_lora_expand(
375
num_slice_c=num_slices,
376
top_k=1,
377
MUL_ROUTED_WEIGHT=mul_routed_weight,
378
+ IS_PRIMARY=False,
379
**expand_config,
380
)
381
for i in range(num_slices):
vllm/lora/ops/triton_ops/kernel_utils.py
+21
-7
@@ -22,6 +22,7 @@ def mm_k(
22
SPLIT_K: tl.constexpr,
23
CAST_TYPE: tl.constexpr,
24
b_dtype: tl.constexpr,
25
+ USE_GDC: tl.constexpr,
26
):
27
"""
28
Given a_ptr and b_ptr, that identify the rows of A (m x k) and columns of
@@ -45,19 +46,25 @@ def mm_k(
46
CAST_TYPE: if True, cast the values from the A matrix to the B
47
matrix dtype.
48
b_dtype: datatype of the B matrix
49
+ USE_GDC: Whether to use PDL. True indicates use.
50
"""
51
accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
52
for k in range(tl.cdiv(K, BLOCK_K * SPLIT_K)):
53
if EVEN_K:
52
- tiled_a = tl.load(a_ptr)
54
+ # pre-fetech lora weight
55
tiled_b = tl.load(b_ptr)
56
+ if USE_GDC:
57
+ tl.extra.cuda.gdc_wait()
58
+ tiled_a = tl.load(a_ptr)
59
else:
55
- tiled_a = tl.load(
56
- a_ptr, mask=offset_k[None, :] < K - k * (BLOCK_K * SPLIT_K), other=0
57
- )
60
tiled_b = tl.load(
61
b_ptr, mask=offset_k[:, None] < K - k * (BLOCK_K * SPLIT_K), other=0
62
)
63
+ if USE_GDC:
64
+ tl.extra.cuda.gdc_wait()
65
+ tiled_a = tl.load(
66
+ a_ptr, mask=offset_k[None, :] < K - k * (BLOCK_K * SPLIT_K), other=0
67
+ )
68
if CAST_TYPE:
69
tiled_a = tiled_a.to(b_dtype)
70
accumulator += tl.dot(
@@ -102,6 +109,7 @@ def do_expand_kernel(
109
EVEN_K: tl.constexpr,
110
CAST_TYPE: tl.constexpr,
111
ADD_INPUTS: tl.constexpr,
112
+ USE_GDC: tl.constexpr,
113
):
114
"""
115
Given an array of integers that identifies the rows of A, ram,
@@ -154,6 +162,7 @@ def do_expand_kernel(
162
163
# Compute the block matrix product.
164
SPLIT_K = 1
165
+
166
accumulator = mm_k(
167
a_ptr,
168
b_ptr,
@@ -168,6 +177,7 @@ def do_expand_kernel(
177
SPLIT_K,
178
CAST_TYPE,
179
cur_lora_ptr.dtype.element_ty,
180
+ USE_GDC,
181
)
182
183
tiled_c = accumulator.to(cur_lora_ptr.dtype.element_ty)
@@ -223,6 +233,7 @@ def do_shrink_kernel(
233
EVEN_K: tl.constexpr,
234
SPLIT_K: tl.constexpr,
235
SLICE_NUM: tl.constexpr,
236
+ USE_GDC: tl.constexpr,
237
):
238
"""
239
Given an array of integers that identifies the rows of A, ram,
@@ -272,8 +283,11 @@ def do_shrink_kernel(
283
SPLIT_K,
284
False,
285
cur_lora_ptr.dtype.element_ty,
286
+ False, # USE_GDC is always False in shrink kernel
287
)
276
-
288
+ # GDC launch dependents hints the runtime system to launch dependent kernels.
289
+ if USE_GDC:
290
+ tl.extra.cuda.gdc_launch_dependents()
291
# Identify the C output pointers to store the results of the accumulator.
292
offset_cn = tl.arange(0, BLOCK_N) + pid_n * BLOCK_N
293
offset_cm = tl.arange(0, BLOCK_M)
@@ -284,10 +298,10 @@ def do_shrink_kernel(
298
+ offset_cn[None, :] * output_d2_stride
299
)
300
c_mask = (offset_cm[:, None] < M_LEN) & (offset_cn[None, :] < N)
287
-
301
accumulator *= scaling
302
+
303
# handles write-back with reduction-splitting
304
if SPLIT_K == 1:
305
tl.store(c_ptr, accumulator, mask=c_mask)
306
else:
293
- tl.atomic_add(c_ptr, accumulator, mask=c_mask)
307
+ tl.atomic_add(c_ptr, accumulator, mask=c_mask, sem="relaxed")
vllm/lora/ops/triton_ops/lora_expand_op.py
+7
-1
@@ -14,6 +14,8 @@ from vllm.lora.ops.triton_ops.utils import _get_lora_b_ptr, get_lora_op_configs
14
from vllm.triton_utils import tl, triton
15
from vllm.utils.torch_utils import direct_register_custom_op
16
17
+from .utils import supports_pdl
18
+
19
20
@triton.jit
21
def _lora_expand_kernel(
@@ -45,6 +47,7 @@ def _lora_expand_kernel(
47
CAST_TYPE: tl.constexpr,
48
SLICE_NUM: tl.constexpr,
49
SAME_STRIDE: tl.constexpr,
50
+ USE_GDC: tl.constexpr,
51
):
52
cta_n_num = tl.cdiv(N, BLOCK_N)
53
cta_m_num = tl.cdiv(M, BLOCK_M)
@@ -121,6 +124,7 @@ def _lora_expand_kernel(
124
EVEN_K,
125
CAST_TYPE,
126
ADD_INPUTS,
127
+ USE_GDC,
128
)
129
130
@@ -236,7 +240,7 @@ def _lora_expand(
240
# thread blocks simply exit.
241
MAX_LORAS,
242
)
239
-
243
+ use_gdc = supports_pdl(inputs.device)
244
_lora_expand_kernel[grid](
245
inputs,
246
lora_ptr_tensor,
@@ -266,9 +270,11 @@ def _lora_expand(
270
CAST_TYPE,
271
NUM_SLICES,
272
same_stride,
273
+ use_gdc,
274
num_warps=NUM_WARPS,
275
num_ctas=NUM_CTAS,
276
num_stages=NUM_STAGES,
277
+ launch_pdl=use_gdc,
278
)
279
280
return
vllm/lora/ops/triton_ops/lora_shrink_op.py
+7
-2
@@ -14,6 +14,8 @@ from vllm.lora.ops.triton_ops.utils import _get_lora_a_ptr, get_lora_op_configs
14
from vllm.triton_utils import tl, triton
15
from vllm.utils.torch_utils import direct_register_custom_op
16
17
+from .utils import supports_pdl
18
+
19
20
@triton.jit
21
def _lora_shrink_kernel(
@@ -43,6 +45,7 @@ def _lora_shrink_kernel(
45
SPLIT_K: tl.constexpr,
46
GROUP_SIZE_M: tl.constexpr,
47
SLICE_NUM: tl.constexpr,
48
+ USE_GDC: tl.constexpr,
49
):
50
cta_n_num = tl.cdiv(N, BLOCK_N)
51
cta_m_num = tl.cdiv(M, BLOCK_M)
@@ -83,7 +86,6 @@ def _lora_shrink_kernel(
86
cta_lora_seq_indices = (
87
token_indices_sorted_by_lora_ids + lora_m_indices_start + cta_m_offset
88
)
86
-
89
# Load all relevant row indices.
90
offset_m = tl.arange(0, BLOCK_M) % cta_m_len
91
ram = tl.load(cta_lora_seq_indices + offset_m)
@@ -118,6 +120,7 @@ def _lora_shrink_kernel(
120
EVEN_K,
121
SPLIT_K,
122
SLICE_NUM,
123
+ USE_GDC,
124
)
125
126
@@ -217,7 +220,7 @@ def _lora_shrink(
220
# thread blocks exit early.
221
MAX_LORAS,
222
)
220
-
223
+ use_gdc = supports_pdl(inputs.device)
224
_lora_shrink_kernel[grid](
225
inputs,
226
lora_ptr_tensor,
@@ -245,9 +248,11 @@ def _lora_shrink(
248
SPLIT_K,
249
GROUP_SIZE_M,
250
NUM_SLICES,
251
+ use_gdc,
252
num_warps=NUM_WARPS,
253
num_ctas=NUM_CTAS,
254
num_stages=NUM_STAGES,
255
+ launch_pdl=use_gdc,
256
)
257
258
return
vllm/lora/ops/triton_ops/utils.py
+11
@@ -3,6 +3,7 @@
3
4
import functools
5
import json
6
+from functools import lru_cache
7
from pathlib import Path
8
from typing import Any
9
@@ -10,6 +11,7 @@ import torch
11
12
from vllm import envs
13
from vllm.logger import init_logger
14
+from vllm.platforms import current_platform
15
16
logger = init_logger(__name__)
17
@@ -282,3 +284,12 @@ def get_lora_op_configs(
284
285
assert config_data is not None
286
return config_data
287
+
288
+
289
+@lru_cache
290
+def supports_pdl(device: torch.device | None = None) -> bool:
291
+ """
292
+ Refer to: https://github.com/triton-lang/triton/blob/v3.5.0/python/tutorials/11-programmatic-dependent-launch.py
293
+ """
294
+ # PDL requires compute capability SM90 or above
295
+ return current_platform.is_cuda() and current_platform.has_device_capability(90)