[Test][LoRA] Speed up the LoRA test job (#52331)
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
stefankoncarevic committed
Aug 14, 2026 at 19:08 UTC
83ded8d839f60d97f5b020b598001cd9081150f6
4 files changed
+104
-126
tests/lora/test_fused_moe_lora_kernel.py
+4
@@ -22,6 +22,10 @@ from vllm.platforms import current_platform
22
from vllm.utils.network_utils import get_open_port
23
from vllm.utils.torch_utils import set_random_seed
24
25
+# The tensor-parallel cases build their distributed environment inside processes
26
+# spawned by torch.multiprocessing.spawn, so this process never has one to tear down.
27
+pytestmark = pytest.mark.skip_global_cleanup
28
+
29
30
@pytest.fixture(autouse=True)
31
def reset_device(reset_default_device):
tests/lora/test_layers.py
+12
-7
@@ -56,14 +56,19 @@ TOLERANCES = {
56
torch.bfloat16: (3e-2, 2e-2),
57
}
58
59
-pytestmark = pytest.mark.skipif(
60
- not (
61
- current_platform.is_cuda_alike()
62
- or current_platform.is_cpu()
63
- or current_platform.is_xpu()
59
+pytestmark = [
60
+ pytest.mark.skipif(
61
+ not (
62
+ current_platform.is_cuda_alike()
63
+ or current_platform.is_cpu()
64
+ or current_platform.is_xpu()
65
+ ),
66
+ reason="Backend not supported",
67
),
65
- reason="Backend not supported",
66
-)
68
+ # Tests here either take dist_init, which tears the distributed environment
69
+ # down itself, or never build one, so the global cleanup only repeats it.
70
+ pytest.mark.skip_global_cleanup,
71
+]
72
73
DEVICE_TYPE = current_platform.device_type
74
DEVICES = (
tests/lora/test_punica_ops.py
+37
-27
@@ -15,6 +15,14 @@ from .utils import PunicaTensors, assert_close, generate_data_for_nslices
15
16
DEVICE_TYPE = current_platform.device_type
17
18
+# On XPU, oneDNN/oneMKL can return wrong results for these reference matmuls
19
+# after many Triton kernel launches, so the reference stays on CPU there.
20
+_REF_ON_CPU = current_platform.is_xpu()
21
+
22
+
23
+def _to_ref_device(tensor: torch.Tensor) -> torch.Tensor:
24
+ return tensor.cpu() if _REF_ON_CPU else tensor
25
+
26
27
@pytest.fixture(autouse=True)
28
def reset_device(reset_default_device):
@@ -33,10 +41,10 @@ def dynamo_reset():
41
yield
42
43
36
-def _cpu_bgmv_shrink(
44
+def _bgmv_shrink(
45
inputs, lora_weight, output, seq_len_tensor, lora_indices, scaling=1.0
46
):
39
- """Memory-efficient shrink reference: per-LoRA matmul loop on CPU.
47
+ """Memory-efficient shrink reference: per-LoRA matmul loop.
48
output[mask] = scaling * inputs[mask] @ weight.T"""
49
exploded = torch.repeat_interleave(lora_indices, seq_len_tensor)
50
for lid in exploded.unique():
@@ -48,7 +56,7 @@ def _cpu_bgmv_shrink(
56
output[mask] = scaling * (inp @ w.T)
57
58
51
-def _cpu_bgmv_expand(
59
+def _bgmv_expand(
60
inputs,
61
lora_weight,
62
output,
@@ -57,7 +65,7 @@ def _cpu_bgmv_expand(
65
offset=0,
66
add_inputs=False,
67
):
60
- """Memory-efficient expand reference: per-LoRA matmul loop on CPU.
68
+ """Memory-efficient expand reference: per-LoRA matmul loop.
69
output[mask, offset:offset+n] (+)= inputs[mask] @ weight.T"""
70
exploded = torch.repeat_interleave(lora_indices, seq_len_tensor)
71
for lid in exploded.unique():
@@ -88,21 +96,22 @@ def sgmv_shrink_for_nslices(
96
num_tokens: int,
97
scaling: float,
98
):
91
- """CPU reference for sgmv_shrink using per-LoRA matmul loop."""
92
- inp_cpu = inputs_tensor.cpu()
93
- seq_cpu = seq_len_tensor.cpu()
94
- idx_cpu = prompt_lora_mapping.cpu()
95
- out_cpu = out_tensor.cpu()
99
+ """Reference for sgmv_shrink using per-LoRA matmul loop."""
100
+ inputs = _to_ref_device(inputs_tensor)
101
+ seq_len = _to_ref_device(seq_len_tensor)
102
+ mapping = _to_ref_device(prompt_lora_mapping)
103
+ out = _to_ref_device(out_tensor)
104
for index in range(nslices):
97
- _cpu_bgmv_shrink(
98
- inp_cpu,
99
- lora_weights_lst[index].cpu(),
100
- out_cpu[index],
101
- seq_cpu,
102
- idx_cpu,
105
+ _bgmv_shrink(
106
+ inputs,
107
+ _to_ref_device(lora_weights_lst[index]),
108
+ out[index],
109
+ seq_len,
110
+ mapping,
111
scaling=scaling,
112
)
105
- out_tensor.copy_(out_cpu)
113
+ if _REF_ON_CPU:
114
+ out_tensor.copy_(out)
115
116
117
def sgmv_expand_for_nslices(
@@ -119,21 +128,22 @@ def sgmv_expand_for_nslices(
128
num_tokens: int,
129
add_inputs: bool,
130
) -> None:
122
- """CPU reference for sgmv_expand using per-LoRA matmul loop."""
123
- seq_cpu = seq_len_tensor.cpu()
124
- idx_cpu = prompt_lora_mapping.cpu()
125
- out_cpu = out_tensor.cpu()
131
+ """Reference for sgmv_expand using per-LoRA matmul loop."""
132
+ seq_len = _to_ref_device(seq_len_tensor)
133
+ mapping = _to_ref_device(prompt_lora_mapping)
134
+ out = _to_ref_device(out_tensor)
135
for index in range(nslices):
127
- _cpu_bgmv_expand(
128
- inputs_tensor[index].cpu(),
129
- lora_weights_lst[index].cpu(),
130
- out_cpu,
131
- seq_cpu,
132
- idx_cpu,
136
+ _bgmv_expand(
137
+ _to_ref_device(inputs_tensor[index]),
138
+ _to_ref_device(lora_weights_lst[index]),
139
+ out,
140
+ seq_len,
141
+ mapping,
142
offset=hidden_size * index,
143
add_inputs=add_inputs,
144
)
136
- out_tensor.copy_(out_cpu)
145
+ if _REF_ON_CPU:
146
+ out_tensor.copy_(out)
147
148
149
_dict_lock = Lock()
tests/lora/test_punica_ops_fp8.py
+51
-92
@@ -17,6 +17,7 @@ from threading import Lock
17
18
import pytest
19
import torch
20
+import torch.nn.functional as F
21
22
import vllm.lora.ops.torch_ops as torch_ops
23
import vllm.lora.ops.triton_ops as triton_ops
@@ -31,6 +32,8 @@ from vllm.lora.ops.triton_ops.utils import _LORA_A_PTR_DICT, _LORA_B_PTR_DICT
32
from vllm.platforms import current_platform
33
from vllm.utils.torch_utils import set_random_seed
34
35
+pytestmark = pytest.mark.skip_global_cleanup
36
+
37
DEVICE_TYPE = current_platform.device_type
38
DEVICES = [f"{DEVICE_TYPE}:{0}"]
39
SEED = [0]
@@ -206,43 +209,32 @@ def quantize_to_fp8_blockwise(
209
if tensor.ndim == 2:
210
M, K = tensor.shape
211
n_blocks_k = math.ceil(K / group_k)
209
- scale = torch.zeros(M, n_blocks_k, dtype=torch.float32, device=tensor.device)
210
- fp8_tensor = torch.zeros_like(tensor, dtype=FP8_DTYPE)
211
- for m in range(M):
212
- for bk in range(n_blocks_k):
213
- k_start = bk * group_k
214
- k_end = min(k_start + group_k, K)
215
- block = tensor[m, k_start:k_end].float()
216
- amax = block.abs().max().clamp(min=1e-12)
217
- s = (amax / FP8_MAX).to(torch.float32)
218
- scale[m, bk] = s
219
- fp8_tensor[m, k_start:k_end] = (
220
- (block / s).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
221
- )
222
- return fp8_tensor, scale
212
+ # Zero padding leaves the per-block amax unchanged, so the padded
213
+ # columns are dropped again after quantizing.
214
+ padded = F.pad(tensor.float(), (0, n_blocks_k * group_k - K))
215
+ blocks = padded.view(M, n_blocks_k, group_k)
216
+ scale = blocks.abs().amax(dim=-1).clamp(min=1e-12) / FP8_MAX
217
+ fp8_tensor = (
218
+ (blocks / scale.unsqueeze(-1)).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
219
+ )
220
+ return fp8_tensor.view(M, -1)[:, :K].contiguous(), scale
221
elif tensor.ndim == 3:
222
L, N, K = tensor.shape
223
n_blocks_n = math.ceil(N / group_n)
224
n_blocks_k = math.ceil(K / group_k)
227
- scale = torch.zeros(
228
- L, n_blocks_n, n_blocks_k, dtype=torch.float32, device=tensor.device
225
+ padded = F.pad(
226
+ tensor.float(),
227
+ (0, n_blocks_k * group_k - K, 0, n_blocks_n * group_n - N),
228
+ )
229
+ blocks = padded.view(L, n_blocks_n, group_n, n_blocks_k, group_k)
230
+ scale = blocks.abs().amax(dim=(2, 4)).clamp(min=1e-12) / FP8_MAX
231
+ fp8_tensor = (
232
+ (blocks / scale[:, :, None, :, None])
233
+ .clamp(FP8_MIN, FP8_MAX)
234
+ .to(FP8_DTYPE)
235
+ .view(L, n_blocks_n * group_n, n_blocks_k * group_k)
236
)
230
- fp8_tensor = torch.zeros_like(tensor, dtype=FP8_DTYPE)
231
- for li in range(L):
232
- for bn in range(n_blocks_n):
233
- for bk in range(n_blocks_k):
234
- n_start = bn * group_n
235
- n_end = min(n_start + group_n, N)
236
- k_start = bk * group_k
237
- k_end = min(k_start + group_k, K)
238
- block = tensor[li, n_start:n_end, k_start:k_end].float()
239
- amax = block.abs().max().clamp(min=1e-12)
240
- s = (amax / FP8_MAX).to(torch.float32)
241
- scale[li, bn, bk] = s
242
- fp8_tensor[li, n_start:n_end, k_start:k_end] = (
243
- (block / s).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
244
- )
245
- return fp8_tensor, scale
237
+ return fp8_tensor[:, :N, :K].contiguous(), scale
238
else:
239
raise ValueError(f"Unsupported tensor ndim: {tensor.ndim}")
240
@@ -301,33 +293,23 @@ def dequantize_fp8_blockwise(
293
"""Dequantize FP8 tensor with block-wise scale back to output_dtype."""
294
if fp8_tensor.ndim == 2:
295
M, K = fp8_tensor.shape
304
- out = torch.zeros(M, K, dtype=output_dtype, device=fp8_tensor.device)
296
n_blocks_k = math.ceil(K / group_k)
306
- for m in range(M):
307
- for bk in range(n_blocks_k):
308
- k_start = bk * group_k
309
- k_end = min(k_start + group_k, K)
310
- out[m, k_start:k_end] = (
311
- fp8_tensor[m, k_start:k_end].float() * scale[m, bk].float()
312
- ).to(output_dtype)
313
- return out
297
+ padded = F.pad(fp8_tensor.float(), (0, n_blocks_k * group_k - K))
298
+ blocks = padded.view(M, n_blocks_k, group_k)
299
+ out = (blocks * scale.float().unsqueeze(-1)).to(output_dtype)
300
+ return out.view(M, -1)[:, :K].contiguous()
301
elif fp8_tensor.ndim == 3:
302
L, N, K = fp8_tensor.shape
316
- out = torch.zeros(L, N, K, dtype=output_dtype, device=fp8_tensor.device)
303
n_blocks_n = math.ceil(N / group_n)
304
n_blocks_k = math.ceil(K / group_k)
319
- for l_idx in range(L):
320
- for bn in range(n_blocks_n):
321
- for bk in range(n_blocks_k):
322
- n_start = bn * group_n
323
- n_end = min(n_start + group_n, N)
324
- k_start = bk * group_k
325
- k_end = min(k_start + group_k, K)
326
- out[l_idx, n_start:n_end, k_start:k_end] = (
327
- fp8_tensor[l_idx, n_start:n_end, k_start:k_end].float()
328
- * scale[l_idx, bn, bk].float()
329
- ).to(output_dtype)
330
- return out
305
+ padded = F.pad(
306
+ fp8_tensor.float(),
307
+ (0, n_blocks_k * group_k - K, 0, n_blocks_n * group_n - N),
308
+ )
309
+ blocks = padded.view(L, n_blocks_n, group_n, n_blocks_k, group_k)
310
+ out = (blocks * scale.float()[:, :, None, :, None]).to(output_dtype)
311
+ out = out.view(L, n_blocks_n * group_n, n_blocks_k * group_k)
312
+ return out[:, :N, :K].contiguous()
313
else:
314
raise ValueError(f"Unsupported tensor ndim: {fp8_tensor.ndim}")
315
@@ -519,45 +501,22 @@ def generate_fp8_expand_data(
501
# shared across slices. Compute shared scale across slices, then quantize.
502
# First compute per-token-per-block scale across all slices
503
n_blocks_k = math.ceil(rank / group_k)
522
- a_scale = torch.zeros(
523
- total_tokens, n_blocks_k, dtype=torch.float32, device=device
504
+ padded = F.pad(inputs_bf16.float(), (0, n_blocks_k * group_k - rank))
505
+ blocks = padded.view(nslices, total_tokens, n_blocks_k, group_k)
506
+ # Take the block amax across all slices so every slice shares the scale.
507
+ a_scale = blocks.abs().amax(dim=-1).clamp(min=1e-12).amax(dim=0) / FP8_MAX
508
+ fp8_blocks = (
509
+ (blocks / a_scale[None, :, :, None]).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
510
+ )
511
+ inputs_fp8 = fp8_blocks.view(nslices, total_tokens, -1)[
512
+ :, :, :rank
513
+ ].contiguous()
514
+ inputs_dequant = (
515
+ (fp8_blocks.float() * a_scale[None, :, :, None])
516
+ .to(dtype)
517
+ .view(nslices, total_tokens, -1)[:, :, :rank]
518
+ .contiguous()
519
)
525
- for m in range(total_tokens):
526
- for bk in range(n_blocks_k):
527
- k_start = bk * group_k
528
- k_end = min(k_start + group_k, rank)
529
- # Max across all slices for this token and block
530
- block_amax = torch.tensor(0.0, device=device)
531
- for s in range(nslices):
532
- block = inputs_bf16[s, m, k_start:k_end].float()
533
- block_amax = torch.max(
534
- block_amax, block.abs().max().clamp(min=1e-12)
535
- )
536
- a_scale[m, bk] = (block_amax / FP8_MAX).to(torch.float32)
537
-
538
- # Quantize all slices with the shared scale
539
- inputs_fp8_list = []
540
- inputs_dequant_list = []
541
- for s in range(nslices):
542
- slice_2d = inputs_bf16[s] # (total_tokens, rank)
543
- fp8_slice = torch.zeros_like(slice_2d, dtype=FP8_DTYPE)
544
- dequant_slice = torch.zeros_like(slice_2d)
545
- for m in range(total_tokens):
546
- for bk in range(n_blocks_k):
547
- k_start = bk * group_k
548
- k_end = min(k_start + group_k, rank)
549
- block = slice_2d[m, k_start:k_end].float()
550
- s_val = a_scale[m, bk]
551
- fp8_slice[m, k_start:k_end] = (
552
- (block / s_val).clamp(FP8_MIN, FP8_MAX).to(FP8_DTYPE)
553
- )
554
- dequant_slice[m, k_start:k_end] = (
555
- fp8_slice[m, k_start:k_end].float() * s_val.float()
556
- ).to(dtype)
557
- inputs_fp8_list.append(fp8_slice)
558
- inputs_dequant_list.append(dequant_slice)
559
- inputs_fp8 = torch.stack(inputs_fp8_list, dim=0)
560
- inputs_dequant = torch.stack(inputs_dequant_list, dim=0)
520
elif quant_mode == "per_tensor":
521
# Per-tensor: kernel loads a single scalar from a_scale_ptr
522
inputs_fp8_2d, a_scale = quantize_to_fp8_per_tensor(inputs_2d_all)