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 7, 12:45 UTC
Add TMA support to fused_moe_lora kernel (#32195)
Signed-off-by: gnovack <gnovack@amazon.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
gnovack committed
Feb 28, 2026 at 18:55 UTC
3ecd0bf9fccc425c015f7723b6a7730c0dda2970
5 files changed
+279
-52
tests/lora/test_fused_moe_lora_kernel.py
+25
-10
@@ -231,17 +231,22 @@ def use_torch(
231
lora_a_stacked,
232
lora_b_stacked,
233
top_k_num,
234
+ num_slices=1,
235
):
236
outputs = []
237
for i in range(hidden_states.shape[0]):
237
- lora_idx = token_lora_mapping[i]
238
- expert_ids = topk_ids[i]
239
- lora_a = lora_a_stacked[0][lora_idx][expert_ids]
240
- lora_b = lora_b_stacked[0][lora_idx][expert_ids]
241
- tensors = [
242
- hidden_states[i] @ lora_a[x].T @ lora_b[x].T for x in range(top_k_num)
243
- ]
244
- outputs.append(torch.stack(tensors, dim=0))
238
+ slice_tensors = []
239
+ for slice_id in range(num_slices):
240
+ lora_idx = token_lora_mapping[i]
241
+ expert_ids = topk_ids[i]
242
+ lora_a = lora_a_stacked[slice_id][lora_idx][expert_ids]
243
+ lora_b = lora_b_stacked[slice_id][lora_idx][expert_ids]
244
+ tensors = [
245
+ hidden_states[i] @ lora_a[x].T @ lora_b[x].T for x in range(top_k_num)
246
+ ]
247
+ slice_tensors.append(torch.stack(tensors, dim=0))
248
+
249
+ outputs.append(torch.concat(slice_tensors, dim=-1))
250
return torch.stack(outputs, dim=0)
251
252
@@ -259,6 +264,7 @@ SEED = [42]
264
@pytest.mark.parametrize("K", [2048])
265
@pytest.mark.parametrize("max_lora_rank", [16, 32, 64])
266
@pytest.mark.parametrize("block_size", [16])
267
+@pytest.mark.parametrize("num_slices", [1, 2])
268
@pytest.mark.parametrize("dtype", DTYPES)
269
@pytest.mark.parametrize("device", DEVICES)
270
@pytest.mark.parametrize("seed", SEED)
@@ -271,6 +277,7 @@ def test_fused_moe_lora_kernel(
277
K,
278
max_lora_rank,
279
block_size,
280
+ num_slices,
281
dtype,
282
device,
283
seed,
@@ -295,17 +302,19 @@ def test_fused_moe_lora_kernel(
302
),
303
dtype=dtype,
304
)
305
+ for _ in range(num_slices)
306
]
307
lora_b_stacked = [
308
torch.rand(
309
(
310
max_loras,
311
num_experts,
304
- N,
312
+ N // num_slices,
313
max_lora_rank,
314
),
315
dtype=dtype,
316
)
317
+ for _ in range(num_slices)
318
]
319
hidden_states = torch.rand(
320
(
@@ -340,6 +349,7 @@ def test_fused_moe_lora_kernel(
349
lora_a_stacked,
350
lora_b_stacked,
351
top_k_num,
352
+ num_slices,
353
)
354
355
torch.testing.assert_close(output, output2, atol=1e-2, rtol=1e-2)
@@ -434,6 +444,7 @@ def use_fused_moe_lora_kernel_naive(
444
@pytest.mark.parametrize("K", [2048])
445
@pytest.mark.parametrize("max_lora_rank", [16, 32])
446
@pytest.mark.parametrize("block_size", [16])
447
+@pytest.mark.parametrize("num_slices", [1, 2])
448
@pytest.mark.parametrize("dtype", DTYPES)
449
@pytest.mark.parametrize("device", DEVICES)
450
@pytest.mark.parametrize("seed", SEED)
@@ -446,6 +457,7 @@ def test_fused_moe_lora_kernel_naive_block_assignment(
457
K,
458
max_lora_rank,
459
block_size,
460
+ num_slices,
461
dtype,
462
device,
463
seed,
@@ -484,17 +496,19 @@ def test_fused_moe_lora_kernel_naive_block_assignment(
496
),
497
dtype=dtype,
498
)
499
+ for _ in range(num_slices)
500
]
501
lora_b_stacked = [
502
torch.rand(
503
(
504
max_loras,
505
num_experts,
493
- N,
506
+ N // num_slices,
507
max_lora_rank,
508
),
509
dtype=dtype,
510
)
511
+ for _ in range(num_slices)
512
]
513
hidden_states = torch.rand(
514
(
@@ -529,6 +543,7 @@ def test_fused_moe_lora_kernel_naive_block_assignment(
543
lora_a_stacked,
544
lora_b_stacked,
545
top_k_num,
546
+ num_slices,
547
)
548
549
torch.testing.assert_close(output, output_ref, atol=1e-2, rtol=1e-2)
tests/lora/test_olmoe_tp.py
+39
@@ -2,7 +2,11 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
+import shutil
6
+
7
import pytest
8
+import torch
9
+from safetensors.torch import load_file, save_file
10
11
import vllm
12
from vllm.lora.request import LoRARequest
@@ -122,6 +126,41 @@ def test_olmoe_lora_mixed(olmoe_lora_files):
126
generate_and_test(llm, olmoe_lora_files, lora_id=[1, None, 3, None])
127
128
129
+def test_olmoe_lora_mixed_random(olmoe_lora_files, tmp_path):
130
+ # Create a dummy LoRA with random weights based on the real one
131
+ random_lora_path = tmp_path / "random_lora"
132
+ shutil.copytree(olmoe_lora_files, random_lora_path)
133
+
134
+ weights_path = random_lora_path / "adapter_model.safetensors"
135
+ weights = load_file(str(weights_path))
136
+ random_weights = {k: torch.randn_like(v) for k, v in weights.items()}
137
+ save_file(random_weights, str(weights_path))
138
+
139
+ llm = vllm.LLM(
140
+ MODEL_PATH,
141
+ max_model_len=1024,
142
+ enable_lora=True,
143
+ max_loras=4,
144
+ enforce_eager=True,
145
+ trust_remote_code=True,
146
+ enable_chunked_prefill=True,
147
+ )
148
+
149
+ prompts = [
150
+ PROMPT_TEMPLATE.format(context="How many candidates are there?"),
151
+ PROMPT_TEMPLATE.format(context="Count the number of candidates."),
152
+ ]
153
+
154
+ lora_requests = [
155
+ LoRARequest("real", 1, olmoe_lora_files),
156
+ LoRARequest("random", 2, str(random_lora_path)),
157
+ ]
158
+
159
+ sampling_params = vllm.SamplingParams(temperature=0, max_tokens=64)
160
+ outputs = llm.generate(prompts, sampling_params, lora_request=lora_requests)
161
+ assert outputs[0].outputs[0].text.strip().startswith(EXPECTED_LORA_OUTPUT[0])
162
+
163
+
164
@pytest.mark.parametrize("fully_sharded_loras", [False, True])
165
@multi_gpu_test(num_gpus=2)
166
def test_olmoe_lora_tp2(olmoe_lora_files, fully_sharded_loras):
vllm/lora/ops/triton_ops/fused_moe_lora_op.py
+196
-42
@@ -8,9 +8,10 @@ from vllm.distributed import (
8
tensor_model_parallel_all_reduce,
9
)
10
from vllm.triton_utils import tl, triton
11
+from vllm.triton_utils.allocation import set_triton_allocator
12
from vllm.utils.torch_utils import direct_register_custom_op
13
13
-from .utils import supports_pdl
14
+from .utils import supports_pdl, supports_tma
15
16
17
@triton.jit
@@ -70,6 +71,37 @@ def _get_token_offs(
71
)
72
73
74
+@triton.jit
75
+def _get_c_ptrs(
76
+ cur_c_ptr,
77
+ lora_id,
78
+ pid_m,
79
+ offs,
80
+ offs_token,
81
+ offs_cn,
82
+ stride_cm,
83
+ stride_cn,
84
+ EM: tl.constexpr,
85
+ BLOCK_SIZE_M: tl.constexpr,
86
+ sort_c: tl.constexpr,
87
+):
88
+ # When sort_c is true, store the output in c_ptr using token order defined
89
+ # in sorted_token_ids_ptr; otherwise, use the original token order from the prompt
90
+ if sort_c:
91
+ offs_token_id = pid_m * BLOCK_SIZE_M + offs
92
+ c_ptrs = (
93
+ cur_c_ptr
94
+ + lora_id * EM * stride_cm
95
+ + stride_cm * offs_token_id[:, None]
96
+ + stride_cn * offs_cn[None, :]
97
+ )
98
+ else:
99
+ c_ptrs = (
100
+ cur_c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
101
+ )
102
+ return c_ptrs
103
+
104
+
105
_LORA_PTR_DICT: dict[tuple[int, ...], torch.tensor] = {}
106
107
@@ -125,7 +157,9 @@ def _adjust_kernel_inputs(
157
)
158
def _fused_moe_lora_kernel(
159
a_ptr,
160
+ a_desc,
161
b_ptr,
162
+ b_desc,
163
c_ptr,
164
topk_weights_ptr,
165
sorted_token_ids_ptr,
@@ -177,6 +211,18 @@ def _fused_moe_lora_kernel(
211
USE_GDC: tl.constexpr,
212
launch_pdl: tl.constexpr,
213
IS_PRIMARY: tl.constexpr,
214
+ USE_TMA: tl.constexpr,
215
+ # sort_c determines whether tokens are stored in C in the order determined
216
+ # by sorted_token_ids to enable later TMA loads from this tensor.
217
+ #
218
+ # When USE_TMA is enabled, the parameter combinations are:
219
+ # a_desc | b_desc | sort_c | Use Case
220
+ # --------|---------|--------|-----------------------------
221
+ # yes | yes | False | expand kernel (num_slices=1)
222
+ # no | yes | True | shrink kernel (num_slices=1)
223
+ # yes | no | False | expand kernel (num_slices>1)
224
+ # no | no | True | shrink kernel (num_slices>1)
225
+ sort_c: tl.constexpr,
226
):
227
pid = tl.program_id(axis=0)
228
slice_id = tl.program_id(axis=1)
@@ -250,58 +296,90 @@ def _fused_moe_lora_kernel(
296
cur_b_ptr = tl.load(b_ptr + slice_id).to(tl.pointer_type(c_ptr.dtype.element_ty))
297
cur_c_ptr = c_ptr + (slice_id % num_slice_c) * slice_c_size
298
253
- # remove modulo wrap-around
254
- offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int32)
299
offs_k = pid_sk * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
300
token_mask = offs_token < num_valid_tokens
301
258
- # get a_ptrs,b_ptrs
259
- a_ptrs = cur_a_ptr + (
260
- offs_token[:, None] // token_mapping_factor * stride_am
261
- + offs_k[None, :] * stride_ak
262
- )
302
+ if USE_TMA and a_desc is not None:
303
+ # Expand path - with TMA enabled, load from A using TMA descriptor
304
+ offs_am = (
305
+ slice_id * max_loras * EM
306
+ + lora_id * EM
307
+ + pid_m * BLOCK_SIZE_M // token_mapping_factor
308
+ )
309
+ offs_ak = pid_sk * BLOCK_SIZE_K
310
+ else:
311
+ # Shrink path - load hidden states based on order defined in
312
+ # 'sorted_token_ids_ptr' then store them in c_ptr in this same sorted order
313
+ tl.static_assert(a_desc is None, "a_desc must be none")
314
+ a_ptrs = cur_a_ptr + (
315
+ offs_token[:, None] // token_mapping_factor * stride_am
316
+ + offs_k[None, :] * stride_ak
317
+ )
318
264
- b_ptrs = (
265
- cur_b_ptr
266
- + lora_id * stride_bl
267
- + expert_id * stride_be
268
- + offs_k[:, None] * stride_bk
269
- + offs_bn[None, :] * stride_bn
270
- )
319
+ if USE_TMA:
320
+ offs_bn = pid_n * BLOCK_SIZE_N
321
+ offs_bk = pid_sk * BLOCK_SIZE_K
322
+ if b_desc is None:
323
+ # Note(@gnovack) - Allocation of TMA descriptors on-device
324
+ # can cause conflicts when running in parallel via PDL
325
+ if USE_GDC and not IS_PRIMARY:
326
+ tl.extra.cuda.gdc_wait()
327
+
328
+ b_desc = tl.make_tensor_descriptor(
329
+ cur_b_ptr,
330
+ shape=[max_loras, num_experts, N, K],
331
+ strides=[stride_bl, stride_be, stride_bn, stride_bk],
332
+ block_shape=[1, 1, BLOCK_SIZE_N, BLOCK_SIZE_K],
333
+ )
334
+ else:
335
+ offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N).to(tl.int32)
336
+ b_ptrs = (
337
+ cur_b_ptr
338
+ + lora_id * stride_bl
339
+ + expert_id * stride_be
340
+ + offs_k[:, None] * stride_bk
341
+ + offs_bn[None, :] * stride_bn
342
+ )
343
344
if USE_GDC and IS_PRIMARY:
345
# GDC launch dependents hints the runtime system to launch dependent kernels.
346
tl.extra.cuda.gdc_launch_dependents()
347
276
- # accumulator
348
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
349
350
if USE_GDC and not IS_PRIMARY:
351
tl.extra.cuda.gdc_wait()
352
353
for k in range(0, grid_k):
283
- k_remaining = K - k * (BLOCK_SIZE_K * SPLIT_K)
284
- # GDC wait waits for ALL programs in the prior kernel to complete
285
- # before continuing.
354
+ cur_k_offset = k * (BLOCK_SIZE_K * SPLIT_K)
355
+ k_remaining = K - cur_k_offset
356
# pre-fetch lora weight
287
- # add (offs_bn < N) mask; optional .ca for B
288
- b_mask = (offs_k[:, None] < k_remaining) & (offs_bn[None, :] < N)
289
- if USE_B_L2_CACHE:
290
- b = tl.load(b_ptrs, mask=b_mask, other=0.0, cache_modifier=".ca")
357
+ if b_desc is not None:
358
+ b = (
359
+ b_desc.load([lora_id, expert_id, offs_bn, offs_bk + cur_k_offset])
360
+ .reshape(BLOCK_SIZE_N, BLOCK_SIZE_K)
361
+ .T
362
+ )
363
else:
292
- b = tl.load(b_ptrs, mask=b_mask, other=0.0)
293
-
294
- if USE_GDC and not IS_PRIMARY:
295
- tl.extra.cuda.gdc_wait()
296
- a = tl.load(
297
- a_ptrs,
298
- mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),
299
- other=0.0,
300
- )
364
+ # add (offs_bn < N) mask; optional .ca for B
365
+ b_mask = (offs_k[:, None] < k_remaining) & (offs_bn[None, :] < N)
366
+ if USE_B_L2_CACHE:
367
+ b = tl.load(b_ptrs, mask=b_mask, other=0.0, cache_modifier=".ca")
368
+ else:
369
+ b = tl.load(b_ptrs, mask=b_mask, other=0.0)
370
+ b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk
371
+
372
+ if a_desc is not None:
373
+ a = a_desc.load([offs_am, offs_ak + cur_k_offset])
374
+ else:
375
+ a = tl.load(
376
+ a_ptrs,
377
+ mask=token_mask[:, None] & (offs_k[None, :] < k_remaining),
378
+ other=0.0,
379
+ )
380
+ a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
381
+
382
accumulator += tl.dot(a, b)
302
- # Advance the ptrs to the next K block.
303
- a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
304
- b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk
383
384
if MUL_ROUTED_WEIGHT:
385
moe_weight = tl.load(topk_weights_ptr + offs_token, mask=token_mask, other=0.0)
@@ -309,7 +387,19 @@ def _fused_moe_lora_kernel(
387
accumulator = accumulator.to(c_ptr.dtype.element_ty)
388
# Write back the block of the output
389
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
312
- c_ptrs = cur_c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
390
+ c_ptrs = _get_c_ptrs(
391
+ cur_c_ptr,
392
+ lora_id,
393
+ pid_m,
394
+ offs,
395
+ offs_token,
396
+ offs_cn,
397
+ stride_cm,
398
+ stride_cn,
399
+ EM,
400
+ BLOCK_SIZE_M,
401
+ sort_c,
402
+ )
403
c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
404
405
if SPLIT_K == 1:
@@ -357,6 +447,7 @@ def _fused_moe_lora_shrink(
447
num_active_loras: int,
448
mul_routed_weight: bool = False,
449
use_gdc: bool = False,
450
+ use_tma: bool = False,
451
) -> None:
452
w1_lora_a_stacked = lora_a_stacked[0]
453
shrink_config = {
@@ -369,6 +460,7 @@ def _fused_moe_lora_shrink(
460
"SPLIT_K": split_k,
461
"USE_GDC": use_gdc,
462
"launch_pdl": use_gdc, # triton kernel metadata
463
+ "USE_TMA": use_tma,
464
}
465
466
b_ptr = _get_ptr(lora_a_stacked, device)
@@ -383,9 +475,20 @@ def _fused_moe_lora_shrink(
475
len(lora_a_stacked),
476
grid_lora_dim,
477
)
478
+
479
+ a_desc = None
480
+ b_desc = None
481
+ if use_tma and num_slices == 1:
482
+ b_desc = triton.tools.tensor_descriptor.TensorDescriptor.from_tensor(
483
+ lora_a_stacked[0],
484
+ [1, 1, shrink_config["BLOCK_SIZE_N"], shrink_config["BLOCK_SIZE_K"]],
485
+ )
486
+
487
_fused_moe_lora_kernel[grid](
488
qcurr_hidden_states,
489
+ a_desc,
490
b_ptr,
491
+ b_desc,
492
a_intermediate_cache1,
493
topk_weights,
494
sorted_token_ids,
@@ -407,8 +510,8 @@ def _fused_moe_lora_shrink(
510
w1_lora_a_stacked.stride(1),
511
w1_lora_a_stacked.stride(3),
512
w1_lora_a_stacked.stride(2),
410
- a_intermediate_cache1.stride(2),
411
- a_intermediate_cache1.stride(3),
513
+ a_intermediate_cache1.stride(-2),
514
+ a_intermediate_cache1.stride(-1),
515
stride_tl,
516
stride_el,
517
slice_a_size=qcurr_hidden_states.numel(),
@@ -419,7 +522,8 @@ def _fused_moe_lora_shrink(
522
naive_block_assignment=sorted_token_ids is None,
523
MUL_ROUTED_WEIGHT=False,
524
ADD_INPUTS=False,
422
- USE_B_L2_CACHE=True, # new
525
+ USE_B_L2_CACHE=True,
526
+ sort_c=use_tma and sorted_token_ids is not None,
527
IS_PRIMARY=True,
528
**shrink_config,
529
)
@@ -462,6 +566,7 @@ def _fused_moe_lora_expand(
566
mul_routed_weight: bool = False,
567
offset: int = 0,
568
use_gdc: bool = False,
569
+ use_tma: bool = False,
570
) -> None:
571
b_ptr = _get_ptr(lora_b_stacked, device)
572
K = max_lora_rank
@@ -470,7 +575,7 @@ def _fused_moe_lora_expand(
575
w1_lora_b_stacked = lora_b_stacked[0]
576
577
a_intermediate_cache1 = a_intermediate_cache1.view(
473
- -1, a_intermediate_cache1.shape[3]
578
+ -1, a_intermediate_cache1.shape[-1]
579
)
580
581
expand_config = {
@@ -483,6 +588,7 @@ def _fused_moe_lora_expand(
588
"SPLIT_K": 1, # Set split_k = 1 for expand calls
589
"USE_GDC": use_gdc,
590
"launch_pdl": use_gdc, # triton kernel metadata
591
+ "USE_TMA": use_tma,
592
}
593
594
grid_lora_dim, stride_tl, stride_el = _adjust_kernel_inputs(
@@ -498,10 +604,27 @@ def _fused_moe_lora_expand(
604
# Fast path: directly accumulate into the corresponding slice interval of output.
605
out_view = output[:, :, offset : offset + num_slices * N]
606
slice_c_size = N * out_view.stride(2)
607
+ a_desc = None
608
+ b_desc = None
609
+ if use_tma:
610
+ if sorted_token_ids is not None:
611
+ a_desc = triton.tools.tensor_descriptor.TensorDescriptor.from_tensor(
612
+ a_intermediate_cache1,
613
+ [expand_config["BLOCK_SIZE_M"], expand_config["BLOCK_SIZE_K"]],
614
+ )
615
+ if num_slices == 1:
616
+ b_desc = triton.tools.tensor_descriptor.TensorDescriptor.from_tensor(
617
+ lora_b_stacked[0],
618
+ [1, 1, expand_config["BLOCK_SIZE_N"], expand_config["BLOCK_SIZE_K"]],
619
+ )
620
+ else:
621
+ b_desc = None
622
623
_fused_moe_lora_kernel[grid](
624
a_intermediate_cache1,
625
+ a_desc,
626
b_ptr,
627
+ b_desc,
628
out_view,
629
topk_weights,
630
sorted_token_ids,
@@ -535,7 +658,8 @@ def _fused_moe_lora_expand(
658
naive_block_assignment=sorted_token_ids is None,
659
MUL_ROUTED_WEIGHT=mul_routed_weight,
660
ADD_INPUTS=True,
538
- USE_B_L2_CACHE=True, # new
661
+ USE_B_L2_CACHE=True,
662
+ sort_c=False,
663
IS_PRIMARY=False,
664
**expand_config,
665
)
@@ -616,8 +740,34 @@ def _fused_moe_lora(
740
else num_tokens * shrink_block_size_m
741
)
742
743
+ # TMA is not currently compatiple with fully_sharded due to the non-determinism
744
+ # of token id sorting across ranks.
745
+ use_tma = supports_tma(device) and not fully_sharded
746
+
747
+ intermediate_cache_shape = (
748
+ num_slices,
749
+ M,
750
+ top_k_num,
751
+ max_lora_rank,
752
+ )
753
+ if use_tma:
754
+ if num_slices > 1:
755
+ # if num_slices > 1, we construct TMA descriptors for LoRA
756
+ # weights within the kernel, which requires us to first set an allocator
757
+ set_triton_allocator(device)
758
+
759
+ # When storing intermediate data in sorted order for TMA, we
760
+ # need an extra 'num_active_loras' dim in the cache to avoid conflicts
761
+ if sorted_token_ids is not None:
762
+ intermediate_cache_shape = (
763
+ num_slices,
764
+ sorted_token_ids.shape[0],
765
+ EM,
766
+ max_lora_rank,
767
+ )
768
+
769
a_intermediate_cache1 = torch.zeros(
620
- (num_slices, M, top_k_num, max_lora_rank),
770
+ intermediate_cache_shape,
771
dtype=output.dtype,
772
device=device,
773
)
@@ -654,6 +804,7 @@ def _fused_moe_lora(
804
num_active_loras,
805
mul_routed_weight,
806
use_gdc=use_gdc,
807
+ use_tma=use_tma,
808
)
809
810
if fully_sharded:
@@ -703,6 +854,7 @@ def _fused_moe_lora(
854
mul_routed_weight,
855
offset,
856
use_gdc=use_gdc,
857
+ use_tma=use_tma,
858
)
859
860
@@ -772,6 +924,7 @@ def _fused_moe_lora_shrink_fake(
924
num_active_loras: int,
925
mul_routed_weight: bool = False,
926
use_gdc: bool = False,
927
+ use_tma: bool = False,
928
) -> None:
929
return
930
@@ -809,6 +962,7 @@ def _fused_moe_lora_expand_fake(
962
mul_routed_weight: bool = False,
963
offset: int = 0,
964
use_gdc: bool = False,
965
+ use_tma: bool = False,
966
) -> None:
967
return
968
vllm/lora/ops/triton_ops/utils.py
+6
@@ -316,3 +316,9 @@ def supports_pdl(device: torch.device | None = None) -> bool:
316
and current_platform.has_device_capability(90)
317
and not envs.VLLM_LORA_DISABLE_PDL
318
)
319
+
320
+
321
+@lru_cache
322
+def supports_tma(device: torch.device | None = None) -> bool:
323
+ # TMA requires compute capability SM90 or above
324
+ return current_platform.is_cuda() and current_platform.has_device_capability(90)
vllm/triton_utils/allocation.py
new
+13
@@ -0,0 +1,13 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from vllm.triton_utils import triton
7
+
8
+
9
+def set_triton_allocator(device: torch.device):
10
+ def alloc_fn(size: int, alignment: int, stream: int | None):
11
+ return torch.empty(size, device=device, dtype=torch.int8)
12
+
13
+ triton.set_allocator(alloc_fn)