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
Remove all cases of `fmt: on/off` (#26253)
Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Harry Mellor committed
Oct 5, 2025 at 17:18 UTC
557b2e961d8e6fe47c67be61d089275102573590
5 files changed
+217
-157
benchmarks/kernels/deepgemm/benchmark_fp8_block_dense_gemm.py
+117
-109
@@ -1,6 +1,5 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-# fmt: off
3
# ruff: noqa: E501
4
import time
5
@@ -20,19 +19,21 @@ from vllm.utils.deep_gemm import (
19
)
20
21
23
-def benchmark_shape(m: int,
24
- n: int,
25
- k: int,
26
- warmup: int = 100,
27
- repeat: int = 10000,
28
- verbose: bool = False) -> dict:
22
+def benchmark_shape(
23
+ m: int,
24
+ n: int,
25
+ k: int,
26
+ warmup: int = 100,
27
+ repeat: int = 10000,
28
+ verbose: bool = False,
29
+) -> dict:
30
"""Benchmark all implementations for a specific (m, n, k) shape."""
31
if verbose:
32
print(f"\n=== Benchmarking shape: m={m}, n={n}, k={k} ===")
33
34
# Create test tensors
34
- A = torch.randn((m, k), device='cuda', dtype=torch.bfloat16)
35
- B = torch.randn((n, k), device='cuda', dtype=torch.bfloat16)
35
+ A = torch.randn((m, k), device="cuda", dtype=torch.bfloat16)
36
+ B = torch.randn((n, k), device="cuda", dtype=torch.bfloat16)
37
38
# Reference result in BF16
39
torch.cuda.synchronize()
@@ -49,34 +50,39 @@ def benchmark_shape(m: int,
50
# Pre-quantize A for all implementations
51
A_deepgemm, A_scale_deepgemm = per_token_group_quant_fp8(A, block_size[1])
52
A_scale_deepgemm = get_col_major_tma_aligned_tensor(A_scale_deepgemm)
52
- C_deepgemm = torch.empty((m, n), device='cuda', dtype=torch.bfloat16)
53
+ C_deepgemm = torch.empty((m, n), device="cuda", dtype=torch.bfloat16)
54
A_vllm, A_scale_vllm = per_token_group_quant_fp8(A, block_size[1])
55
A_vllm_cutlass, A_scale_vllm_cutlass = per_token_group_quant_fp8(
55
- A, block_size[1], column_major_scales=True)
56
+ A, block_size[1], column_major_scales=True
57
+ )
58
59
# === DeepGEMM Implementation ===
60
def deepgemm_gemm():
59
- fp8_gemm_nt((A_deepgemm, A_scale_deepgemm),
60
- (B_deepgemm, B_scale_deepgemm),
61
- C_deepgemm)
61
+ fp8_gemm_nt(
62
+ (A_deepgemm, A_scale_deepgemm), (B_deepgemm, B_scale_deepgemm), C_deepgemm
63
+ )
64
return C_deepgemm
65
66
# === vLLM Triton Implementation ===
67
def vllm_triton_gemm():
66
- return w8a8_triton_block_scaled_mm(A_vllm,
67
- B_vllm,
68
- A_scale_vllm,
69
- B_scale_vllm,
70
- block_size,
71
- output_dtype=torch.bfloat16)
68
+ return w8a8_triton_block_scaled_mm(
69
+ A_vllm,
70
+ B_vllm,
71
+ A_scale_vllm,
72
+ B_scale_vllm,
73
+ block_size,
74
+ output_dtype=torch.bfloat16,
75
+ )
76
77
# === vLLM CUTLASS Implementation ===
78
def vllm_cutlass_gemm():
75
- return ops.cutlass_scaled_mm(A_vllm_cutlass,
76
- B_vllm.T,
77
- scale_a=A_scale_vllm_cutlass,
78
- scale_b=B_scale_vllm.T,
79
- out_dtype=torch.bfloat16)
79
+ return ops.cutlass_scaled_mm(
80
+ A_vllm_cutlass,
81
+ B_vllm.T,
82
+ scale_a=A_scale_vllm_cutlass,
83
+ scale_b=B_scale_vllm.T,
84
+ out_dtype=torch.bfloat16,
85
+ )
86
87
# Run correctness check first
88
if verbose:
@@ -93,26 +99,23 @@ def benchmark_shape(m: int,
99
print(f"DeepGEMM vs Reference difference: {deepgemm_diff:.6f}")
100
print(f"vLLM Triton vs Reference difference: {vllm_triton_diff:.6f}")
101
print(f"vLLM CUTLASS vs Reference difference: {vllm_cutlass_diff:.6f}")
96
- print("vLLM Triton vs DeepGEMM difference: "
97
- f"{calc_diff(C_vllm_triton, C_deepgemm):.6f}")
98
- print("vLLM CUTLASS vs DeepGEMM difference: "
99
- f"{calc_diff(C_vllm_cutlass, C_deepgemm):.6f}")
102
+ print(
103
+ "vLLM Triton vs DeepGEMM difference: "
104
+ f"{calc_diff(C_vllm_triton, C_deepgemm):.6f}"
105
+ )
106
+ print(
107
+ "vLLM CUTLASS vs DeepGEMM difference: "
108
+ f"{calc_diff(C_vllm_cutlass, C_deepgemm):.6f}"
109
+ )
110
111
# Benchmark implementations
112
implementations = {
113
"DeepGEMM": deepgemm_gemm,
114
"vLLM Triton": vllm_triton_gemm,
105
- "vLLM CUTLASS": vllm_cutlass_gemm
115
+ "vLLM CUTLASS": vllm_cutlass_gemm,
116
}
117
108
- benchmark_results = {
109
- "shape": {
110
- "m": m,
111
- "n": n,
112
- "k": k
113
- },
114
- "implementations": {}
115
- }
118
+ benchmark_results = {"shape": {"m": m, "n": n, "k": k}, "implementations": {}}
119
120
for name, func in implementations.items():
121
# Warmup
@@ -140,38 +143,36 @@ def benchmark_shape(m: int,
143
"tflops": tflops,
144
"gb_s": gb_s,
145
"diff": {
143
- "DeepGEMM":
144
- 0.0 if name == "DeepGEMM" else calc_diff(func(), C_deepgemm),
145
- "Reference":
146
- deepgemm_diff if name == "DeepGEMM" else
147
- (vllm_triton_diff
148
- if name == "vLLM Triton" else vllm_cutlass_diff)
149
- }
146
+ "DeepGEMM": 0.0
147
+ if name == "DeepGEMM"
148
+ else calc_diff(func(), C_deepgemm),
149
+ "Reference": deepgemm_diff
150
+ if name == "DeepGEMM"
151
+ else (vllm_triton_diff if name == "vLLM Triton" else vllm_cutlass_diff),
152
+ },
153
}
154
155
if verbose:
153
- print(
154
- f"{name}: {avg_time_ms:.3f} ms, {tflops:.2f} TFLOPS, {gb_s:.2f} GB/s"
155
- )
156
+ print(f"{name}: {avg_time_ms:.3f} ms, {tflops:.2f} TFLOPS, {gb_s:.2f} GB/s")
157
158
# Calculate speedups
159
baseline = benchmark_results["implementations"]["DeepGEMM"]["time_ms"]
160
for name, data in benchmark_results["implementations"].items():
161
if name != "DeepGEMM":
162
speedup = baseline / data["time_ms"]
162
- benchmark_results["implementations"][name][
163
- "speedup_vs_deepgemm"] = speedup
163
+ benchmark_results["implementations"][name]["speedup_vs_deepgemm"] = speedup
164
if verbose:
165
- print(f"DeepGEMM is {1/speedup:.2f}x "
166
- f"{'faster' if 1/speedup > 1 else 'slower'} than {name}")
165
+ print(
166
+ f"DeepGEMM is {1 / speedup:.2f}x "
167
+ f"{'faster' if 1 / speedup > 1 else 'slower'} than {name}"
168
+ )
169
168
- vllm_triton_time = benchmark_results["implementations"]["vLLM Triton"][
169
- "time_ms"]
170
- vllm_cutlass_time = benchmark_results["implementations"]["vLLM CUTLASS"][
171
- "time_ms"]
170
+ vllm_triton_time = benchmark_results["implementations"]["vLLM Triton"]["time_ms"]
171
+ vllm_cutlass_time = benchmark_results["implementations"]["vLLM CUTLASS"]["time_ms"]
172
cutlass_vs_triton = vllm_triton_time / vllm_cutlass_time
173
- benchmark_results["implementations"]["vLLM CUTLASS"][
174
- "speedup_vs_triton"] = cutlass_vs_triton
173
+ benchmark_results["implementations"]["vLLM CUTLASS"]["speedup_vs_triton"] = (
174
+ cutlass_vs_triton
175
+ )
176
if verbose:
177
print(
178
f"vLLM CUTLASS is {cutlass_vs_triton:.2f}x "
@@ -183,8 +184,7 @@ def benchmark_shape(m: int,
184
185
def format_table_row(values, widths):
186
"""Format a row with specified column widths."""
186
- return "| " + " | ".join(f"{val:{w}}"
187
- for val, w in zip(values, widths)) + " |"
187
+ return "| " + " | ".join(f"{val:{w}}" for val, w in zip(values, widths)) + " |"
188
189
190
def print_table(headers, rows, title=None):
@@ -292,38 +292,50 @@ def run_benchmarks(verbose: bool = False):
292
for result in all_results:
293
shape = result["shape"]
294
impl_data = result["implementations"]["DeepGEMM"]
295
- deepgemm_rows.append([
296
- shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
297
- f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}"
298
- ])
295
+ deepgemm_rows.append(
296
+ [
297
+ shape["m"],
298
+ shape["n"],
299
+ shape["k"],
300
+ f"{impl_data['time_us']:.1f}",
301
+ f"{impl_data['tflops']:.1f}",
302
+ f"{impl_data['gb_s']:.1f}",
303
+ ]
304
+ )
305
300
- print_table(deepgemm_headers,
301
- deepgemm_rows,
302
- title="DeepGEMM Implementation:")
306
+ print_table(deepgemm_headers, deepgemm_rows, title="DeepGEMM Implementation:")
307
308
# Print vLLM Triton table
305
- triton_headers = [
306
- "m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM"
307
- ]
309
+ triton_headers = ["m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM"]
310
triton_rows = []
311
for result in all_results:
312
shape = result["shape"]
313
impl_data = result["implementations"]["vLLM Triton"]
314
speedup = impl_data.get("speedup_vs_deepgemm", 1.0)
313
- triton_rows.append([
314
- shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
315
- f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}",
316
- format_speedup(speedup)
317
- ])
315
+ triton_rows.append(
316
+ [
317
+ shape["m"],
318
+ shape["n"],
319
+ shape["k"],
320
+ f"{impl_data['time_us']:.1f}",
321
+ f"{impl_data['tflops']:.1f}",
322
+ f"{impl_data['gb_s']:.1f}",
323
+ format_speedup(speedup),
324
+ ]
325
+ )
326
319
- print_table(triton_headers,
320
- triton_rows,
321
- title="vLLM Triton Implementation:")
327
+ print_table(triton_headers, triton_rows, title="vLLM Triton Implementation:")
328
329
# Print vLLM CUTLASS table
330
cutlass_headers = [
325
- "m", "n", "k", "Time (μs)", "TFLOPS", "GB/s", "vs DeepGEMM",
326
- "vs Triton"
331
+ "m",
332
+ "n",
333
+ "k",
334
+ "Time (μs)",
335
+ "TFLOPS",
336
+ "GB/s",
337
+ "vs DeepGEMM",
338
+ "vs Triton",
339
]
340
cutlass_rows = []
341
for result in all_results:
@@ -331,28 +343,27 @@ def run_benchmarks(verbose: bool = False):
343
impl_data = result["implementations"]["vLLM CUTLASS"]
344
vs_deepgemm = impl_data.get("speedup_vs_deepgemm", 1.0)
345
vs_triton = impl_data.get("speedup_vs_triton", 1.0)
334
- cutlass_rows.append([
335
- shape["m"], shape["n"], shape["k"], f"{impl_data['time_us']:.1f}",
336
- f"{impl_data['tflops']:.1f}", f"{impl_data['gb_s']:.1f}",
337
- format_speedup(vs_deepgemm),
338
- format_speedup(vs_triton)
339
- ])
346
+ cutlass_rows.append(
347
+ [
348
+ shape["m"],
349
+ shape["n"],
350
+ shape["k"],
351
+ f"{impl_data['time_us']:.1f}",
352
+ f"{impl_data['tflops']:.1f}",
353
+ f"{impl_data['gb_s']:.1f}",
354
+ format_speedup(vs_deepgemm),
355
+ format_speedup(vs_triton),
356
+ ]
357
+ )
358
341
- print_table(cutlass_headers,
342
- cutlass_rows,
343
- title="vLLM CUTLASS Implementation:")
359
+ print_table(cutlass_headers, cutlass_rows, title="vLLM CUTLASS Implementation:")
360
361
# Calculate and print averages
362
print("\n===== AVERAGE PERFORMANCE =====")
363
364
implementations = ["DeepGEMM", "vLLM Triton", "vLLM CUTLASS"]
365
avg_metrics = {
350
- impl: {
351
- "tflops": 0,
352
- "gb_s": 0,
353
- "time_ms": 0
354
- }
355
- for impl in implementations
366
+ impl: {"tflops": 0, "gb_s": 0, "time_ms": 0} for impl in implementations
367
}
368
369
for result in all_results:
@@ -370,9 +381,9 @@ def run_benchmarks(verbose: bool = False):
381
avg_tflops = avg_metrics[impl]["tflops"] / num_shapes
382
avg_mem_bw = avg_metrics[impl]["gb_s"] / num_shapes
383
avg_time = avg_metrics[impl]["time_ms"] / num_shapes
373
- avg_rows.append([
374
- impl, f"{avg_tflops:.2f}", f"{avg_mem_bw:.2f}", f"{avg_time:.2f}"
375
- ])
384
+ avg_rows.append(
385
+ [impl, f"{avg_tflops:.2f}", f"{avg_mem_bw:.2f}", f"{avg_time:.2f}"]
386
+ )
387
388
print_table(avg_headers, avg_rows)
389
@@ -380,21 +391,19 @@ def run_benchmarks(verbose: bool = False):
391
avg_speedups = {
392
"DeepGEMM vs vLLM Triton": 0,
393
"DeepGEMM vs vLLM CUTLASS": 0,
383
- "vLLM CUTLASS vs vLLM Triton": 0
394
+ "vLLM CUTLASS vs vLLM Triton": 0,
395
}
396
397
for result in all_results:
398
deepgemm_time = result["implementations"]["DeepGEMM"]["time_ms"]
399
vllm_triton_time = result["implementations"]["vLLM Triton"]["time_ms"]
389
- vllm_cutlass_time = result["implementations"]["vLLM CUTLASS"][
390
- "time_ms"]
400
+ vllm_cutlass_time = result["implementations"]["vLLM CUTLASS"]["time_ms"]
401
392
- avg_speedups[
393
- "DeepGEMM vs vLLM Triton"] += vllm_triton_time / deepgemm_time
394
- avg_speedups[
395
- "DeepGEMM vs vLLM CUTLASS"] += vllm_cutlass_time / deepgemm_time
396
- avg_speedups[
397
- "vLLM CUTLASS vs vLLM Triton"] += vllm_triton_time / vllm_cutlass_time
402
+ avg_speedups["DeepGEMM vs vLLM Triton"] += vllm_triton_time / deepgemm_time
403
+ avg_speedups["DeepGEMM vs vLLM CUTLASS"] += vllm_cutlass_time / deepgemm_time
404
+ avg_speedups["vLLM CUTLASS vs vLLM Triton"] += (
405
+ vllm_triton_time / vllm_cutlass_time
406
+ )
407
408
print("\n===== AVERAGE SPEEDUPS =====")
409
speedup_headers = ["Comparison", "Speedup"]
@@ -412,8 +421,7 @@ def run_benchmarks(verbose: bool = False):
421
422
for result in all_results:
423
for impl in implementations:
415
- avg_diff[impl] += result["implementations"][impl]["diff"][
416
- "Reference"]
424
+ avg_diff[impl] += result["implementations"][impl]["diff"]["Reference"]
425
426
diff_headers = ["Implementation", "Avg Diff vs Reference"]
427
diff_rows = []
tests/kernels/mamba/test_mamba_ssm_ssd.py
+47
-24
@@ -442,14 +442,22 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
442
B_chunked = torch.zeros_like(B)[:chunked_input_seq_len, ...]
443
C_chunked = torch.zeros_like(C)[:chunked_input_seq_len, ...]
444
for i in range(num_sequences):
445
- # fmt: off
446
- chunk_f = lambda x, i: x[cu_seqlens[i]:cu_seqlens[i] + chunked_seqlens[i], ...] # noqa: E501
445
+ chunk_f = lambda x, i: x[
446
+ cu_seqlens[i] : cu_seqlens[i] + chunked_seqlens[i], ...
447
+ ]
448
448
- X_chunked[chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(X, i) # noqa: E501
449
- dt_chunked[chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(dt, i) # noqa: E501
450
- B_chunked[chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(B, i) # noqa: E501
451
- C_chunked[chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1], ...] = chunk_f(C, i) # noqa: E501
452
- # fmt: on
449
+ X_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
450
+ X, i
451
+ )
452
+ dt_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
453
+ dt, i
454
+ )
455
+ B_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
456
+ B, i
457
+ )
458
+ C_chunked[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...] = chunk_f(
459
+ C, i
460
+ )
461
462
cu_chunk_seqlens, last_chunk_indices, seq_idx_chunks = (
463
compute_varlen_chunk_metadata(chunked_cu_seqlens, chunk_size)
@@ -481,27 +489,42 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens):
489
dim=0,
490
)
491
remaining_chunked_input_seq_len = remaining_chunked_cu_seqlens[-1]
484
- # fmt: off
485
- remaining_X_chunked = torch.zeros_like(X)[:remaining_chunked_input_seq_len, ...] # noqa: E501
486
- remaining_dt_chunked = torch.zeros_like(dt)[:remaining_chunked_input_seq_len, ...] # noqa: E501
487
- remaining_B_chunked = torch.zeros_like(B)[:remaining_chunked_input_seq_len, ...] # noqa: E501
488
- remaining_C_chunked = torch.zeros_like(C)[:remaining_chunked_input_seq_len, ...] # noqa: E501
492
+ remaining_X_chunked = torch.zeros_like(X)[:remaining_chunked_input_seq_len, ...]
493
+ remaining_dt_chunked = torch.zeros_like(dt)[:remaining_chunked_input_seq_len, ...]
494
+ remaining_B_chunked = torch.zeros_like(B)[:remaining_chunked_input_seq_len, ...]
495
+ remaining_C_chunked = torch.zeros_like(C)[:remaining_chunked_input_seq_len, ...]
496
for i in range(num_sequences):
490
- remaining_chunk_f = lambda x, i: x[cu_seqlens[i] + chunked_seqlens[i]:cu_seqlens[i+1], ...] # noqa: E501
491
-
492
- remaining_X_chunked[remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(X, i) # noqa: E501
493
- remaining_dt_chunked[remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(dt, i) # noqa: E501
494
- remaining_B_chunked[remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(B, i) # noqa: E501
495
- remaining_C_chunked[remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1], ...] = remaining_chunk_f(C, i) # noqa: E501
497
+ remaining_chunk_f = lambda x, i: x[
498
+ cu_seqlens[i] + chunked_seqlens[i] : cu_seqlens[i + 1], ...
499
+ ]
500
+
501
+ remaining_X_chunked[
502
+ remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
503
+ ] = remaining_chunk_f(X, i)
504
+ remaining_dt_chunked[
505
+ remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
506
+ ] = remaining_chunk_f(dt, i)
507
+ remaining_B_chunked[
508
+ remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
509
+ ] = remaining_chunk_f(B, i)
510
+ remaining_C_chunked[
511
+ remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1], ...
512
+ ] = remaining_chunk_f(C, i)
513
514
# assert input chunking is correct
498
- concat_chunk_f = lambda pt1, pt2, i: torch.cat([
499
- pt1[chunked_cu_seqlens[i]:chunked_cu_seqlens[i+1],...],
500
- pt2[remaining_chunked_cu_seqlens[i]:remaining_chunked_cu_seqlens[i+1],...],
515
+ concat_chunk_f = lambda pt1, pt2, i: torch.cat(
516
+ [
517
+ pt1[chunked_cu_seqlens[i] : chunked_cu_seqlens[i + 1], ...],
518
+ pt2[
519
+ remaining_chunked_cu_seqlens[i] : remaining_chunked_cu_seqlens[i + 1],
520
+ ...,
521
+ ],
522
],
502
- dim=0)
503
- concat_batch_f = lambda pt1, pt2: torch.cat([concat_chunk_f(pt1, pt2, i) for i in range(num_sequences)], dim=0) # noqa: E501
504
- # fmt: on
523
+ dim=0,
524
+ )
525
+ concat_batch_f = lambda pt1, pt2: torch.cat(
526
+ [concat_chunk_f(pt1, pt2, i) for i in range(num_sequences)], dim=0
527
+ )
528
529
assert concat_batch_f(X_chunked, remaining_X_chunked).equal(X)
530
assert concat_batch_f(dt_chunked, remaining_dt_chunked).equal(dt)
tests/v1/engine/test_fast_incdec_prefix_err.py
+5
-3
@@ -182,8 +182,9 @@ def test_fast_inc_detok_invalid_utf8_err_case():
182
finished = i == len(test_tokens) - 1
183
output += detokenizer.get_next_output_text(finished, delta=True)
184
185
- # fmt: off
186
- assert output == r'''[
185
+ assert (
186
+ output
187
+ == r"""[
188
{
189
"source": "Résultats",
190
"source_type": "CONCEPT",
@@ -191,4 +192,5 @@ def test_fast_inc_detok_invalid_utf8_err_case():
192
"target": "Israël",
193
"target_type": "ORGANIZATION",
194
"target_description": "Pays qui a obtenu à sa frontière libanaise « un niveau de calme inédit depuis les années 1960 »",
194
- "relationship": "Obtention d'un niveau de'''
195
+ "relationship": "Obtention d'un niveau de"""
196
+ )
vllm/model_executor/models/qwen2_5_vl.py
+4
-6
@@ -1398,12 +1398,10 @@ class Qwen2_5_VLForConditionalGeneration(
1398
# Tensors
1399
input_ids_t = torch.as_tensor(input_ids, device=device, dtype=torch.long)
1400
1401
- # fmt: off
1402
- mm_embeddings_out = [mm[:, :-4] for mm in
1403
- multimodal_embeddings]
1404
- mm_embeddings_pos = [mm[:, -4:].permute(1, 0).long() for mm in
1405
- multimodal_embeddings]
1406
- # fmt: in
1401
+ mm_embeddings_out = [mm[:, :-4] for mm in multimodal_embeddings]
1402
+ mm_embeddings_pos = [
1403
+ mm[:, -4:].permute(1, 0).long() for mm in multimodal_embeddings
1404
+ ]
1405
1406
positions, mrope_positions_delta = recompute_mrope_positions(
1407
input_ids_t,
vllm/model_executor/models/voxtral.py
+44
-15
@@ -516,14 +516,18 @@ class VoxtralForConditionalGeneration(
516
)
517
518
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
519
- # fmt: off
519
remapping_rules = [
520
(r"mm_whisper_embeddings\.(.*)", r"\1"),
521
(r"audio_language_projection\.(.*)", r"audio_language_adapter.\1"),
523
- (r"audio_language_adapter\.0\.weight", r"audio_language_adapter.w_in.weight"), # noqa: E501
524
- (r"audio_language_adapter\.2\.weight", r"audio_language_adapter.w_out.weight"), # noqa: E501
522
+ (
523
+ r"audio_language_adapter\.0\.weight",
524
+ r"audio_language_adapter.w_in.weight",
525
+ ),
526
+ (
527
+ r"audio_language_adapter\.2\.weight",
528
+ r"audio_language_adapter.w_out.weight",
529
+ ),
530
]
526
- # fmt: on
531
532
audio_params = dict(
533
nn.ModuleDict(
@@ -678,19 +682,44 @@ class AudioLanguageAdapter(nn.Module):
682
class VoxtralEncoderModel(nn.Module):
683
packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
684
681
- # fmt: off
685
mistral_remapping = [
683
- (r"whisper_encoder\.conv_layers\.0\.(weight|bias)", r"whisper_encoder.conv1.\1"), # noqa: E501
684
- (r"whisper_encoder\.conv_layers\.1\.(weight|bias)", r"whisper_encoder.conv2.\1"), # noqa: E501
685
- (r"whisper_encoder\.transformer\.layers\.(\d+)\.attention\.w([qkv])\.(weight|bias)", r"whisper_encoder.layers.\1.self_attn.\2_proj.\3"), # noqa: E501
686
- (r"whisper_encoder\.transformer\.layers\.(\d+)\.attention\.wo\.(weight|bias)", r"whisper_encoder.layers.\1.self_attn.out_proj.\2"), # noqa: E501
687
- (r"whisper_encoder\.transformer\.layers\.(\d+)\.attention_norm\.(weight|bias)", r"whisper_encoder.layers.\1.self_attn_layer_norm.\2"), # noqa: E501
688
- (r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w1\.(weight|bias)", r"whisper_encoder.layers.\1.mlp.fc1.\2"), # noqa: E501
689
- (r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w2\.(weight|bias)", r"whisper_encoder.layers.\1.mlp.fc2.\2"), # noqa: E501
690
- (r"whisper_encoder\.transformer\.layers\.(\d+)\.ffn_norm\.(weight|bias)", r"whisper_encoder.layers.\1.final_layer_norm.\2"), # noqa: E501
691
- (r"whisper_encoder\.transformer\.norm\.(weight|bias)", r"whisper_encoder.layer_norm.\1"), # noqa: E501
686
+ (
687
+ r"whisper_encoder\.conv_layers\.0\.(weight|bias)",
688
+ r"whisper_encoder.conv1.\1",
689
+ ),
690
+ (
691
+ r"whisper_encoder\.conv_layers\.1\.(weight|bias)",
692
+ r"whisper_encoder.conv2.\1",
693
+ ),
694
+ (
695
+ r"whisper_encoder\.transformer\.layers\.(\d+)\.attention\.w([qkv])\.(weight|bias)", # noqa: E501
696
+ r"whisper_encoder.layers.\1.self_attn.\2_proj.\3",
697
+ ),
698
+ (
699
+ r"whisper_encoder\.transformer\.layers\.(\d+)\.attention\.wo\.(weight|bias)", # noqa: E501
700
+ r"whisper_encoder.layers.\1.self_attn.out_proj.\2",
701
+ ),
702
+ (
703
+ r"whisper_encoder\.transformer\.layers\.(\d+)\.attention_norm\.(weight|bias)", # noqa: E501
704
+ r"whisper_encoder.layers.\1.self_attn_layer_norm.\2",
705
+ ),
706
+ (
707
+ r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w1\.(weight|bias)", # noqa: E501
708
+ r"whisper_encoder.layers.\1.mlp.fc1.\2",
709
+ ),
710
+ (
711
+ r"whisper_encoder\.transformer\.layers\.(\d+)\.feed_forward\.w2\.(weight|bias)", # noqa: E501
712
+ r"whisper_encoder.layers.\1.mlp.fc2.\2",
713
+ ),
714
+ (
715
+ r"whisper_encoder\.transformer\.layers\.(\d+)\.ffn_norm\.(weight|bias)",
716
+ r"whisper_encoder.layers.\1.final_layer_norm.\2",
717
+ ),
718
+ (
719
+ r"whisper_encoder\.transformer\.norm\.(weight|bias)",
720
+ r"whisper_encoder.layer_norm.\1",
721
+ ),
722
]
693
- # fmt: on
723
724
def __init__(
725
self,