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 6, 04:30 UTC
[Sparse24] [Deprecation] Remove Sparse24 CT integration and kernels (#36799)
Signed-off-by: Kyle Sayers <kylesayrs@gmail.com>
Kyle Sayers committed
Mar 23, 2026 at 16:03 UTC
38364a7e32c923cfbd71f4d749fb2791c1741bd4
17 files changed
+9
-2674
.buildkite/lm-eval-harness/configs/SparseLlama3.1_2of4_fp8_compressed.yaml
deleted
-12
@@ -1,12 +0,0 @@
1
-# For vllm script, with -t option (tensor parallel size).
2
-# bash ./run-lm-eval-gsm-vllm-baseline.sh -m nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_fp8-BitM -b "auto" -t 2
3
-model_name: "nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_fp8-BitM"
4
-tasks:
5
-- name: "gsm8k"
6
- metrics:
7
- - name: "exact_match,strict-match"
8
- value: 0.6353
9
- - name: "exact_match,flexible-extract"
10
- value: 0.637
11
-limit: null
12
-num_fewshot: null
CMakeLists.txt
-26
@@ -343,7 +343,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
343
"csrc/quantization/w8a8/cutlass/scaled_mm_entry.cu"
344
"csrc/quantization/fp4/nvfp4_quant_entry.cu"
345
"csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
346
- "csrc/sparse/cutlass/sparse_scaled_mm_entry.cu"
346
"csrc/cutlass_extensions/common.cpp"
347
"csrc/quantization/w8a8/fp8/per_token_group_quant.cu"
348
"csrc/quantization/w8a8/int8/per_token_group_quant.cu")
@@ -619,31 +618,6 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
618
endif()
619
endif()
620
622
- #
623
- # 2:4 Sparse Kernels
624
-
625
- # The 2:4 sparse kernels cutlass_scaled_sparse_mm and cutlass_compressor
626
- # require CUDA 12.2 or later (and only work on Hopper).
627
- cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}")
628
- if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.2 AND SCALED_MM_ARCHS)
629
- set(SRCS "csrc/sparse/cutlass/sparse_scaled_mm_c3x.cu")
630
- set_gencode_flags_for_srcs(
631
- SRCS "${SRCS}"
632
- CUDA_ARCHS "${SCALED_MM_ARCHS}")
633
- list(APPEND VLLM_EXT_SRC "${SRCS}")
634
- list(APPEND VLLM_GPU_FLAGS "-DENABLE_SPARSE_SCALED_MM_C3X=1")
635
- message(STATUS "Building sparse_scaled_mm_c3x for archs: ${SCALED_MM_ARCHS}")
636
- else()
637
- if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.2 AND SCALED_MM_ARCHS)
638
- message(STATUS "Not building sparse_scaled_mm_c3x kernels as CUDA Compiler version is "
639
- "not >= 12.2, we recommend upgrading to CUDA 12.2 or later "
640
- "if you intend on running FP8 sparse quantized models on Hopper.")
641
- else()
642
- message(STATUS "Not building sparse_scaled_mm_c3x as no compatible archs found "
643
- "in CUDA target architectures")
644
- endif()
645
- endif()
646
-
621
# The nvfp4_scaled_mm_sm120 kernels for Geforce Blackwell SM120 require
622
# CUDA 12.8 or later
623
if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 13.0)
benchmarks/cutlass_benchmarks/sparse_benchmarks.py
deleted
-517
@@ -1,517 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-
4
-import argparse
5
-import copy
6
-import itertools
7
-import pickle as pkl
8
-import time
9
-from collections.abc import Callable, Iterable
10
-
11
-import torch
12
-import torch.utils.benchmark as TBenchmark
13
-from torch.utils.benchmark import Measurement as TMeasurement
14
-from utils import make_rand_sparse_tensors
15
-from weight_shapes import WEIGHT_SHAPES
16
-
17
-from vllm import _custom_ops as ops
18
-from vllm.utils.argparse_utils import FlexibleArgumentParser
19
-
20
-DEFAULT_MODELS = list(WEIGHT_SHAPES.keys())
21
-DEFAULT_BATCH_SIZES = [1, 16, 32, 64, 128, 256, 512]
22
-DEFAULT_TP_SIZES = [1]
23
-
24
-
25
-# bench
26
-def bench_fn(
27
- label: str, sub_label: str, description: str, fn: Callable, *args, **kwargs
28
-) -> TMeasurement:
29
- min_run_time = 1
30
-
31
- globals = {
32
- "args": args,
33
- "kwargs": kwargs,
34
- "fn": fn,
35
- }
36
- return TBenchmark.Timer(
37
- stmt="fn(*args, **kwargs)",
38
- globals=globals,
39
- label=label,
40
- sub_label=sub_label,
41
- description=description,
42
- ).blocked_autorange(min_run_time=min_run_time)
43
-
44
-
45
-def bench_int8(
46
- dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str
47
-) -> Iterable[TMeasurement]:
48
- assert dtype == torch.int8
49
- b_compressed, e, a, b = make_rand_sparse_tensors(torch.int8, m, n, k)
50
- scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
51
- scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
52
- bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)
53
-
54
- out = ops.cutlass_scaled_sparse_mm(
55
- a, b_compressed, e, scale_a, scale_b, torch.bfloat16
56
- )
57
- out_ref = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)
58
-
59
- if not torch.allclose(out, out_ref):
60
- print("Incorrect results")
61
- print(out)
62
- print(out_ref)
63
- else:
64
- print("Correct results")
65
-
66
- timers = []
67
- # pytorch impl - bfloat16
68
- timers.append(
69
- bench_fn(
70
- label,
71
- sub_label,
72
- "pytorch_bf16_bf16_bf16_matmul-no-scales",
73
- torch.mm,
74
- a.to(dtype=torch.bfloat16),
75
- b.to(dtype=torch.bfloat16),
76
- )
77
- )
78
-
79
- # pytorch impl - float16
80
- timers.append(
81
- bench_fn(
82
- label,
83
- sub_label,
84
- "pytorch_fp16_fp16_fp16_matmul-no-scales",
85
- torch.mm,
86
- a.to(dtype=torch.float16),
87
- b.to(dtype=torch.float16),
88
- )
89
- )
90
-
91
- # cutlass impl
92
- timers.append(
93
- bench_fn(
94
- label,
95
- sub_label,
96
- "cutlass_i8_i8_bf16_scaled_mm",
97
- ops.cutlass_scaled_mm,
98
- a,
99
- b,
100
- scale_a,
101
- scale_b,
102
- torch.bfloat16,
103
- )
104
- )
105
-
106
- # cutlass with bias
107
- timers.append(
108
- bench_fn(
109
- label,
110
- sub_label,
111
- "cutlass_i8_i8_bf16_scaled_mm_bias",
112
- ops.cutlass_scaled_mm,
113
- a,
114
- b,
115
- scale_a,
116
- scale_b,
117
- torch.bfloat16,
118
- bias,
119
- )
120
- )
121
-
122
- # cutlass sparse impl
123
- timers.append(
124
- bench_fn(
125
- label,
126
- sub_label,
127
- "cutlass_i8_i8_bf16_scaled_sparse_mm",
128
- ops.cutlass_scaled_sparse_mm,
129
- a,
130
- b_compressed,
131
- e,
132
- scale_a,
133
- scale_b,
134
- torch.bfloat16,
135
- )
136
- )
137
-
138
- # cutlass sparse with bias
139
- timers.append(
140
- bench_fn(
141
- label,
142
- sub_label,
143
- "cutlass_i8_i8_bf16_scaled_sparse_mm_bias",
144
- ops.cutlass_scaled_sparse_mm,
145
- a,
146
- b_compressed,
147
- e,
148
- scale_a,
149
- scale_b,
150
- torch.bfloat16,
151
- bias,
152
- )
153
- )
154
-
155
- return timers
156
-
157
-
158
-def bench_fp8(
159
- dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str
160
-) -> Iterable[TMeasurement]:
161
- assert dtype == torch.float8_e4m3fn
162
- b_compressed, e, a, b = make_rand_sparse_tensors(torch.float8_e4m3fn, m, n, k)
163
- scale_a = torch.tensor(1.0, device="cuda", dtype=torch.float32)
164
- scale_b = torch.tensor(1.0, device="cuda", dtype=torch.float32)
165
- bias = torch.zeros((n,), device="cuda", dtype=torch.bfloat16)
166
-
167
- out = ops.cutlass_scaled_sparse_mm(
168
- a, b_compressed, e, scale_a, scale_b, torch.bfloat16
169
- )
170
- out_ref = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16)
171
-
172
- if not torch.allclose(out, out_ref):
173
- print("Incorrect results")
174
- print(out)
175
- print(out_ref)
176
- else:
177
- print("Correct results")
178
-
179
- timers = []
180
-
181
- # pytorch impl w. bf16
182
- timers.append(
183
- bench_fn(
184
- label,
185
- sub_label,
186
- "pytorch_bf16_bf16_bf16_matmul-no-scales",
187
- torch.mm,
188
- a.to(dtype=torch.bfloat16, device="cuda"),
189
- b.to(dtype=torch.bfloat16, device="cuda"),
190
- )
191
- )
192
-
193
- # pytorch impl: bf16 output, without fp8 fast accum
194
- timers.append(
195
- bench_fn(
196
- label,
197
- sub_label,
198
- "pytorch_fp8_fp8_bf16_scaled_mm",
199
- torch._scaled_mm,
200
- a,
201
- b,
202
- scale_a=scale_a,
203
- scale_b=scale_b,
204
- out_dtype=torch.bfloat16,
205
- )
206
- )
207
-
208
- # pytorch impl: bf16 output, with fp8 fast accum
209
- timers.append(
210
- bench_fn(
211
- label,
212
- sub_label,
213
- "pytorch_fp8_fp8_bf16_scaled_mm_fast_accum",
214
- torch._scaled_mm,
215
- a,
216
- b,
217
- scale_a=scale_a,
218
- scale_b=scale_b,
219
- out_dtype=torch.bfloat16,
220
- use_fast_accum=True,
221
- )
222
- )
223
-
224
- # pytorch impl: fp16 output, without fp8 fast accum
225
- timers.append(
226
- bench_fn(
227
- label,
228
- sub_label,
229
- "pytorch_fp8_fp8_fp16_scaled_mm",
230
- torch._scaled_mm,
231
- a,
232
- b,
233
- scale_a=scale_a,
234
- scale_b=scale_b,
235
- out_dtype=torch.float16,
236
- )
237
- )
238
-
239
- # pytorch impl: fp16 output, with fp8 fast accum
240
- timers.append(
241
- bench_fn(
242
- label,
243
- sub_label,
244
- "pytorch_fp8_fp8_fp16_scaled_mm_fast_accum",
245
- torch._scaled_mm,
246
- a,
247
- b,
248
- scale_a=scale_a,
249
- scale_b=scale_b,
250
- out_dtype=torch.float16,
251
- use_fast_accum=True,
252
- )
253
- )
254
-
255
- # cutlass impl: bf16 output
256
- timers.append(
257
- bench_fn(
258
- label,
259
- sub_label,
260
- "cutlass_fp8_fp8_bf16_scaled_mm",
261
- ops.cutlass_scaled_mm,
262
- a,
263
- b,
264
- scale_a,
265
- scale_b,
266
- torch.bfloat16,
267
- )
268
- )
269
-
270
- # cutlass impl: bf16 output
271
- timers.append(
272
- bench_fn(
273
- label,
274
- sub_label,
275
- "cutlass_fp8_fp8_bf16_scaled_sparse_mm",
276
- ops.cutlass_scaled_sparse_mm,
277
- a,
278
- b_compressed,
279
- e,
280
- scale_a,
281
- scale_b,
282
- torch.bfloat16,
283
- )
284
- )
285
-
286
- # cutlass impl: fp16 output
287
- timers.append(
288
- bench_fn(
289
- label,
290
- sub_label,
291
- "cutlass_fp8_fp8_fp16_scaled_sparse_mm",
292
- ops.cutlass_scaled_sparse_mm,
293
- a,
294
- b_compressed,
295
- e,
296
- scale_a,
297
- scale_b,
298
- torch.float16,
299
- )
300
- )
301
-
302
- # cutlass impl: bf16 output, with bias
303
- timers.append(
304
- bench_fn(
305
- label,
306
- sub_label,
307
- "cutlass_fp8_fp8_bf16_scaled_sparse_mm_bias",
308
- ops.cutlass_scaled_sparse_mm,
309
- a,
310
- b_compressed,
311
- e,
312
- scale_a,
313
- scale_b,
314
- torch.bfloat16,
315
- bias,
316
- )
317
- )
318
-
319
- # cutlass impl: fp16 output, with bias
320
- timers.append(
321
- bench_fn(
322
- label,
323
- sub_label,
324
- "cutlass_fp8_fp8_fp16_scaled_sparse_mm_bias",
325
- ops.cutlass_scaled_sparse_mm,
326
- a,
327
- b_compressed,
328
- e,
329
- scale_a,
330
- scale_b,
331
- torch.float16,
332
- bias.to(dtype=torch.float16),
333
- )
334
- )
335
-
336
- return timers
337
-
338
-
339
-def bench(
340
- dtype: torch.dtype, m: int, k: int, n: int, label: str, sub_label: str
341
-) -> Iterable[TMeasurement]:
342
- if dtype == torch.int8:
343
- return bench_int8(dtype, m, k, n, label, sub_label)
344
- if dtype == torch.float8_e4m3fn:
345
- return bench_fp8(dtype, m, k, n, label, sub_label)
346
- raise ValueError(
347
- f"Unsupported dtype {dtype}: should be one of torch.int8, torch.float8_e4m3fn."
348
- )
349
-
350
-
351
-# runner
352
-def print_timers(timers: Iterable[TMeasurement]):
353
- compare = TBenchmark.Compare(timers)
354
- compare.print()
355
-
356
-
357
-def run(
358
- dtype: torch.dtype, MKNs: Iterable[tuple[int, int, int]]
359
-) -> Iterable[TMeasurement]:
360
- results = []
361
- for m, k, n in MKNs:
362
- timers = bench(dtype, m, k, n, f"scaled-{dtype}-gemm", f"MKN=({m}x{k}x{n})")
363
- print_timers(timers)
364
- results.extend(timers)
365
-
366
- return results
367
-
368
-
369
-# output makers
370
-def make_output(
371
- data: Iterable[TMeasurement],
372
- MKNs: Iterable[tuple[int, int, int]],
373
- base_description: str,
374
- timestamp=None,
375
-):
376
- print(f"== All Results {base_description} ====")
377
- print_timers(data)
378
-
379
- # pickle all the results
380
- timestamp = int(time.time()) if timestamp is None else timestamp
381
- with open(f"{base_description}-{timestamp}.pkl", "wb") as f:
382
- pkl.dump(data, f)
383
-
384
-
385
-# argparse runners
386
-
387
-
388
-def run_square_bench(args):
389
- dim_sizes = list(range(args.dim_start, args.dim_end + 1, args.dim_increment))
390
- MKNs = list(zip(dim_sizes, dim_sizes, dim_sizes))
391
- data = run(args.dtype, MKNs)
392
-
393
- make_output(data, MKNs, f"square_bench-{args.dtype}")
394
-
395
-
396
-def run_range_bench(args):
397
- dim_sizes = list(range(args.dim_start, args.dim_end, args.dim_increment))
398
- n = len(dim_sizes)
399
- Ms = [args.m_constant] * n if args.m_constant is not None else dim_sizes
400
- Ks = [args.k_constant] * n if args.k_constant is not None else dim_sizes
401
- Ns = [args.n_constant] * n if args.n_constant is not None else dim_sizes
402
- MKNs = list(zip(Ms, Ks, Ns))
403
- data = run(args.dtype, MKNs)
404
-
405
- make_output(data, MKNs, f"range_bench-{args.dtype}")
406
-
407
-
408
-def run_model_bench(args):
409
- print("Benchmarking models:")
410
- for i, model in enumerate(args.models):
411
- print(f"[{i}] {model}")
412
-
413
- def model_shapes(model_name: str, tp_size: int) -> list[tuple[int, int]]:
414
- KNs = []
415
- for KN, tp_split_dim in copy.deepcopy(WEIGHT_SHAPES[model_name]):
416
- KN[tp_split_dim] = KN[tp_split_dim] // tp_size
417
- KNs.append(KN)
418
- return KNs
419
-
420
- model_bench_data = []
421
- models_tps = list(itertools.product(args.models, args.tp_sizes))
422
- for model, tp_size in models_tps:
423
- Ms = args.batch_sizes
424
- KNs = model_shapes(model, tp_size)
425
- MKNs = []
426
- for m in Ms:
427
- for k, n in KNs:
428
- MKNs.append((m, k, n))
429
-
430
- data = run(args.dtype, MKNs)
431
- model_bench_data.append(data)
432
-
433
- # Print all results
434
- for data, model_tp in zip(model_bench_data, models_tps):
435
- model, tp_size = model_tp
436
- print(f"== Results {args.dtype} {model}-TP{tp_size} ====")
437
- print_timers(data)
438
-
439
- timestamp = int(time.time())
440
-
441
- all_data = []
442
- for d in model_bench_data:
443
- all_data.extend(d)
444
- # pickle all data
445
- with open(f"model_bench-{args.dtype}-{timestamp}.pkl", "wb") as f:
446
- pkl.dump(all_data, f)
447
-
448
-
449
-if __name__ == "__main__":
450
-
451
- def to_torch_dtype(dt):
452
- if dt == "int8":
453
- return torch.int8
454
- if dt == "fp8":
455
- return torch.float8_e4m3fn
456
- raise ValueError("unsupported dtype")
457
-
458
- parser = FlexibleArgumentParser(
459
- description="""
460
-Benchmark Cutlass GEMM.
461
-
462
- To run square GEMMs:
463
- python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 square_bench --dim-start 128 --dim-end 512 --dim-increment 64
464
-
465
- To run constant N and K and sweep M:
466
- python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 range_bench --dim-start 128 --dim-end 512 --dim-increment 64 --n-constant 16384 --k-constant 16384
467
-
468
- To run dimensions from a model:
469
- python3 ./benchmarks/cutlass_benchmarks/sparse_benchmarks.py --dtype fp8 model_bench --models meta-llama/Llama-2-7b-hf --batch-sizes 16 --tp-sizes 1
470
-
471
- Output:
472
- - a .pkl file, that is a list of raw torch.benchmark.utils.Measurements for the pytorch and cutlass implementations for the various GEMMs.
473
- """, # noqa: E501
474
- formatter_class=argparse.RawTextHelpFormatter,
475
- )
476
-
477
- parser.add_argument(
478
- "--dtype",
479
- type=to_torch_dtype,
480
- required=True,
481
- help="Available options are ['int8', 'fp8']",
482
- )
483
- subparsers = parser.add_subparsers(dest="cmd")
484
-
485
- square_parser = subparsers.add_parser("square_bench")
486
- square_parser.add_argument("--dim-start", type=int, required=True)
487
- square_parser.add_argument("--dim-end", type=int, required=True)
488
- square_parser.add_argument("--dim-increment", type=int, required=True)
489
- square_parser.set_defaults(func=run_square_bench)
490
-
491
- range_parser = subparsers.add_parser("range_bench")
492
- range_parser.add_argument("--dim-start", type=int, required=True)
493
- range_parser.add_argument("--dim-end", type=int, required=True)
494
- range_parser.add_argument("--dim-increment", type=int, required=True)
495
- range_parser.add_argument("--m-constant", type=int, default=None)
496
- range_parser.add_argument("--n-constant", type=int, default=None)
497
- range_parser.add_argument("--k-constant", type=int, default=None)
498
- range_parser.set_defaults(func=run_range_bench)
499
-
500
- model_parser = subparsers.add_parser("model_bench")
501
- model_parser.add_argument(
502
- "--models",
503
- nargs="+",
504
- type=str,
505
- default=DEFAULT_MODELS,
506
- choices=WEIGHT_SHAPES.keys(),
507
- )
508
- model_parser.add_argument(
509
- "--tp-sizes", nargs="+", type=int, default=DEFAULT_TP_SIZES
510
- )
511
- model_parser.add_argument(
512
- "--batch-sizes", nargs="+", type=int, default=DEFAULT_BATCH_SIZES
513
- )
514
- model_parser.set_defaults(func=run_model_bench)
515
-
516
- args = parser.parse_args()
517
- args.func(args)
benchmarks/cutlass_benchmarks/utils.py
-48
@@ -5,8 +5,6 @@
5
6
import torch
7
8
-import vllm._custom_ops as ops
9
-
8
9
def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
10
finfo = torch.finfo(torch.float8_e4m3fn)
@@ -39,49 +37,3 @@ def make_rand_tensors(
37
return to_fp8(a), to_fp8(b)
38
39
raise ValueError("unsupported dtype")
42
-
43
-
44
-def prune_to_2_4(tensor):
45
- # Reshape tensor to [N, 4] where N is number of groups of 4
46
- original_shape = tensor.shape
47
- reshaped = tensor.reshape(-1, 4)
48
-
49
- # Get indices of top 2 absolute values in each group of 4
50
- _, indices = torch.topk(torch.abs(reshaped), k=2, dim=1)
51
-
52
- # Create binary mask
53
- mask = torch.zeros_like(reshaped)
54
- mask.scatter_(dim=1, index=indices, src=torch.ones_like(indices, dtype=mask.dtype))
55
-
56
- # Apply mask and reshape back
57
- pruned = reshaped * mask
58
-
59
- # Turn all -0.0 to 0.0
60
- pruned[pruned == -0.0] = 0.0
61
-
62
- return pruned.reshape(original_shape)
63
-
64
-
65
-def make_rand_sparse_tensors(
66
- dtype: torch.dtype, m: int, n: int, k: int
67
-) -> tuple[torch.Tensor, torch.Tensor]:
68
- a = torch.randn((m, k), device="cuda") * 5
69
- b = torch.randn((n, k), device="cuda").t() * 5
70
-
71
- b = prune_to_2_4(b.t()).t()
72
-
73
- if dtype == torch.int8:
74
- a, b = to_int8(a), to_int8(b)
75
- elif dtype == torch.float8_e4m3fn:
76
- a, b = to_fp8(a), to_fp8(b)
77
- elif dtype == torch.float16:
78
- a, b = to_fp16(a), to_fp16(b)
79
- elif dtype == torch.bfloat16:
80
- a, b = to_bf16(a), to_bf16(b)
81
- else:
82
- raise ValueError("unsupported dtype")
83
-
84
- b_compressed, e = ops.cutlass_sparse_compress(b.t())
85
-
86
- # Compressed B, Metadata, Original A, B
87
- return b_compressed, e, a, b
csrc/ops.h
-10
@@ -285,16 +285,6 @@ void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a,
285
std::optional<torch::Tensor> const& azp,
286
std::optional<torch::Tensor> const& bias);
287
288
-bool cutlass_sparse_scaled_mm_supported(int64_t cuda_device_capability);
289
-
290
-void cutlass_scaled_sparse_mm(torch::Tensor& out, torch::Tensor const& a,
291
- torch::Tensor const& b, torch::Tensor const& e,
292
- torch::Tensor const& a_scales,
293
- torch::Tensor const& b_scales,
294
- std::optional<torch::Tensor> const& bias);
295
-
296
-std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
297
-
288
std::tuple<torch::Tensor, torch::Tensor> scaled_fp4_quant_func(
289
torch::Tensor const& input, torch::Tensor const& input_scale,
290
bool is_sf_swizzled_layout);
csrc/sparse/cutlass/sparse_compressor_c3x.cuh
deleted
-90
@@ -1,90 +0,0 @@
1
-#pragma once
2
-
3
-// clang-format will break include orders
4
-// clang-format off
5
-#include <cudaTypedefs.h>
6
-
7
-#if defined CUDA_VERSION && CUDA_VERSION >= 12020
8
-#include "sparse_scaled_mm_c3x.cuh"
9
-
10
-#include "cutlass/numeric_conversion.h"
11
-#include "cutlass/transform/device/transform_universal_adapter.hpp"
12
-#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp"
13
-#include "cutlass/epilogue/collective/default_epilogue.hpp"
14
-
15
-// clang-format on
16
-
17
-using namespace cute;
18
-using namespace vllm;
19
-
20
-using CompressorResult = std::tuple<torch::Tensor, torch::Tensor>;
21
-/// Make A structured sparse by replacing elements with 0 and compress it
22
-template <typename Gemm>
23
-CompressorResult cutlass_sparse_compress(torch::Tensor const& a) {
24
- // Checks for conformality
25
- TORCH_CHECK(a.dtype() == torch::kInt8 || a.dtype() == torch::kFloat8_e4m3fn ||
26
- a.dtype() == torch::kFloat16 || a.dtype() == torch::kBFloat16);
27
- TORCH_CHECK(a.dim() == 2)
28
- // Check for strides and alignment
29
- TORCH_CHECK(a.stride(0) % 4 == 0) // Required for semi-structured sparsity
30
- TORCH_CHECK(a.stride(1) == 1)
31
-
32
- using GemmKernel = typename Gemm::KernelType;
33
- using ElementA = typename Gemm::ElementAB;
34
- using ElementE = typename GemmKernel::CollectiveMainloop::ElementE;
35
-
36
- int m = a.size(0);
37
- int k = a.size(1);
38
- using ProblemShape = typename GemmKernel::ProblemShape;
39
- ProblemShape prob_shape{m, 1, k, 1};
40
-
41
- int64_t lda = a.stride(0);
42
- using StrideA = Stride<int64_t, Int<1>, int64_t>;
43
- StrideA a_stride{lda, Int<1>{}, 0};
44
-
45
- using CompressorUtility = typename Gemm::CompressorUtility;
46
- CompressorUtility compressor_utility(prob_shape, a_stride);
47
-
48
- // Allocate buffers for the metadata E and the compressed matrix A
49
- int ME = compressor_utility.get_metadata_m_physical();
50
- int KE = compressor_utility.get_metadata_k_physical();
51
- int MC = compressor_utility.get_tensorA_m_physical();
52
- int KC = compressor_utility.get_tensorA_k_physical();
53
-
54
- auto const a_meta_options =
55
- torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
56
- auto const a_nzs_options =
57
- torch::TensorOptions().dtype(a.dtype()).device(a.device());
58
-
59
- auto a_meta = torch::zeros({ME, KE}, a_meta_options);
60
- auto a_nzs = torch::zeros({MC, KC}, a_nzs_options);
61
-
62
- auto a_ptr = static_cast<ElementA*>(a.data_ptr());
63
- auto a_nzs_ptr = static_cast<ElementA*>(a_nzs.data_ptr());
64
- auto a_meta_ptr = static_cast<ElementE*>(a_meta.data_ptr());
65
-
66
- cutlass::KernelHardwareInfo hw_info;
67
- hw_info.device_id = a.device().index();
68
- hw_info.sm_count =
69
- cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
70
- hw_info.device_id);
71
-
72
- using Compressor = typename Gemm::Compressor;
73
- typename Compressor::Arguments arguments{
74
- prob_shape, {a_ptr, a_stride, a_nzs_ptr, a_meta_ptr}, {hw_info}};
75
-
76
- Compressor compressor_op;
77
- size_t workspace_size = Compressor::get_workspace_size(arguments);
78
- auto const workspace_options =
79
- torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
80
- auto workspace = torch::empty(workspace_size, workspace_options);
81
-
82
- CUTLASS_CHECK(compressor_op.can_implement(arguments));
83
- CUTLASS_CHECK(compressor_op.initialize(arguments, workspace.data_ptr()));
84
- CUTLASS_CHECK(compressor_op.run());
85
- CUDA_CHECK(cudaDeviceSynchronize());
86
-
87
- return {a_meta, a_nzs};
88
-}
89
-
90
-#endif
csrc/sparse/cutlass/sparse_scaled_mm_c3x.cu
deleted
-307
@@ -1,307 +0,0 @@
1
-// clang-format will break include orders
2
-// clang-format off
3
-#include <cudaTypedefs.h>
4
-
5
-#if defined CUDA_VERSION && CUDA_VERSION >= 12020
6
-#include "sparse_scaled_mm_c3x.cuh"
7
-// clang-format on
8
-
9
-using namespace cute;
10
-using namespace vllm;
11
-
12
-struct GemmCallerTraits {
13
- using return_type = void;
14
-
15
- template <typename GemmConfig, typename... Args>
16
- static return_type invoke(Args&&... args) {
17
- return cutlass_sparse_gemm_caller<GemmConfig>(std::forward<Args>(args)...);
18
- }
19
-};
20
-
21
-struct GemmCompressorTraits {
22
- using return_type = CompressorResult;
23
-
24
- template <typename GemmConfig, typename... Args>
25
- static return_type invoke(Args&&... args) {
26
- return cutlass_sparse_compress<GemmConfig>(std::forward<Args>(args)...);
27
- }
28
-};
29
-
30
-template <typename InType, typename OutType,
31
- template <typename, typename, typename> typename Epilogue,
32
- typename DispatchFunc, typename... Args>
33
-typename DispatchFunc::return_type cutlass_gemm_sm90_fp8_dispatch(
34
- uint32_t m, uint32_t n, Args&&... args) {
35
- static_assert(std::is_same_v<InType, cutlass::float_e4m3_t>);
36
-
37
- using Cutlass3xGemmDefault =
38
- typename sm90_config_default<InType, OutType, Epilogue>::Cutlass3xGemm;
39
- using Cutlass3xGemmM64 =
40
- typename sm90_fp8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
41
- using Cutlass3xGemmM128 =
42
- typename sm90_fp8_config_M128<InType, OutType, Epilogue>::Cutlass3xGemm;
43
- using Cutlass3xGemmM256 =
44
- typename sm90_fp8_config_M256<InType, OutType, Epilogue>::Cutlass3xGemm;
45
- using Cutlass3xGemmM512 =
46
- typename sm90_fp8_config_M512<InType, OutType, Epilogue>::Cutlass3xGemm;
47
-
48
- using Cutlass3xGemm1 =
49
- typename sm90_fp8_config_1<InType, OutType, Epilogue>::Cutlass3xGemm;
50
- using Cutlass3xGemm2 =
51
- typename sm90_fp8_config_2<InType, OutType, Epilogue>::Cutlass3xGemm;
52
- using Cutlass3xGemm3 =
53
- typename sm90_fp8_config_3<InType, OutType, Epilogue>::Cutlass3xGemm;
54
- using Cutlass3xGemm4 =
55
- typename sm90_fp8_config_4<InType, OutType, Epilogue>::Cutlass3xGemm;
56
- using Cutlass3xGemm5 =
57
- typename sm90_fp8_config_5<InType, OutType, Epilogue>::Cutlass3xGemm;
58
- using Cutlass3xGemm6 =
59
- typename sm90_fp8_config_6<InType, OutType, Epilogue>::Cutlass3xGemm;
60
- using Cutlass3xGemm7 =
61
- typename sm90_fp8_config_7<InType, OutType, Epilogue>::Cutlass3xGemm;
62
- using Cutlass3xGemm8 =
63
- typename sm90_fp8_config_8<InType, OutType, Epilogue>::Cutlass3xGemm;
64
-
65
- uint32_t const mp2 =
66
- std::max(static_cast<uint32_t>(64), next_pow_2(m)); // next power of 2
67
-
68
- if (mp2 <= 64) {
69
- if (n == 28672) {
70
- return DispatchFunc::template invoke<Cutlass3xGemm2>(
71
- std::forward<Args>(args)...);
72
- } else if (n == 4096 || n == 6144) {
73
- return DispatchFunc::template invoke<Cutlass3xGemm1>(
74
- std::forward<Args>(args)...);
75
- }
76
- } else if (mp2 <= 128) {
77
- if (n == 4096) {
78
- return DispatchFunc::template invoke<Cutlass3xGemm3>(
79
- std::forward<Args>(args)...);
80
- } else if (n == 28672) {
81
- return DispatchFunc::template invoke<Cutlass3xGemm5>(
82
- std::forward<Args>(args)...);
83
- } else if (n == 6144) {
84
- return DispatchFunc::template invoke<Cutlass3xGemm4>(
85
- std::forward<Args>(args)...);
86
- }
87
- } else if (mp2 <= 256) {
88
- if (n == 4096) {
89
- return DispatchFunc::template invoke<Cutlass3xGemm6>(
90
- std::forward<Args>(args)...);
91
- } else if (n == 28672) {
92
- return DispatchFunc::template invoke<Cutlass3xGemm8>(
93
- std::forward<Args>(args)...);
94
- } else if (n == 6144) {
95
- return DispatchFunc::template invoke<Cutlass3xGemm7>(
96
- std::forward<Args>(args)...);
97
- }
98
- } else {
99
- if (n == 6144 || n == 28672) {
100
- return DispatchFunc::template invoke<Cutlass3xGemm8>(
101
- std::forward<Args>(args)...);
102
- } else if (n == 4096) {
103
- return DispatchFunc::template invoke<Cutlass3xGemm7>(
104
- std::forward<Args>(args)...);
105
- }
106
- }
107
-
108
- // Otherwise the default heuristic
109
- if (mp2 <= 64) {
110
- // n in [1, 64]
111
- return DispatchFunc::template invoke<Cutlass3xGemmM64>(
112
- std::forward<Args>(args)...);
113
- } else if (mp2 <= 128) {
114
- // n in (64, 128]
115
- return DispatchFunc::template invoke<Cutlass3xGemmM128>(
116
- std::forward<Args>(args)...);
117
- } else if (mp2 <= 256) {
118
- // n in (128, 256]
119
- return DispatchFunc::template invoke<Cutlass3xGemmM256>(
120
- std::forward<Args>(args)...);
121
- } else {
122
- // n in (256, inf)
123
- return DispatchFunc::template invoke<Cutlass3xGemmM512>(
124
- std::forward<Args>(args)...);
125
- }
126
-}
127
-
128
-template <typename InType, typename OutType,
129
- template <typename, typename, typename> typename Epilogue,
130
- typename DispatchFunc, typename... Args>
131
-typename DispatchFunc::return_type cutlass_gemm_sm90_16bit_dispatch(
132
- uint32_t m, uint32_t n, Args&&... args) {
133
- using Cutlass3xGemmDefault =
134
- typename sm90_config_default<InType, OutType, Epilogue>::Cutlass3xGemm;
135
-
136
- return DispatchFunc::template invoke<Cutlass3xGemmDefault>(
137
- std::forward<Args>(args)...);
138
-}
139
-
140
-template <typename InType, typename OutType,
141
- template <typename, typename, typename> typename Epilogue,
142
- typename DispatchFunc, typename... Args>
143
-typename DispatchFunc::return_type cutlass_gemm_sm90_int8_dispatch(
144
- uint32_t m, uint32_t n, Args&&... args) {
145
- static_assert(std::is_same_v<InType, int8_t>);
146
-
147
- using Cutlass3xGemmDefault =
148
- typename sm90_config_default<InType, OutType, Epilogue>::Cutlass3xGemm;
149
- using Cutlass3xGemmM128 =
150
- typename sm90_int8_config_M128<InType, OutType, Epilogue>::Cutlass3xGemm;
151
- using Cutlass3xGemmM64 =
152
- typename sm90_int8_config_M64<InType, OutType, Epilogue>::Cutlass3xGemm;
153
- using Cutlass3xGemmM32NBig =
154
- typename sm90_int8_config_M32_NBig<InType, OutType,
155
- Epilogue>::Cutlass3xGemm;
156
- using Cutlass3xGemmM32NSmall =
157
- typename sm90_int8_config_M32_NSmall<InType, OutType,
158
- Epilogue>::Cutlass3xGemm;
159
-
160
- bool const is_small_n = n < 8192;
161
- uint32_t const mp2 =
162
- std::max(static_cast<uint32_t>(32), next_pow_2(m)); // next power of 2
163
-
164
- if (mp2 <= 32) {
165
- // m in [1, 32]
166
- if (is_small_n) {
167
- return DispatchFunc::template invoke<Cutlass3xGemmM32NSmall>(
168
- std::forward<Args>(args)...);
169
- } else {
170
- return DispatchFunc::template invoke<Cutlass3xGemmM32NBig>(
171
- std::forward<Args>(args)...);
172
- }
173
- } else if (mp2 <= 64) {
174
- // m in (32, 64]
175
- return DispatchFunc::template invoke<Cutlass3xGemmM64>(
176
- std::forward<Args>(args)...);
177
- } else if (mp2 <= 128) {
178
- // m in (64, 128]
179
- return DispatchFunc::template invoke<Cutlass3xGemmM128>(
180
- std::forward<Args>(args)...);
181
- } else {
182
- // m in (128, inf)
183
- return DispatchFunc::template invoke<Cutlass3xGemmDefault>(
184
- std::forward<Args>(args)...);
185
- }
186
-}
187
-
188
-// Dispatch to GEMM implementations based on element types
189
-template <template <typename, typename, typename> typename Epilogue,
190
- typename... EpilogueArgs>
191
-void cutlass_scaled_sparse_mm_sm90_epilogue(torch::Tensor& out,
192
- torch::Tensor const& a,
193
- torch::Tensor const& bt_nzs,
194
- torch::Tensor const& bt_meta,
195
- EpilogueArgs&&... epilogue_args) {
196
- uint32_t const m = out.size(0);
197
- uint32_t const n = out.size(1);
198
-
199
- // TODO: add dispatch functions to all of these
200
- TORCH_CHECK(bt_meta.dtype() == torch::kUInt8);
201
- if (a.dtype() == torch::kInt8) {
202
- TORCH_CHECK(bt_nzs.dtype() == torch::kInt8);
203
-
204
- if (out.dtype() == torch::kBFloat16) {
205
- return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::bfloat16_t,
206
- Epilogue, GemmCallerTraits>(
207
- m, n, out, a, bt_nzs, bt_meta,
208
- std::forward<EpilogueArgs>(epilogue_args)...);
209
- } else {
210
- TORCH_CHECK(out.dtype() == torch::kFloat16);
211
- return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::half_t, Epilogue,
212
- GemmCallerTraits>(
213
- m, n, out, a, bt_nzs, bt_meta,
214
- std::forward<EpilogueArgs>(epilogue_args)...);
215
- }
216
- } else if (a.dtype() == torch::kFloat8_e4m3fn) {
217
- TORCH_CHECK(bt_nzs.dtype() == torch::kFloat8_e4m3fn);
218
-
219
- if (out.dtype() == torch::kBFloat16) {
220
- return cutlass_gemm_sm90_fp8_dispatch<cutlass::float_e4m3_t,
221
- cutlass::bfloat16_t, Epilogue,
222
- GemmCallerTraits>(
223
- m, n, out, a, bt_nzs, bt_meta,
224
- std::forward<EpilogueArgs>(epilogue_args)...);
225
- } else {
226
- TORCH_CHECK(out.dtype() == torch::kFloat16);
227
- return cutlass_gemm_sm90_fp8_dispatch<
228
- cutlass::float_e4m3_t, cutlass::half_t, Epilogue, GemmCallerTraits>(
229
- m, n, out, a, bt_nzs, bt_meta,
230
- std::forward<EpilogueArgs>(epilogue_args)...);
231
- }
232
- } else if (a.dtype() == torch::kFloat16) {
233
- TORCH_CHECK(bt_nzs.dtype() == torch::kFloat16);
234
- TORCH_CHECK(out.dtype() == torch::kFloat16);
235
-
236
- return cutlass_gemm_sm90_16bit_dispatch<cutlass::half_t, cutlass::half_t,
237
- Epilogue, GemmCallerTraits>(
238
- m, n, out, a, bt_nzs, bt_meta,
239
- std::forward<EpilogueArgs>(epilogue_args)...);
240
- } else { // a.dtype() == torch::kBFloat16
241
- TORCH_CHECK(a.dtype() == torch::kBFloat16);
242
- TORCH_CHECK(bt_nzs.dtype() == torch::kBFloat16);
243
- TORCH_CHECK(out.dtype() == torch::kBFloat16);
244
-
245
- return cutlass_gemm_sm90_16bit_dispatch<
246
- cutlass::bfloat16_t, cutlass::bfloat16_t, Epilogue, GemmCallerTraits>(
247
- m, n, out, a, bt_nzs, bt_meta,
248
- std::forward<EpilogueArgs>(epilogue_args)...);
249
- }
250
-}
251
-
252
-void cutlass_scaled_sparse_mm_sm90(torch::Tensor& out, torch::Tensor const& a,
253
- torch::Tensor const& bt_nzs,
254
- torch::Tensor const& bt_meta,
255
- torch::Tensor const& a_scales,
256
- torch::Tensor const& b_scales,
257
- std::optional<torch::Tensor> const& bias) {
258
- TORCH_CHECK(bt_meta.dtype() == torch::kUInt8);
259
- TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
260
- TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
261
-
262
- if (bias) {
263
- TORCH_CHECK(bias->dtype() == out.dtype(),
264
- "CUTLASS scaled_mm bias dtype must match output dtype ",
265
- out.dtype());
266
- return cutlass_scaled_sparse_mm_sm90_epilogue<
267
- c3x::ScaledEpilogueColumnBias>(out, a, bt_nzs, bt_meta, b_scales,
268
- a_scales, *bias);
269
- } else {
270
- return cutlass_scaled_sparse_mm_sm90_epilogue<c3x::ScaledEpilogue>(
271
- out, a, bt_nzs, bt_meta, b_scales, a_scales);
272
- }
273
-}
274
-
275
-CompressorResult cutlass_sparse_compress_sm90(torch::Tensor const& a) {
276
- // These m and n variables are fordispatching to different GEMM algorithms.
277
- uint32_t const m = 1; // Set M to 1 for compression
278
- uint32_t const n = a.size(1);
279
-
280
- // Note: For correctness, the compressed format must be invariant in:
281
- // - M, the flattened number of tokens
282
- // - Whether output dtype is fp16 or bf16
283
- // - CUTLASS epilogues
284
-
285
- if (a.dtype() == torch::kInt8) {
286
- return cutlass_gemm_sm90_int8_dispatch<int8_t, cutlass::bfloat16_t,
287
- c3x::TrivialEpilogue,
288
- GemmCompressorTraits>(m, n, a);
289
- } else if (a.dtype() == torch::kFloat8_e4m3fn) {
290
- return cutlass_gemm_sm90_fp8_dispatch<
291
- cutlass::float_e4m3_t, cutlass::bfloat16_t, c3x::TrivialEpilogue,
292
- GemmCompressorTraits>(m, n, a);
293
- } else if (a.dtype() == torch::kFloat16) {
294
- return cutlass_gemm_sm90_16bit_dispatch<
295
- cutlass::bfloat16_t, cutlass::bfloat16_t, c3x::TrivialEpilogue,
296
- GemmCompressorTraits>(m, n, a);
297
- } else {
298
- TORCH_CHECK(a.dtype() == torch::kBFloat16,
299
- "cutlass_sparse_compress only supports int8, fp8_e4m3, fp16, "
300
- "and bf16 datatypes");
301
- return cutlass_gemm_sm90_16bit_dispatch<cutlass::half_t, cutlass::half_t,
302
- c3x::TrivialEpilogue,
303
- GemmCompressorTraits>(m, n, a);
304
- }
305
-}
306
-
307
-#endif
csrc/sparse/cutlass/sparse_scaled_mm_c3x.cuh
deleted
-570
@@ -1,570 +0,0 @@
1
-#pragma once
2
-
3
-// clang-format will break include orders
4
-// clang-format off
5
-#include <cudaTypedefs.h>
6
-
7
-#include <torch/all.h>
8
-
9
-#include <ATen/cuda/CUDAContext.h>
10
-
11
-#include "cuda_utils.h"
12
-
13
-#include "cutlass/cutlass.h"
14
-
15
-#include "cutlass/gemm/device/gemm_universal_adapter.h"
16
-#include "cutlass/epilogue/collective/collective_builder.hpp"
17
-#include "cutlass/gemm/collective/collective_builder.hpp"
18
-
19
-#include "cutlass/transform/device/transform_universal_adapter.hpp"
20
-#include "cutlass/transform/kernel/sparse_gemm_compressor.hpp"
21
-
22
-#include "core/math.hpp"
23
-#include "cutlass_extensions/cute_utils.cuh"
24
-#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
25
-#include "cutlass_extensions/common.hpp"
26
-#include "cutlass_extensions/torch_utils.hpp"
27
-// clang-format on
28
-
29
-using namespace cute;
30
-
31
-/*
32
- This file defines 2:4 sparse GEMM operations using the CUTLASS 3.x API,
33
- for NVIDIA GPUs with sm90a (Hopper) or later.
34
-*/
35
-
36
-namespace {
37
-
38
-// A wrapper for the GEMM kernel that is used to guard against compilation on
39
-// architectures that will never use the kernel. The purpose of this is to
40
-// reduce the size of the compiled binary.
41
-// __CUDA_ARCH__ is not defined in host code, so this lets us smuggle the ifdef
42
-// into code that will be executed on the device where it is defined.
43
-template <typename Kernel>
44
-struct enable_sm90_or_later : Kernel {
45
- template <typename... Args>
46
- CUTLASS_DEVICE void operator()(Args&&... args) {
47
-#if defined __CUDA_ARCH__ && __CUDA_ARCH__ >= 900
48
- Kernel::operator()(std::forward<Args>(args)...);
49
-#endif
50
- }
51
-};
52
-
53
-using GemmUniversalMode = cutlass::gemm::GemmUniversalMode;
54
-
55
-/*
56
- * cutlass_sparse_3x_gemm defines a 2:4 sparse GEMM kernel via CUTLASS
57
- * for SM90 Hopper systems.
58
- */
59
-template <typename ElementAB_, typename ElementD_,
60
- template <typename, typename, typename> typename Epilogue_,
61
- typename TileShape, typename ClusterShape, typename KernelSchedule,
62
- typename EpilogueSchedule>
63
-struct cutlass_sparse_3x_gemm {
64
- using ElementAB = ElementAB_;
65
- using ElementD = ElementD_;
66
- using ElementAcc =
67
- typename std::conditional<std::is_same_v<ElementAB, int8_t>, int32_t,
68
- float>::type;
69
-
70
- using Epilogue = Epilogue_<ElementAcc, ElementD, TileShape>;
71
-
72
- using ElementC = void;
73
- using LayoutC = cutlass::layout::RowMajor;
74
- using LayoutC_Transpose =
75
- typename cutlass::layout::LayoutTranspose<LayoutC>::type;
76
-
77
- using EVTCompute = typename Epilogue::EVTCompute;
78
-
79
- // These are the minimum alignments needed for the kernels to compile
80
- static constexpr int AlignmentAB =
81
- 128 / cutlass::sizeof_bits<ElementAB>::value;
82
- static constexpr int AlignmentCD =
83
- 128 / cutlass::sizeof_bits<ElementD>::value;
84
-
85
- using CollectiveEpilogue =
86
- typename cutlass::epilogue::collective::CollectiveBuilder<
87
- cutlass::arch::Sm90, cutlass::arch::OpClassTensorOp, TileShape,
88
- ClusterShape, cutlass::epilogue::collective::EpilogueTileAuto,
89
- ElementAcc, float, ElementC, LayoutC_Transpose, AlignmentCD, ElementD,
90
- LayoutC_Transpose, AlignmentCD, EpilogueSchedule,
91
- EVTCompute>::CollectiveOp;
92
-
93
- static constexpr size_t CEStorageSize =
94
- sizeof(typename CollectiveEpilogue::SharedStorage);
95
- using Stages = typename cutlass::gemm::collective::StageCountAutoCarveout<
96
- static_cast<int>(CEStorageSize)>;
97
-
98
- // clang-format off
99
- using CollectiveMainloop =
100
- typename cutlass::gemm::collective::CollectiveBuilder<
101
- cutlass::arch::Sm90, cutlass::arch::OpClassSparseTensorOp,
102
- ElementAB, cutlass::layout::RowMajor, AlignmentAB,
103
- ElementAB, cutlass::layout::ColumnMajor, AlignmentAB,
104
- ElementAcc, TileShape, ClusterShape,
105
- Stages,
106
- KernelSchedule>::CollectiveOp;
107
- // clang-format on
108
-
109
- using KernelType = enable_sm90_or_later<cutlass::gemm::kernel::GemmUniversal<
110
- cute::Shape<int, int, int, int>, CollectiveMainloop, CollectiveEpilogue,
111
- cutlass::gemm::PersistentScheduler>>;
112
-
113
- struct GemmKernel : public KernelType {};
114
-
115
- // Sparse compressor definitions
116
- using SparseConfig = typename GemmKernel::CollectiveMainloop::SparseConfig;
117
- using LayoutTagA = cutlass::layout::RowMajor;
118
- using CompressorUtility =
119
- cutlass::transform::kernel::StructuredSparseCompressorUtility<
120
- typename GemmKernel::ProblemShape, ElementAB, LayoutTagA,
121
- SparseConfig>;
122
- using CompressorKernel =
123
- cutlass::transform::kernel::StructuredSparseCompressor<
124
- typename GemmKernel::ProblemShape, ElementAB, LayoutTagA,
125
- SparseConfig, cutlass::arch::Sm90>;
126
- using Compressor =
127
- cutlass::transform::device::TransformUniversalAdapter<CompressorKernel>;
128
-};
129
-
130
-/*
131
- * This class defines kernel to compress a 2:4 sparse matrix.
132
- * The particular format is defined by the Gemm template parameter,
133
- * which is a cutlass_sparse_3x_gemm.
134
- */
135
-using CompressorResult = std::tuple<torch::Tensor, torch::Tensor>;
136
-/// Make A structured sparse by replacing elements with 0 and compress it
137
-template <typename Gemm>
138
-CompressorResult cutlass_sparse_compress(torch::Tensor const& a) {
139
- // Checks for conformality
140
- TORCH_CHECK(a.dtype() == torch::kInt8 || a.dtype() == torch::kFloat8_e4m3fn ||
141
- a.dtype() == torch::kFloat16 || a.dtype() == torch::kBFloat16);
142
- TORCH_CHECK(a.dim() == 2)
143
- // Check for strides and alignment
144
- TORCH_CHECK(a.stride(0) % 4 == 0) // Required for semi-structured sparsity
145
- TORCH_CHECK(a.stride(1) == 1)
146
-
147
- using GemmKernel = typename Gemm::KernelType;
148
- using ElementA = typename Gemm::ElementAB;
149
- using ElementE = typename GemmKernel::CollectiveMainloop::ElementE;
150
-
151
- int m = a.size(0);
152
- int k = a.size(1);
153
- using ProblemShape = typename GemmKernel::ProblemShape;
154
- ProblemShape prob_shape{m, 1, k, 1};
155
-
156
- int64_t lda = a.stride(0);
157
- using StrideA = Stride<int64_t, Int<1>, int64_t>;
158
- StrideA a_stride{lda, Int<1>{}, 0};
159
-
160
- using CompressorUtility = typename Gemm::CompressorUtility;
161
- CompressorUtility compressor_utility(prob_shape, a_stride);
162
-
163
- // Allocate buffers for the metadata E and the compressed matrix A
164
- int ME = compressor_utility.get_metadata_m_physical();
165
- int KE = compressor_utility.get_metadata_k_physical();
166
- int MC = compressor_utility.get_tensorA_m_physical();
167
- int KC = compressor_utility.get_tensorA_k_physical();
168
-
169
- auto const a_meta_options =
170
- torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
171
- auto const a_nzs_options =
172
- torch::TensorOptions().dtype(a.dtype()).device(a.device());
173
-
174
- auto a_meta = torch::zeros({ME, KE}, a_meta_options);
175
- auto a_nzs = torch::zeros({MC, KC}, a_nzs_options);
176
-
177
- auto a_ptr = static_cast<ElementA*>(a.data_ptr());
178
- auto a_nzs_ptr = static_cast<ElementA*>(a_nzs.data_ptr());
179
- auto a_meta_ptr = static_cast<ElementE*>(a_meta.data_ptr());
180
-
181
- cutlass::KernelHardwareInfo hw_info;
182
- hw_info.device_id = a.device().index();
183
- hw_info.sm_count =
184
- cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
185
- hw_info.device_id);
186
-
187
- using Compressor = typename Gemm::Compressor;
188
- typename Compressor::Arguments arguments{
189
- prob_shape, {a_ptr, a_stride, a_nzs_ptr, a_meta_ptr}, {hw_info}};
190
-
191
- Compressor compressor_op;
192
- size_t workspace_size = Compressor::get_workspace_size(arguments);
193
- auto const workspace_options =
194
- torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
195
- auto workspace = torch::empty(workspace_size, workspace_options);
196
-
197
- CUTLASS_CHECK(compressor_op.can_implement(arguments));
198
- CUTLASS_CHECK(compressor_op.initialize(arguments, workspace.data_ptr()));
199
- CUTLASS_CHECK(compressor_op.run());
200
- CUDA_CHECK(cudaDeviceSynchronize());
201
-
202
- return {a_meta, a_nzs};
203
-}
204
-
205
-template <typename Gemm, typename... EpilogueArgs>
206
-void cutlass_sparse_gemm_caller(torch::Tensor& out, torch::Tensor const& a,
207
- torch::Tensor const& bt_nzs,
208
- torch::Tensor const& bt_meta,
209
- EpilogueArgs&&... epilogue_params) {
210
- using ElementAB = typename Gemm::ElementAB;
211
- using ElementD = typename Gemm::ElementD;
212
-
213
- // Interface stride expected from the argument a (will get transposed)
214
- // We compute C^T = B^T * A^T, but we assume B is transposed before
215
- // compression and hence the bt_* naming
216
- using LayoutB = typename Gemm::GemmKernel::CollectiveMainloop::LayoutA;
217
- using LayoutE = typename Gemm::GemmKernel::CollectiveMainloop::LayoutE;
218
-
219
- // M, N, K after transposition
220
- int32_t m = out.size(1);
221
- int32_t n = out.size(0);
222
- int32_t k = a.size(1);
223
-
224
- int64_t lda = a.stride(0);
225
- int64_t ldc = out.stride(0);
226
-
227
- using StrideA = Stride<int64_t, Int<1>, int64_t>;
228
- using StrideC = Stride<Int<1>, int64_t, int64_t>;
229
-
230
- StrideA a_stride{lda, Int<1>{}, Int<0>{}};
231
- StrideC c_stride{Int<1>{}, ldc, Int<0>{}};
232
-
233
- using GemmKernel = typename Gemm::GemmKernel;
234
- typename GemmKernel::ProblemShape prob_shape{m, n, k, 1};
235
-
236
- using ElementE = typename GemmKernel::CollectiveMainloop::ElementE;
237
- using SparseConfig = typename GemmKernel::CollectiveMainloop::SparseConfig;
238
-
239
- LayoutB b_layout = SparseConfig::fill_layoutA(prob_shape);
240
- LayoutE e_layout = SparseConfig::fill_layoutE(prob_shape);
241
-
242
- auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
243
- auto b_ptr = static_cast<ElementAB*>(bt_nzs.data_ptr());
244
- auto e_ptr = static_cast<ElementE*>(bt_meta.data_ptr());
245
- typename GemmKernel::MainloopArguments mainloop_args{
246
- b_ptr, b_layout, a_ptr, a_stride, e_ptr, e_layout};
247
-
248
- auto c_ptr = static_cast<ElementD*>(out.data_ptr());
249
- typename GemmKernel::EpilogueArguments epilogue_args{
250
- Gemm::Epilogue::prepare_args(
251
- std::forward<EpilogueArgs>(epilogue_params)...),
252
- c_ptr, c_stride, c_ptr, c_stride};
253
-
254
- typename GemmKernel::Arguments args{cutlass::gemm::GemmUniversalMode::kGemm,
255
- prob_shape, mainloop_args, epilogue_args};
256
-
257
- // Launch the CUTLASS GEMM kernel.
258
- using GemmOp = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
259
- GemmOp gemm_op;
260
- CUTLASS_CHECK(gemm_op.can_implement(args));
261
-
262
- size_t workspace_size = gemm_op.get_workspace_size(args);
263
- auto const workspace_options =
264
- torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
265
- auto workspace = torch::empty(workspace_size, workspace_options);
266
-
267
- auto stream = at::cuda::getCurrentCUDAStream(a.get_device());
268
-
269
- cutlass::Status status = gemm_op.run(args, workspace.data_ptr(), stream);
270
- CUTLASS_CHECK(status);
271
-}
272
-
273
-//////////////////////////////////////////////////
274
-// Gemm Configs are defined below
275
-//////////////////////////////////////////////////
276
-
277
-template <typename InType, typename OutType,
278
- template <typename, typename, typename> typename Epilogue>
279
-struct sm90_config_default {};
280
-
281
-template <typename OutType,
282
- template <typename, typename, typename> typename Epilogue>
283
-struct sm90_config_default<half_t, OutType, Epilogue> {
284
- using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecialized;
285
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
286
- using TileShape = Shape<_128, _128, _128>;
287
- using ClusterShape = Shape<_1, _1, _1>;
288
- using Cutlass3xGemm =
289
- cutlass_sparse_3x_gemm<half_t, OutType, Epilogue, TileShape, ClusterShape,
290
- KernelSchedule, EpilogueSchedule>;
291
-};
292
-
293
-template <typename OutType,
294
- template <typename, typename, typename> typename Epilogue>
295
-struct sm90_config_default<cutlass::bfloat16_t, OutType, Epilogue> {
296
- using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecialized;
297
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
298
- using TileShape = Shape<_128, _128, _128>;
299
- using ClusterShape = Shape<_1, _1, _1>;
300
- using Cutlass3xGemm =
301
- cutlass_sparse_3x_gemm<cutlass::bfloat16_t, OutType, Epilogue, TileShape,
302
- ClusterShape, KernelSchedule, EpilogueSchedule>;
303
-};
304
-
305
-//////////////////////// Cherry-Picking Kernels ////////////////////////
306
-template <typename InType, typename OutType,
307
- template <typename, typename, typename> typename Epilogue>
308
-struct sm90_fp8_config_1 {
309
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
310
- using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
311
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
312
- using TileShape = Shape<_64, _64, _256>;
313
- using ClusterShape = Shape<_8, _1, _1>;
314
- using Cutlass3xGemm =
315
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
316
- KernelSchedule, EpilogueSchedule>;
317
-};
318
-
319
-template <typename InType, typename OutType,
320
- template <typename, typename, typename> typename Epilogue>
321
-struct sm90_fp8_config_2 {
322
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
323
- using KernelSchedule =
324
- cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
325
- using EpilogueSchedule =
326
- typename cutlass::epilogue::TmaWarpSpecializedCooperative;
327
- using TileShape = Shape<_128, _64, _256>;
328
- using ClusterShape = Shape<_8, _1, _1>;
329
- using Cutlass3xGemm =
330
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
331
- KernelSchedule, EpilogueSchedule>;
332
-};
333
-
334
-template <typename InType, typename OutType,
335
- template <typename, typename, typename> typename Epilogue>
336
-struct sm90_fp8_config_3 {
337
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
338
- using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
339
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
340
- using TileShape = Shape<_64, _64, _256>;
341
- using ClusterShape = Shape<_1, _2, _1>;
342
- using Cutlass3xGemm =
343
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
344
- KernelSchedule, EpilogueSchedule>;
345
-};
346
-
347
-template <typename InType, typename OutType,
348
- template <typename, typename, typename> typename Epilogue>
349
-struct sm90_fp8_config_4 {
350
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
351
- using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
352
- using EpilogueSchedule =
353
- typename cutlass::epilogue::TmaWarpSpecializedCooperative;
354
- using TileShape = Shape<_64, _128, _256>;
355
- using ClusterShape = Shape<_8, _1, _1>;
356
- using Cutlass3xGemm =
357
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
358
- KernelSchedule, EpilogueSchedule>;
359
-};
360
-
361
-template <typename InType, typename OutType,
362
- template <typename, typename, typename> typename Epilogue>
363
-struct sm90_fp8_config_5 {
364
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
365
- using KernelSchedule =
366
- cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum;
367
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
368
- using TileShape = Shape<_128, _128, _256>;
369
- using ClusterShape = Shape<_8, _1, _1>;
370
- using Cutlass3xGemm =
371
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
372
- KernelSchedule, EpilogueSchedule>;
373
-};
374
-
375
-template <typename InType, typename OutType,
376
- template <typename, typename, typename> typename Epilogue>
377
-struct sm90_fp8_config_6 {
378
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
379
- using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
380
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
381
- using TileShape = Shape<_64, _128, _256>;
382
- using ClusterShape = Shape<_1, _2, _1>;
383
- using Cutlass3xGemm =
384
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
385
- KernelSchedule, EpilogueSchedule>;
386
-};
387
-
388
-template <typename InType, typename OutType,
389
- template <typename, typename, typename> typename Epilogue>
390
-struct sm90_fp8_config_7 {
391
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
392
- using KernelSchedule =
393
- cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
394
- using EpilogueSchedule =
395
- typename cutlass::epilogue::TmaWarpSpecializedCooperative;
396
- using TileShape = Shape<_128, _128, _256>;
397
- using ClusterShape = Shape<_1, _1, _1>;
398
- using Cutlass3xGemm =
399
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
400
- KernelSchedule, EpilogueSchedule>;
401
-};
402
-
403
-template <typename InType, typename OutType,
404
- template <typename, typename, typename> typename Epilogue>
405
-struct sm90_fp8_config_8 {
406
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
407
- using KernelSchedule =
408
- cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
409
- using EpilogueSchedule =
410
- typename cutlass::epilogue::TmaWarpSpecializedCooperative;
411
- using TileShape = Shape<_128, _256, _128>;
412
- using ClusterShape = Shape<_8, _1, _1>;
413
- using Cutlass3xGemm =
414
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
415
- KernelSchedule, EpilogueSchedule>;
416
-};
417
-////////////////////////////////////////////////////////////////////////
418
-
419
-template <typename OutType,
420
- template <typename, typename, typename> typename Epilogue>
421
-struct sm90_config_default<cutlass::float_e4m3_t, OutType, Epilogue> {
422
- // M in (128, inf)
423
- using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
424
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
425
- using TileShape = Shape<_128, _128, _128>;
426
- using ClusterShape = Shape<_1, _2, _1>;
427
- using Cutlass3xGemm =
428
- cutlass_sparse_3x_gemm<cutlass::float_e4m3_t, OutType, Epilogue,
429
- TileShape, ClusterShape, KernelSchedule,
430
- EpilogueSchedule>;
431
-};
432
-
433
-template <typename InType, typename OutType,
434
- template <typename, typename, typename> typename Epilogue>
435
-struct sm90_fp8_config_M64 {
436
- // M in [1, 64]
437
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
438
- using KernelSchedule = cutlass::gemm::KernelTmaWarpSpecializedFP8FastAccum;
439
- using EpilogueSchedule =
440
- typename cutlass::epilogue::TmaWarpSpecializedCooperative;
441
- using TileShape = Shape<_64, _64, _256>;
442
- using ClusterShape = Shape<_1, _1, _1>;
443
-
444
- using Cutlass3xGemm =
445
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
446
- KernelSchedule, EpilogueSchedule>;
447
-};
448
-
449
-template <typename InType, typename OutType,
450
- template <typename, typename, typename> typename Epilogue>
451
-struct sm90_fp8_config_M128 {
452
- // M in (64, 128]
453
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
454
- using KernelSchedule =
455
- cutlass::gemm::KernelTmaWarpSpecializedPingpongFP8FastAccum;
456
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
457
- using TileShape = Shape<_64, _128, _256>;
458
- using ClusterShape = Shape<_1, _1, _1>;
459
-
460
- using Cutlass3xGemm =
461
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
462
- KernelSchedule, EpilogueSchedule>;
463
-};
464
-
465
-template <typename InType, typename OutType,
466
- template <typename, typename, typename> typename Epilogue>
467
-struct sm90_fp8_config_M256 {
468
- // M in (128, 256]
469
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
470
- using KernelSchedule =
471
- cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
472
- using EpilogueSchedule =
473
- typename cutlass::epilogue::TmaWarpSpecializedCooperative;
474
- using TileShape = Shape<_128, _128, _256>;
475
- using ClusterShape = Shape<_1, _1, _1>;
476
-
477
- using Cutlass3xGemm =
478
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
479
- KernelSchedule, EpilogueSchedule>;
480
-};
481
-
482
-template <typename InType, typename OutType,
483
- template <typename, typename, typename> typename Epilogue>
484
-struct sm90_fp8_config_M512 {
485
- // M in (256, ]
486
- static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
487
- using KernelSchedule =
488
- cutlass::gemm::KernelTmaWarpSpecializedCooperativeFP8FastAccum;
489
- using EpilogueSchedule =
490
- typename cutlass::epilogue::TmaWarpSpecializedCooperative;
491
- using TileShape = Shape<_128, _128, _256>;
492
- using ClusterShape = Shape<_1, _1, _1>;
493
-
494
- using Cutlass3xGemm =
495
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
496
- KernelSchedule, EpilogueSchedule>;
497
-};
498
-
499
-template <typename OutType,
500
- template <typename, typename, typename> typename Epilogue>
501
-struct sm90_config_default<int8_t, OutType, Epilogue> {
502
- // For M > 128 and any N
503
- using KernelSchedule =
504
- typename cutlass::gemm::KernelTmaWarpSpecializedPingpong;
505
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
506
- using TileShape = Shape<_128, _128, _128>;
507
- using ClusterShape = Shape<_2, _1, _1>;
508
- using Cutlass3xGemm =
509
- cutlass_sparse_3x_gemm<int8_t, OutType, Epilogue, TileShape, ClusterShape,
510
- KernelSchedule, EpilogueSchedule>;
511
-};
512
-
513
-template <typename InType, typename OutType,
514
- template <typename, typename, typename> typename Epilogue>
515
-struct sm90_int8_config_M128 {
516
- // For M in (64, 128] and any N
517
- static_assert(std::is_same<InType, int8_t>());
518
- using KernelSchedule =
519
- typename cutlass::gemm::KernelTmaWarpSpecializedPingpong;
520
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
521
- using TileShape = Shape<_64, _128, _128>;
522
- using ClusterShape = Shape<_2, _1, _1>;
523
- using Cutlass3xGemm =
524
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
525
- KernelSchedule, EpilogueSchedule>;
526
-};
527
-
528
-template <typename InType, typename OutType,
529
- template <typename, typename, typename> typename Epilogue>
530
-struct sm90_int8_config_M64 {
531
- // For M in (32, 64] and any N
532
- static_assert(std::is_same<InType, int8_t>());
533
- using KernelSchedule = typename cutlass::gemm::KernelTmaWarpSpecialized;
534
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
535
- using TileShape = Shape<_64, _64, _256>;
536
- using ClusterShape = Shape<_1, _1, _1>;
537
- using Cutlass3xGemm =
538
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
539
- KernelSchedule, EpilogueSchedule>;
540
-};
541
-
542
-template <typename InType, typename OutType,
543
- template <typename, typename, typename> typename Epilogue>
544
-struct sm90_int8_config_M32_NBig {
545
- // For M in [1, 32] and N >= 8192
546
- static_assert(std::is_same<InType, int8_t>());
547
- using KernelSchedule = typename cutlass::gemm::KernelTmaWarpSpecialized;
548
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
549
- using TileShape = Shape<_64, _128, _256>;
550
- using ClusterShape = Shape<_1, _4, _1>;
551
- using Cutlass3xGemm =
552
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
553
- KernelSchedule, EpilogueSchedule>;
554
-};
555
-
556
-template <typename InType, typename OutType,
557
- template <typename, typename, typename> typename Epilogue>
558
-struct sm90_int8_config_M32_NSmall {
559
- // For M in [1, 32] and N < 8192
560
- static_assert(std::is_same<InType, int8_t>());
561
- using KernelSchedule = typename cutlass::gemm::KernelTmaWarpSpecialized;
562
- using EpilogueSchedule = typename cutlass::epilogue::TmaWarpSpecialized;
563
- using TileShape = Shape<_64, _64, _256>;
564
- using ClusterShape = Shape<_1, _8, _1>;
565
- using Cutlass3xGemm =
566
- cutlass_sparse_3x_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
567
- KernelSchedule, EpilogueSchedule>;
568
-};
569
-
570
-} // namespace
csrc/sparse/cutlass/sparse_scaled_mm_entry.cu
deleted
-104
@@ -1,104 +0,0 @@
1
-#include <cudaTypedefs.h>
2
-
3
-#include <c10/cuda/CUDAGuard.h>
4
-#include <torch/all.h>
5
-
6
-#include "cutlass_extensions/common.hpp"
7
-
8
-bool cutlass_sparse_scaled_mm_supported(int64_t cuda_device_capability) {
9
- // sparse CUTLASS kernels need exactly hopper and are not forward compatible
10
- // CUDA 12.2 and SM90 (Hopper)
11
-
12
-#if defined CUDA_VERSION
13
- return CUDA_VERSION >= 12020 && cuda_device_capability == 90;
14
-#endif
15
-
16
- return false;
17
-}
18
-
19
-#if defined ENABLE_SPARSE_SCALED_MM_C3X && ENABLE_SPARSE_SCALED_MM_C3X
20
-void cutlass_scaled_sparse_mm_sm90(torch::Tensor& c, torch::Tensor const& a,
21
- torch::Tensor const& b,
22
- torch::Tensor const& e,
23
- torch::Tensor const& a_scales,
24
- torch::Tensor const& b_scales,
25
- std::optional<torch::Tensor> const& bias);
26
-
27
-using CompressorResult = std::tuple<torch::Tensor, torch::Tensor>;
28
-CompressorResult cutlass_sparse_compress_sm90(torch::Tensor const& a);
29
-#endif
30
-
31
-void cutlass_scaled_sparse_mm(torch::Tensor& c, torch::Tensor const& a,
32
- torch::Tensor const& bt_nzs,
33
- torch::Tensor const& bt_meta,
34
- torch::Tensor const& a_scales,
35
- torch::Tensor const& b_scales,
36
- std::optional<torch::Tensor> const& bias) {
37
- // Checks for conformality
38
- TORCH_CHECK(a.dim() == 2 && bt_nzs.dim() == 2 && c.dim() == 2);
39
- TORCH_CHECK(c.size(1) == bt_nzs.size(0) && bt_nzs.size(1) * 2 == a.size(1) &&
40
- a.size(0) == c.size(0));
41
- TORCH_CHECK(a_scales.numel() == 1 || a_scales.numel() == a.size(0));
42
- TORCH_CHECK(b_scales.numel() == 1 || b_scales.numel() == bt_nzs.size(0));
43
-
44
- // Check for strides and alignment
45
- TORCH_CHECK(a.stride(1) == 1 && bt_nzs.stride(1) == 1 &&
46
- c.stride(1) == 1); // Row-major
47
- TORCH_CHECK(c.stride(0) % 16 == 0); // 16 Byte Alignment
48
- TORCH_CHECK(bt_nzs.stride(0) % 16 == 0); // 16 Byte Alignment
49
- TORCH_CHECK(a_scales.is_contiguous() && b_scales.is_contiguous());
50
-
51
- if (bias) {
52
- TORCH_CHECK(bias->numel() == bt_nzs.size(0) && bias->is_contiguous() &&
53
- bias->dim() == 1);
54
- }
55
-
56
- at::cuda::OptionalCUDAGuard const device_guard(device_of(a));
57
- int32_t version_num = get_sm_version_num();
58
-
59
- // Guard against compilation issues for sm90 kernels
60
-#if defined ENABLE_SPARSE_SCALED_MM_C3X && ENABLE_SPARSE_SCALED_MM_C3X
61
- // We build for 9.0a which is not forward compatible, so restrict this to
62
- // Hopper only
63
- if (version_num == 90) {
64
- cutlass_scaled_sparse_mm_sm90(c, a, bt_nzs, bt_meta, a_scales, b_scales,
65
- bias);
66
- return;
67
- }
68
-#endif
69
-
70
- TORCH_CHECK_NOT_IMPLEMENTED(
71
- false,
72
- "No compiled cutlass_scaled_sparse_mm for a compute capability less than "
73
- "CUDA device capability: ",
74
- version_num);
75
-}
76
-
77
-std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a) {
78
- // Check for strides and alignment
79
- TORCH_CHECK(a.stride(1) == 1); // Row-major
80
- TORCH_CHECK(a.stride(0) % 8 == 0); // 8 Byte Alignment for Compression
81
-
82
- at::cuda::OptionalCUDAGuard const device_guard(device_of(a));
83
- int32_t version_num = get_sm_version_num();
84
-
85
- // Guard against compilation issues for sm90 kernels
86
-#if defined ENABLE_SPARSE_SCALED_MM_C3X && ENABLE_SPARSE_SCALED_MM_C3X
87
- // We build for 9.0a which is not forward compatible, so restrict this to
88
- // Hopper only
89
- if (version_num == 90) {
90
- std::vector<torch::Tensor> result_tensors;
91
-
92
- auto [a_meta, a_nzs] = cutlass_sparse_compress_sm90(a);
93
- result_tensors.push_back(std::move(a_nzs));
94
- result_tensors.push_back(std::move(a_meta));
95
- return result_tensors;
96
- }
97
-#endif
98
-
99
- TORCH_CHECK_NOT_IMPLEMENTED(
100
- false,
101
- "No compiled cutlass_sparse_compress for a compute capability equal to "
102
- "CUDA device capability: ",
103
- version_num);
104
-}
csrc/torch_bindings.cpp
-20
@@ -523,26 +523,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
523
ops.impl("cutlass_scaled_mm_supports_block_fp8",
524
&cutlass_scaled_mm_supports_block_fp8);
525
526
- // Check if cutlass sparse scaled_mm is supported for CUDA devices of the
527
- // given capability
528
- ops.def(
529
- "cutlass_sparse_scaled_mm_supported(int cuda_device_capability) -> bool");
530
- ops.impl("cutlass_sparse_scaled_mm_supported",
531
- &cutlass_sparse_scaled_mm_supported);
532
-
533
- // CUTLASS sparse GEMM, supporting symmetric per-tensor or per-row/column
534
- // quantization, as well as bias
535
- ops.def(
536
- "cutlass_scaled_sparse_mm(Tensor! out, Tensor a,"
537
- " Tensor bt_nzs,"
538
- " Tensor bt_meta, Tensor a_scales,"
539
- " Tensor b_scales, Tensor? bias) -> ()");
540
- ops.impl("cutlass_scaled_sparse_mm", torch::kCUDA, &cutlass_scaled_sparse_mm);
541
-
542
- // CUTLASS sparse matrix compressor
543
- ops.def("cutlass_sparse_compress(Tensor a) -> Tensor[]");
544
- ops.impl("cutlass_sparse_compress", &cutlass_sparse_compress);
545
-
526
// SM100 CUTLASS MLA decode
527
ops.def(
528
"sm100_cutlass_mla_decode(Tensor! out, Tensor! lse, Tensor q_nope,"
tests/kernels/quantization/test_cutlass_2of4_sparse.py
deleted
-238
@@ -1,238 +0,0 @@
1
-# SPDX-License-Identifier: Apache-2.0
2
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-"""Tests for sparse cutlass kernels
4
-
5
-Run `pytest tests/kernels/quantization/test_cutlass_2of4_sparse.py`.
6
-"""
7
-
8
-import pytest
9
-import torch
10
-
11
-from tests.kernels.utils import baseline_scaled_mm, to_fp8, to_int8
12
-from vllm import _custom_ops as ops
13
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
14
- sparse_cutlass_supported,
15
-)
16
-from vllm.platforms import current_platform
17
-
18
-CUDA_DEVICES = [
19
- f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
20
-]
21
-
22
-capability = current_platform.get_device_capability()
23
-capability = capability[0] * 10 + capability[1]
24
-
25
-
26
-def to_bf16(tensor: torch.Tensor) -> torch.Tensor:
27
- return tensor.to(dtype=torch.bfloat16)
28
-
29
-
30
-def to_fp16(tensor: torch.Tensor) -> torch.Tensor:
31
- return tensor.to(dtype=torch.float16)
32
-
33
-
34
-def prune_to_2_4(tensor):
35
- # Reshape tensor to [N, 4] where N is number of groups of 4
36
- original_shape = tensor.shape
37
- reshaped = tensor.reshape(-1, 4)
38
-
39
- # Get indices of top 2 absolute values in each group of 4
40
- _, indices = torch.topk(torch.abs(reshaped), k=2, dim=1)
41
-
42
- # Create binary mask
43
- mask = torch.zeros_like(reshaped)
44
- mask.scatter_(dim=1, index=indices, src=torch.ones_like(indices, dtype=mask.dtype))
45
-
46
- # Apply mask and reshape back
47
- pruned = reshaped * mask
48
-
49
- # Turn all -0.0 to 0.0
50
- pruned[pruned == -0.0] = 0.0
51
-
52
- return pruned.reshape(original_shape)
53
-
54
-
55
-# This function checks that applying an identity matrix multiplication
56
-# to the compressed weights yields the original uncompressed weights.
57
-def check_compress_decompress_invariance(
58
- dtype: torch.dtype,
59
- b: torch.Tensor,
60
- b_compressed: torch.Tensor,
61
- b_metadata: torch.Tensor,
62
-):
63
- # For float16 and bfloat16, cutlass_scaled_sparse_mm's output must be the
64
- # same dtype as its inputs. This line addresses that constraint while
65
- # arbitrarily using bfloat16 for the int8/fp8 cases.
66
- out_dtype = torch.float16 if dtype is torch.float16 else torch.bfloat16
67
-
68
- eye = torch.eye(b.shape[0], device="cuda", dtype=dtype)
69
- eye_scale = torch.ones(1, device="cuda", dtype=torch.float32)
70
- b_decomp = ops.cutlass_scaled_sparse_mm(
71
- eye, b_compressed, b_metadata, eye_scale, eye_scale, out_dtype=out_dtype
72
- )
73
-
74
- torch.testing.assert_close(b.to(dtype=out_dtype), b_decomp)
75
-
76
-
77
-def make_rand_sparse_tensors(
78
- dtype: torch.dtype, m: int, n: int, k: int
79
-) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
80
- a = torch.randn((m, k), device="cuda")
81
- b = torch.randn((n, k), device="cuda").t()
82
-
83
- if dtype == torch.int8:
84
- # ensure A and B aren't all zeros after rounding
85
- a = a * 5.0
86
- b = b * 5.0
87
-
88
- b = prune_to_2_4(b.t()).t()
89
-
90
- if dtype == torch.int8:
91
- a, b = to_int8(a), to_int8(b)
92
- elif dtype == torch.float8_e4m3fn:
93
- a, b = to_fp8(a), to_fp8(b)
94
- elif dtype == torch.float16:
95
- a, b = to_fp16(a), to_fp16(b)
96
- elif dtype == torch.bfloat16:
97
- a, b = to_bf16(a), to_bf16(b)
98
- else:
99
- raise ValueError("unsupported dtype")
100
-
101
- b_compressed, e = ops.cutlass_sparse_compress(b.t())
102
- check_compress_decompress_invariance(dtype, b, b_compressed, e)
103
-
104
- # Compressed B, Metadata, Original A, B
105
- return b_compressed, e, a, b
106
-
107
-
108
-@pytest.mark.skipif(
109
- not sparse_cutlass_supported(),
110
- reason="Sparse CUTLASS is not supported on this GPU type.",
111
-)
112
-# Test working with a subset of A and B for sparse matmul
113
-def test_cutlass_sparse_subset():
114
- big_m = 1024
115
- m, n, k = 512, 512, 512
116
-
117
- # Create tensors
118
- b_comp, e, whole_a, b = make_rand_sparse_tensors(torch.float8_e4m3fn, big_m, n, k)
119
- a = whole_a[0:m, 0:k]
120
- scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
121
- scale_b = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
122
-
123
- out = ops.cutlass_scaled_sparse_mm(
124
- a, b_comp, e, scale_a, scale_b, out_dtype=torch.bfloat16
125
- )
126
- baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype=torch.bfloat16)
127
-
128
- torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
129
-
130
-
131
-MNK_FACTORS = [
132
- (1, 256, 128),
133
- (1, 16384, 1024),
134
- (1, 24576, 512),
135
- (16, 256, 512),
136
- (16, 16384, 128),
137
- (16, 24576, 4096),
138
- (32, 8192, 4096),
139
- (32, 16384, 4096),
140
- (33, 1024, 1024),
141
- (33, 8192, 128),
142
- (64, 2048, 512),
143
- (64, 16384, 1024),
144
- (100, 8192, 512),
145
- (128, 32768, 4096),
146
- (256, 4096, 4096),
147
- (512, 256, 1024),
148
- (512, 8192, 4096),
149
- (512, 16384, 128),
150
- (512, 24576, 128),
151
-]
152
-
153
-
154
-# Test working with a subset of A and B for sparse matmul
155
-@pytest.mark.skipif(
156
- not sparse_cutlass_supported(),
157
- reason="Sparse CUTLASS is not supported on this GPU type.",
158
-)
159
-@pytest.mark.parametrize("m, n, k", MNK_FACTORS)
160
-@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
161
-@pytest.mark.parametrize("use_bias", [True, False])
162
-def test_cutlass_sparse_gemm(
163
- m: int, k: int, n: int, dtype: type[torch.dtype], use_bias: bool
164
-):
165
- # Create tensors
166
- b_comp, e, a, b = make_rand_sparse_tensors(dtype, m, n, k)
167
- scale_a = torch.ones((1, 1), device="cuda", dtype=torch.float32)
168
- scale_b = torch.ones((1, 1), device="cuda", dtype=torch.float32)
169
-
170
- bias = torch.rand((n,), device="cuda", dtype=dtype) if use_bias else None
171
-
172
- out = ops.cutlass_scaled_sparse_mm(
173
- a, b_comp, e, scale_a, scale_b, out_dtype=dtype, bias=bias
174
- )
175
-
176
- baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype=dtype, bias=bias)
177
-
178
- torch.testing.assert_close(out, baseline, rtol=1e-2, atol=3e-1)
179
-
180
-
181
-@pytest.mark.skipif(
182
- not sparse_cutlass_supported(),
183
- reason="Sparse CUTLASS is not supported on this GPU type.",
184
-)
185
-@pytest.mark.parametrize("m, k, n", MNK_FACTORS)
186
-@pytest.mark.skipif(
187
- not current_platform.has_device_capability(89),
188
- reason="FP8 is not supported on this GPU type.",
189
-)
190
-@pytest.mark.parametrize("use_bias", [True, False])
191
-def test_cutlass_sparse_fp8_gemm(m: int, n: int, k: int, use_bias: bool):
192
- # Create tensors
193
- b_comp, e, a, b = make_rand_sparse_tensors(torch.float8_e4m3fn, m, n, k)
194
- scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32)
195
- scale_b = torch.randn((1, 1), device="cuda", dtype=torch.float32)
196
- out_dtype = torch.bfloat16
197
-
198
- bias = torch.rand((n,), device="cuda", dtype=out_dtype) * 10 if use_bias else None
199
-
200
- out = ops.cutlass_scaled_sparse_mm(
201
- a, b_comp, e, scale_a, scale_b, out_dtype=out_dtype, bias=bias
202
- )
203
-
204
- baseline = baseline_scaled_mm(
205
- a, b, scale_a, scale_b, out_dtype=out_dtype, bias=bias
206
- )
207
-
208
- torch.testing.assert_close(out, baseline, rtol=1e-2, atol=3e-1)
209
-
210
-
211
-@pytest.mark.skipif(
212
- not sparse_cutlass_supported(),
213
- reason="Sparse CUTLASS is not supported on this GPU type.",
214
-)
215
-@pytest.mark.parametrize("m,k,n", MNK_FACTORS)
216
-@pytest.mark.parametrize("per_act_token", [True, False])
217
-@pytest.mark.parametrize("per_out_ch", [True, False])
218
-@pytest.mark.parametrize("use_bias", [True, False])
219
-def test_cutlass_sparse_int8_gemm(
220
- m: int, n: int, k: int, per_act_token: bool, per_out_ch: bool, use_bias: bool
221
-):
222
- # Create tensors
223
- b_comp, e, a, b = make_rand_sparse_tensors(torch.int8, m, n, k)
224
- scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32)
225
- scale_b = torch.randn((1, 1), device="cuda", dtype=torch.float32)
226
- out_dtype = torch.bfloat16
227
-
228
- bias = torch.rand((n,), device="cuda", dtype=out_dtype) * 10 if use_bias else None
229
-
230
- out = ops.cutlass_scaled_sparse_mm(
231
- a, b_comp, e, scale_a, scale_b, out_dtype=out_dtype, bias=bias
232
- )
233
-
234
- baseline = baseline_scaled_mm(
235
- a, b, scale_a, scale_b, out_dtype=out_dtype, bias=bias
236
- )
237
-
238
- torch.testing.assert_close(out, baseline, rtol=1e0, atol=2e0)
tests/quantization/test_compressed_tensors.py
-281
@@ -12,7 +12,6 @@ from compressed_tensors.quantization import QuantizationType
12
from tests.models.utils import check_logprobs_close
13
from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod
14
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
15
- CompressedTensors24,
15
CompressedTensorsLinearMethod,
16
CompressedTensorsW4A4Fp4,
17
CompressedTensorsW4A8Fp8,
@@ -27,9 +26,6 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8
26
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
27
cutlass_fp4_supported,
28
)
30
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
31
- sparse_cutlass_supported,
32
-)
29
from vllm.platforms import current_platform
30
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
31
@@ -362,283 +358,6 @@ def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner):
358
assert output
359
360
365
-@pytest.mark.skipif(
366
- not sparse_cutlass_supported(),
367
- reason="Sparse FP8 is not yet supported on this GPU type.",
368
-)
369
-def _test_2of4_quant_models(qkv_proj, weight_strategy, input_strategy, format="dense"):
370
- assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
371
- assert isinstance(qkv_proj.scheme, CompressedTensors24)
372
-
373
- assert qkv_proj.scheme.weight_quant.strategy == weight_strategy
374
- assert qkv_proj.scheme.input_quant.strategy == input_strategy
375
- assert qkv_proj.scheme.quantized
376
- assert qkv_proj.quant_method.quantization_config.sparsity_scheme_map
377
- sparsity_map = qkv_proj.quant_method.quantization_config.sparsity_scheme_map # noqa: E501
378
- assert sparsity_map.get("Linear").format == format
379
- assert sparsity_map.get("Linear").sparsity_structure == "2:4"
380
-
381
-
382
-@pytest.mark.skipif(
383
- not current_platform.is_cuda() or not current_platform.has_device_capability(90),
384
- reason="Sparse FP8 is not yet supported on this GPU type.",
385
-)
386
-@pytest.mark.parametrize(
387
- "args_2of4",
388
- [
389
- (
390
- "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Dynamic-2of4-testing",
391
- "channel",
392
- "token",
393
- ),
394
- (
395
- "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Static-Per-Tensor-testing",
396
- "channel",
397
- "tensor",
398
- ),
399
- (
400
- "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Static-testing",
401
- "tensor",
402
- "tensor",
403
- ),
404
- (
405
- "nm-testing/Meta-Llama-3-8B-Instruct-FP8-Dynamic-IA-Per-Tensor-Weight-testing",
406
- "tensor",
407
- "token",
408
- ),
409
- ],
410
-)
411
-def test_compressed_tensors_2of4_quant_fp8(vllm_runner, args_2of4):
412
- model, weight_strategy, input_strategy = args_2of4
413
- with vllm_runner(model, enforce_eager=True) as llm:
414
-
415
- def check_model(model):
416
- layer = model.model.layers[0]
417
-
418
- qkv_proj = layer.self_attn.qkv_proj
419
- assert qkv_proj.scheme.weights_dtype == torch.float8_e4m3fn
420
- _test_2of4_quant_models(qkv_proj, weight_strategy, input_strategy)
421
-
422
- llm.apply_model(check_model)
423
-
424
- output = llm.generate_greedy("Hello my name is", max_tokens=4)
425
- print(output)
426
- assert output
427
-
428
-
429
-@pytest.mark.skipif(
430
- not current_platform.is_cuda() or not current_platform.has_device_capability(90),
431
- reason="Sparse FP8 is not yet supported on this GPU type.",
432
-)
433
-@pytest.mark.parametrize(
434
- "args_2of4",
435
- [
436
- (
437
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_fp8-BitM",
438
- "channel",
439
- "token",
440
- ),
441
- (
442
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_tensor_act_fp8-BitM",
443
- "channel",
444
- "tensor",
445
- ),
446
- (
447
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_per_tok_dyn_act_fp8-BitM",
448
- "tensor",
449
- "token",
450
- ),
451
- (
452
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_tensor_act_fp8-BitM",
453
- "tensor",
454
- "tensor",
455
- ),
456
- ],
457
-)
458
-def test_compressed_tensors_2of4_quant_fp8_compressed(vllm_runner, args_2of4):
459
- model, weight_strategy, input_strategy = args_2of4
460
- with vllm_runner(model, enforce_eager=True) as llm:
461
-
462
- def check_model(model):
463
- layer = model.model.layers[0]
464
-
465
- qkv_proj = layer.self_attn.qkv_proj
466
- assert qkv_proj.scheme.weights_dtype == torch.float8_e4m3fn
467
- _test_2of4_quant_models(
468
- qkv_proj,
469
- weight_strategy,
470
- input_strategy,
471
- format="sparse-24-bitmask",
472
- )
473
-
474
- llm.apply_model(check_model)
475
-
476
- output = llm.generate_greedy("Hello my name is", max_tokens=4)
477
- print(output)
478
- assert output
479
-
480
-
481
-@pytest.mark.skipif(
482
- not sparse_cutlass_supported(),
483
- reason="cutlass is not yet supported on this GPU type.",
484
-)
485
-@pytest.mark.parametrize(
486
- "args_2of4",
487
- [
488
- (
489
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_per_tok_dyn_act_int8-BitM",
490
- "channel",
491
- "token",
492
- ),
493
- (
494
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-chnl_wts_tensor_act_int8-BitM",
495
- "channel",
496
- "tensor",
497
- ),
498
- (
499
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_per_tok_dyn_act_int8-BitM",
500
- "tensor",
501
- "token",
502
- ),
503
- (
504
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-gsm8k-pruned.2of4-tensor_wts_tensor_act_int8-BitM",
505
- "tensor",
506
- "tensor",
507
- ),
508
- ],
509
-)
510
-def test_compressed_tensors_2of4_quant_int8_compressed(vllm_runner, args_2of4):
511
- model, weight_strategy, input_strategy = args_2of4
512
- with vllm_runner(model, enforce_eager=True) as llm:
513
-
514
- def check_model(model):
515
- layer = model.model.layers[0]
516
-
517
- qkv_proj = layer.self_attn.qkv_proj
518
- assert qkv_proj.scheme.weights_dtype == torch.int8
519
- _test_2of4_quant_models(
520
- qkv_proj,
521
- weight_strategy,
522
- input_strategy,
523
- format="sparse-24-bitmask",
524
- )
525
-
526
- llm.apply_model(check_model)
527
-
528
- output = llm.generate_greedy("Hello my name is", max_tokens=4)
529
- print(output)
530
- assert output
531
-
532
-
533
-@pytest.mark.skipif(
534
- not sparse_cutlass_supported(),
535
- reason="Sparse FP8 is not yet supported on this GPU type.",
536
-)
537
-@pytest.mark.parametrize(
538
- "args_2of4",
539
- [
540
- (
541
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Dynamic-IA-Per-Channel-Weight-testing",
542
- "channel",
543
- "token",
544
- ),
545
- (
546
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Static-testing",
547
- "tensor",
548
- "tensor",
549
- ),
550
- (
551
- "nm-testing/TinyLlama-1.1B-Chat-v1.0-INT8-Dynamic-IA-Per-Tensor-Weight-testing",
552
- "tensor",
553
- "token",
554
- ),
555
- ],
556
-)
557
-def test_compressed_tensors_2of4_quant_int8(vllm_runner, args_2of4):
558
- model, weight_strategy, input_strategy = args_2of4
559
- with vllm_runner(model, enforce_eager=True) as llm:
560
-
561
- def check_model(model):
562
- layer = model.model.layers[0]
563
-
564
- qkv_proj = layer.self_attn.qkv_proj
565
- assert qkv_proj.scheme.weights_dtype == torch.int8
566
- _test_2of4_quant_models(qkv_proj, weight_strategy, input_strategy)
567
-
568
- llm.apply_model(check_model)
569
-
570
- output = llm.generate_greedy("Hello my name is", max_tokens=4)
571
- print(output)
572
- assert output
573
-
574
-
575
-@pytest.mark.skipif(
576
- not sparse_cutlass_supported(),
577
- reason="2of4 Sparse is not yet supported on this GPU type.",
578
-)
579
-@pytest.mark.parametrize(
580
- "args_2of4",
581
- [("nm-testing/TinyLlama-1.1B-Chat-v1.0-2of4-Sparse-Dense-Compressor")],
582
-)
583
-def test_compressed_tensors_2of4_sparse(vllm_runner, args_2of4):
584
- model = args_2of4
585
- with vllm_runner(model, enforce_eager=True) as llm:
586
-
587
- def check_model(model):
588
- layer = model.model.layers[0]
589
-
590
- qkv_proj = layer.self_attn.qkv_proj
591
- assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
592
- assert isinstance(qkv_proj.scheme, CompressedTensors24)
593
-
594
- assert qkv_proj.scheme.weight_quant is None
595
- assert qkv_proj.scheme.input_quant is None
596
- assert not qkv_proj.scheme.quantized
597
- assert qkv_proj.quant_method.quantization_config.sparsity_scheme_map
598
- sparsity_map = qkv_proj.quant_method.quantization_config.sparsity_scheme_map # noqa: E501
599
- assert sparsity_map.get("Linear").format == "dense"
600
- assert sparsity_map.get("Linear").sparsity_structure == "2:4"
601
-
602
- llm.apply_model(check_model)
603
-
604
- output = llm.generate_greedy("Hello my name is", max_tokens=4)
605
- print(output)
606
- assert output
607
-
608
-
609
-@pytest.mark.skipif(
610
- not sparse_cutlass_supported(),
611
- reason="Cutlass is not yet supported on this GPU type.",
612
-)
613
-@pytest.mark.parametrize(
614
- "args_2of4", [("nm-testing/llama2.c-stories42M-pruned2.4-compressed")]
615
-)
616
-def test_compressed_tensors_2of4_sparse_compressed(vllm_runner, args_2of4):
617
- model = args_2of4
618
- with vllm_runner(model, enforce_eager=True) as llm:
619
-
620
- def check_model(model):
621
- layer = model.model.layers[0]
622
-
623
- qkv_proj = layer.self_attn.qkv_proj
624
- assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
625
- assert isinstance(qkv_proj.scheme, CompressedTensors24)
626
-
627
- assert qkv_proj.scheme.weight_quant is None
628
- assert qkv_proj.scheme.input_quant is None
629
- assert not qkv_proj.scheme.quantized
630
- assert qkv_proj.quant_method.quantization_config.sparsity_scheme_map
631
- sparsity_map = qkv_proj.quant_method.quantization_config.sparsity_scheme_map # noqa: E501
632
- assert sparsity_map.get("Linear").format == "sparse-24-bitmask"
633
- assert sparsity_map.get("Linear").sparsity_structure == "2:4"
634
-
635
- llm.apply_model(check_model)
636
-
637
- output = llm.generate_greedy("Hello my name is", max_tokens=4)
638
- print(output)
639
- assert output
640
-
641
-
361
@pytest.mark.skipif(
362
not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform."
363
)
tests/weight_loading/models.txt
-2
@@ -20,8 +20,6 @@ compressed-tensors, nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test, main
20
compressed-tensors, nm-testing/Phi-3-mini-128k-instruct-FP8, main
21
compressed-tensors, neuralmagic/Phi-3-medium-128k-instruct-quantized.w4a16, main
22
#compressed-tensors, mgoin/DeepSeek-Coder-V2-Lite-Instruct-FP8, main
23
-compressed-tensors, nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-FP8-Dynamic-testing, main, 90
24
-compressed-tensors, nm-testing/SparseLlama-3.1-8B-gsm8k-pruned.2of4-W8A8-testing, main, 90
23
awq, casperhansen/mixtral-instruct-awq, main
24
awq_marlin, casperhansen/mixtral-instruct-awq, main
25
fp8, neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV, main
vllm/_custom_ops.py
-92
@@ -876,10 +876,6 @@ def cutlass_scaled_mm_azp(
876
return out.view(*target_shape)
877
878
879
-def cutlass_sparse_scaled_mm_supported(cuda_device_capability: int) -> bool:
880
- return torch.ops._C.cutlass_sparse_scaled_mm_supported(cuda_device_capability)
881
-
882
-
879
def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool:
880
if cuda_device_capability < 90 or cuda_device_capability >= 110:
881
return False
@@ -890,94 +886,6 @@ def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool:
886
return False
887
888
893
-def cutlass_sparse_compress(a: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
894
- """
895
- Compresses a sparse matrix for use with Cutlass sparse operations.
896
-
897
- This function takes a dense tensor and compresses it into two components:
898
- non-zero elements and metadata. The compressed representation is compatible
899
- with Cutlass sparse kernels.
900
-
901
- Args:
902
- a (torch.Tensor):
903
- The input tensor to be compressed. Must have one of the following data types:
904
- - `torch.int8`
905
- - `torch.float8_e4m3fn`
906
- - `torch.bfloat16`
907
- - `torch.float16`
908
-
909
- Returns:
910
- tuple[torch.Tensor, torch.Tensor]:
911
- A tuple containing:
912
- - `a_nzs` (torch.Tensor): A tensor containing non-zero elements of `a`.
913
- - `a_meta` (torch.Tensor): A tensor containing metadata for the sparse representation.
914
-
915
- Raises:
916
- ValueError: If the compression operation fails.
917
-
918
- Notes:
919
- - The `a_meta` tensor has a data type of `torch.uint8`.
920
- - Each metadata element encodes the sparsity of 4 non-zero elements (i.e., `elemsPerMetaElem = 4`).
921
- - The shape of `a_nzs` is `(m, k // 2)`, where `m` and `k` are the dimensions of the input tensor.
922
- - The shape of `a_meta` is `(m, k // 2 // elemsPerMetaElem)`.
923
- """
924
- assert a.dtype in [torch.int8, torch.float8_e4m3fn, torch.bfloat16, torch.float16]
925
- assert a.is_contiguous()
926
-
927
- # a_meta.dtype: torch.uint8 so elemsPerMetaElem = 8b / 2b_per_nz = 4
928
- elemsPerMetaElem = 4
929
- assert a.shape[1] % (2 * elemsPerMetaElem) == 0
930
-
931
- return torch.ops._C.cutlass_sparse_compress(a)
932
-
933
-
934
-def cutlass_scaled_sparse_mm(
935
- a: torch.Tensor,
936
- bt_nzs: torch.Tensor,
937
- bt_meta: torch.Tensor,
938
- scale_a: torch.Tensor,
939
- scale_b: torch.Tensor,
940
- out_dtype: torch.dtype,
941
- bias: torch.Tensor | None = None,
942
-) -> torch.Tensor:
943
- """
944
- Performs a scaled sparse matrix multiplication using Cutlass.
945
-
946
- Steps:
947
- 1. Create a dense matrix `a` of shape (m, k) on the CUDA device:
948
- `a = torch.randn((m, k), device='cuda')`.
949
-
950
- 2. Create a dense matrix `b` of shape (k, n) on the CUDA device:
951
- `b = torch.randn((k, n), device='cuda')`.
952
-
953
- 3. Prune matrix `b` to 2:4 sparsity along the specified dimension:
954
- `b = prune_to_2_4(b, dim=0)`.
955
-
956
- 4. Compress the transposed sparse matrix `b.t()`:
957
- `bt_nzs, bt_meta = cutlass_sparse_compress(b.t())`.
958
-
959
- 5. Perform sparse matrix multiplication using the compressed matrix,
960
- applying scaling factors for `a` and `b`, and the output data type:
961
- `out = cutlass_scaled_sparse_mm(a, bt_nzs, bt_meta, scale_a, scale_b, out_dtype)`.
962
-
963
- Returns:
964
- - The result of the scaled sparse matrix multiplication.
965
- """
966
- assert bt_nzs.shape[0] % 16 == 0 and bt_nzs.shape[1] % 16 == 0
967
- assert out_dtype is torch.bfloat16 or out_dtype is torch.float16
968
- assert bias is None or bias.shape[0] == bt_nzs.shape[0] and bias.dtype == out_dtype
969
-
970
- m = a.shape[0]
971
- n = bt_nzs.shape[0]
972
- out = torch.empty((m, n), dtype=out_dtype, device=a.device)
973
-
974
- torch.ops._C.cutlass_scaled_sparse_mm(
975
- out, a, bt_nzs, bt_meta, scale_a, scale_b, bias
976
- )
977
-
978
- return out
979
-
980
-
889
def get_cutlass_moe_mm_data(
890
topk_ids: torch.Tensor,
891
expert_offsets: torch.Tensor,
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_24.py
+5
-343
@@ -5,39 +5,16 @@ from collections.abc import Callable
5
from typing import Any
6
7
import torch
8
-from compressed_tensors import CompressionFormat, ModelCompressor
8
from compressed_tensors.quantization import (
9
QuantizationArgs,
11
- QuantizationStrategy,
12
- QuantizationType,
10
)
14
-from compressed_tensors.utils import combine_shards
11
16
-from vllm import _custom_ops as ops
17
-from vllm.model_executor.layers.linear import (
18
- MergedColumnParallelLinear,
19
- QKVParallelLinear,
20
-)
12
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
13
CompressedTensorsScheme,
14
)
24
-from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
25
-from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
26
-from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
27
- convert_to_channelwise,
28
- sparse_cutlass_supported,
29
-)
30
-from vllm.model_executor.parameter import (
31
- BasevLLMParameter,
32
- ChannelQuantScaleParameter,
33
- ModelWeightParameter,
34
- PerTensorScaleParameter,
35
-)
15
16
__all__ = ["CompressedTensors24"]
17
39
-from vllm.platforms import current_platform
40
-
18
19
class CompressedTensors24(CompressedTensorsScheme):
20
def __init__(
@@ -47,33 +24,11 @@ class CompressedTensors24(CompressedTensorsScheme):
24
input_quant: QuantizationArgs | None = None,
25
model_compression_config: dict[str, Any] | None = None,
26
):
50
- self.quantized = quantized
51
- self.weight_quant = weight_quant
52
- self.input_quant = input_quant
53
- model_compressor = ModelCompressor.from_compression_config(
54
- model_compression_config
55
- )
56
- self.do_sparse_decompress = (
57
- model_compressor is not None
58
- and model_compressor.sparsity_config.format
59
- == CompressionFormat.sparse_24_bitmask.value
60
- )
61
- if self.do_sparse_decompress:
62
- self.model_compressor = model_compressor
63
-
64
- if (
65
- quantized
66
- and input_quant is not None
67
- and self._get_quant_dtype() == current_platform.fp8_dtype()
68
- ):
69
- static = not input_quant.dynamic
70
- g_shape = GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN
71
- self.quant_fp8 = QuantFP8(static, g_shape)
27
+ raise NotImplementedError("Sparse24 models are no longer supported by vLLM")
28
29
@classmethod
30
def get_min_capability(cls) -> int:
75
- # Only cutlass 3.x kernels are implemented so far
76
- return 90
31
+ raise NotImplementedError("Sparse24 models are no longer supported by vLLM")
32
33
def create_weights(
34
self,
@@ -85,164 +40,10 @@ class CompressedTensors24(CompressedTensorsScheme):
40
weight_loader: Callable,
41
**kwargs,
42
):
88
- if not sparse_cutlass_supported():
89
- raise ValueError(
90
- "Sparse CUTLASS not supported. vLLM must be built with "
91
- "CUDA 12.2 or later to use this feature"
92
- )
93
-
94
- layer.logical_widths = output_partition_sizes
95
- layer.input_size = input_size
96
- layer.input_size_per_partition = input_size_per_partition
97
- self.weights_dtype: torch.dtype = self._get_params_dtype(params_dtype)
98
-
99
- # parameter to store uncompressed weight
100
- weight = ModelWeightParameter(
101
- data=torch.empty(
102
- sum(output_partition_sizes),
103
- input_size_per_partition,
104
- dtype=self.weights_dtype,
105
- ),
106
- input_dim=1,
107
- output_dim=0,
108
- weight_loader=weight_loader,
109
- )
110
- if self.do_sparse_decompress:
111
- assert all(
112
- partition_size % 8 == 0 for partition_size in output_partition_sizes
113
- ), "All partitions must be divisible by 8 for "
114
- "2:4 sparse compressed models"
115
-
116
- shape = BasevLLMParameter(
117
- data=torch.empty(2, 1, dtype=torch.int64),
118
- weight_loader=weight_loader,
119
- )
120
- compressed_weight = ModelWeightParameter(
121
- data=torch.empty(
122
- sum(output_partition_sizes),
123
- input_size_per_partition // 2,
124
- dtype=self.weights_dtype,
125
- ),
126
- input_dim=1,
127
- output_dim=0,
128
- weight_loader=weight_loader,
129
- )
130
-
131
- bitmask = ModelWeightParameter(
132
- data=torch.empty(
133
- sum(output_partition_sizes),
134
- input_size_per_partition // 8,
135
- dtype=torch.uint8,
136
- ),
137
- input_dim=1,
138
- output_dim=0,
139
- weight_loader=weight_loader,
140
- )
141
-
142
- layer.register_parameter("shape", shape)
143
- layer.register_parameter("compressed", compressed_weight)
144
- layer.register_parameter("bitmask", bitmask)
145
-
146
- # Check if quantized, not just 2:4 Sparse
147
- if self.quantized:
148
- if (
149
- self.weight_quant
150
- and self.weight_quant.strategy == QuantizationStrategy.CHANNEL.value
151
- ):
152
- weight_scale = ChannelQuantScaleParameter(
153
- data=torch.empty(
154
- (sum(output_partition_sizes), 1), dtype=torch.float32
155
- ),
156
- output_dim=0,
157
- weight_loader=weight_loader,
158
- )
159
- else:
160
- assert (
161
- self.weight_quant
162
- and self.weight_quant.strategy == QuantizationStrategy.TENSOR.value
163
- )
164
- weight_scale = PerTensorScaleParameter(
165
- data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
166
- weight_loader=weight_loader,
167
- )
168
-
169
- layer.register_parameter("weight_scale", weight_scale)
170
-
171
- # input quant will be non-none
172
- if self.input_quant and not self.input_quant.dynamic:
173
- # register input quant scale
174
- assert self.input_quant.strategy == QuantizationStrategy.TENSOR.value
175
- input_scale = BasevLLMParameter(
176
- data=torch.empty(1, dtype=torch.float32),
177
- weight_loader=weight_loader,
178
- )
179
-
180
- layer.register_parameter("input_scale", input_scale)
181
-
182
- else:
183
- # for sparse-only, pass in 1 for weight/input scales
184
- weight_scale = torch.nn.Parameter(
185
- data=torch.ones(1, dtype=torch.float32), requires_grad=False
186
- )
187
- input_scale = torch.nn.Parameter(
188
- data=torch.ones(1, dtype=torch.float32), requires_grad=False
189
- )
190
- layer.register_parameter("input_scale", input_scale)
191
- layer.register_parameter("weight_scale", weight_scale)
192
-
193
- layer.register_parameter("weight", weight)
43
+ raise NotImplementedError("Sparse24 models are no longer supported by vLLM")
44
45
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
196
- """
197
- Compress weights after loading. Store compressed weight and meta
198
- tensor
199
-
200
- :post-condition: layer.w_compressed and layer.meta are
201
- set to the compressed weight and meta tensor in the
202
- format expected by the Cutlass kernels
203
- :param layer: The layer with the weights to be processed
204
-
205
- """
206
- if self.do_sparse_decompress:
207
- layer.weight.data = self._decompress_bitmask_compressed_weight(
208
- compressed=layer.compressed,
209
- bitmask=layer.bitmask,
210
- layer=layer,
211
- )
212
-
213
- # compressed and bitmask tensors
214
- # are no longer needed after decompression
215
- del layer.compressed
216
- del layer.bitmask
217
-
218
- # torch.compile workaround
219
- if hasattr(layer, "input_scale"):
220
- layer.input_scale = torch.nn.Parameter(
221
- layer.input_scale.data, requires_grad=False
222
- )
223
-
224
- if self.weight_quant:
225
- if self.weight_quant.strategy == QuantizationStrategy.TENSOR.value:
226
- layer.weight_scale = torch.nn.Parameter(
227
- convert_to_channelwise(
228
- weight_scale=layer.weight_scale,
229
- logical_widths=layer.logical_widths,
230
- ),
231
- requires_grad=False,
232
- )
233
- else:
234
- # torch.compile workaround
235
- layer.weight_scale = torch.nn.Parameter(
236
- layer.weight_scale.data, requires_grad=False
237
- )
238
-
239
- # Set all negative zero values to 0 prior to compression
240
- if layer.weight.dtype.is_floating_point and layer.weight.dtype.itemsize >= 2:
241
- layer.weight.data[layer.weight.data == -0.0] = 0.0
242
-
243
- w_compressed, meta = ops.cutlass_sparse_compress(layer.weight.data)
244
- layer.weight = torch.nn.Parameter(w_compressed, requires_grad=False)
245
- layer.meta = torch.nn.Parameter(meta, requires_grad=False)
46
+ raise NotImplementedError("Sparse24 models are no longer supported by vLLM")
47
48
def apply_weights(
49
self,
@@ -250,143 +51,4 @@ class CompressedTensors24(CompressedTensorsScheme):
51
x: torch.Tensor,
52
bias: torch.Tensor | None = None,
53
) -> torch.Tensor:
253
- """
254
- Returns the output tensor for the layer with 2:4
255
- sparse compressed weights, given the input tensor
256
- and bias
257
-
258
- :param layer: The layer with 2:4 sparse compressed
259
- weights to be used for the computation
260
- :param x: The input tensor to the layer
261
- :param bias: The bias to be added to the output tensor
262
- :return: The output tensor of the layer
263
- """
264
- if self.quantized:
265
- scale = getattr(layer, "input_scale", None)
266
-
267
- if self.weights_dtype == torch.int8:
268
- ops_output = ops.scaled_int8_quant(x, scale=scale)
269
- q_input = ops_output[0]
270
- input_scale = ops_output[1]
271
- else:
272
- assert self.weights_dtype == torch.float8_e4m3fn
273
- q_input, input_scale = self.quant_fp8(x, scale=scale)
274
-
275
- else:
276
- # Not quantized, nothing to do with the input_scales, use as is
277
- input_scale = layer.input_scale
278
- q_input = x
279
-
280
- out = ops.cutlass_scaled_sparse_mm(
281
- a=q_input,
282
- bt_nzs=layer.weight,
283
- bt_meta=layer.meta,
284
- scale_a=input_scale,
285
- scale_b=layer.weight_scale,
286
- out_dtype=x.dtype,
287
- bias=bias,
288
- )
289
-
290
- assert out.is_contiguous()
291
- return out
292
-
293
- def _get_params_dtype(self, params_dtype: torch.dtype) -> torch.dtype:
294
- if not self.quantized:
295
- return params_dtype
296
- return self._get_quant_dtype()
297
-
298
- def _get_quant_dtype(self) -> torch.dtype:
299
- assert self.quantized
300
- assert self.weight_quant is not None
301
- assert self.input_quant is not None
302
-
303
- is_8_bits = self.weight_quant.num_bits == self.input_quant.num_bits == 8
304
-
305
- if not is_8_bits:
306
- raise ValueError("Cutlass only supports 8-bit quantization")
307
-
308
- if (
309
- self.weight_quant.type == QuantizationType.FLOAT
310
- and self.input_quant.type == QuantizationType.FLOAT
311
- ):
312
- return torch.float8_e4m3fn
313
-
314
- if (
315
- self.weight_quant.type == QuantizationType.INT
316
- and self.input_quant.type == QuantizationType.INT
317
- ):
318
- return torch.int8
319
-
320
- raise ValueError("Quantization type not supported by Cutlass")
321
-
322
- def _decompress_bitmask_compressed_weight(
323
- self,
324
- compressed: torch.Tensor,
325
- bitmask: torch.Tensor,
326
- layer: torch.nn.Module,
327
- ) -> torch.Tensor:
328
- """
329
- Decompress a compressed 2:4 sparse weight tensor using the bitmask and
330
- return the result.
331
-
332
- This function also supports sharded decompression.
333
-
334
- :param compressed: The 2:4 sparse weight tensor compressed using the
335
- sparse-24-bitmask compressor. This is different from
336
- `cutlass_sparse_compress` which uses a different scheme (2 bits for
337
- every nonzero element that represent the coordinate within the block
338
- of 4). The bitmask compression here uses a bitmask to indicate the
339
- positions of non-zero elements.
340
- :param bitmask: The 2:4 bitmask associated with the compressed weights,
341
- representing the positions of non-zero elements in the compressed
342
- tensor.
343
- :param layer: The layer whose weights need to be processed after
344
- loading.
345
- :return: The decompressed 2:4 sparse weight tensor.
346
- """
347
-
348
- sparsity_compressor = self.model_compressor.sparsity_compressor
349
-
350
- def _process_split(
351
- bitmask_compressed_weight: torch.Tensor,
352
- shape,
353
- bitmask: torch.Tensor,
354
- ) -> torch.Tensor:
355
- weight_data = dict(
356
- compressed=bitmask_compressed_weight,
357
- shape=shape,
358
- bitmask=bitmask,
359
- )
360
- return sparsity_compressor.decompress_weight(weight_data)
361
-
362
- split_weights: list[torch.Tensor] = []
363
- split_bitmask: list[torch.Tensor] = []
364
- split_shape: list[tuple[int, int]] = []
365
-
366
- if isinstance(layer, (QKVParallelLinear, MergedColumnParallelLinear)):
367
- split_weights = torch.split(compressed, layer.logical_widths)
368
- split_bitmask = torch.split(bitmask, layer.logical_widths)
369
- split_shape = [
370
- (out, layer.input_size_per_partition) for out in layer.logical_widths
371
- ]
372
-
373
- if split_weights:
374
- decompressed_shards = [
375
- _process_split(compressed_weight, shape, bitmask)
376
- for compressed_weight, shape, bitmask in zip(
377
- split_weights, split_shape, split_bitmask
378
- )
379
- ]
380
- decompressed = combine_shards(decompressed_shards)
381
- else:
382
- decompressed = sparsity_compressor.decompress_weight(
383
- dict(
384
- compressed=compressed,
385
- shape=(
386
- layer.logical_widths[0],
387
- layer.input_size_per_partition,
388
- ),
389
- bitmask=bitmask,
390
- )
391
- )
392
- return decompressed
54
+ raise NotImplementedError("Sparse24 models are no longer supported by vLLM")
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_scheme.py
+4
-4
@@ -20,7 +20,7 @@ class CompressedTensorsScheme(ABC):
20
"""
21
Get minimum device capability.
22
"""
23
- raise NotImplementedError
23
+ raise NotImplementedError()
24
25
@abstractmethod
26
def create_weights(self, *args, **kwargs):
@@ -28,7 +28,7 @@ class CompressedTensorsScheme(ABC):
28
Weight creation for the particular scheme. Inputs to this function
29
30
"""
31
- raise NotImplementedError
31
+ raise NotImplementedError()
32
33
@abstractmethod
34
def apply_weights(
@@ -44,7 +44,7 @@ class CompressedTensorsScheme(ABC):
44
:param bias: bias parameter
45
46
"""
47
- raise NotImplementedError
47
+ raise NotImplementedError()
48
49
@abstractmethod
50
def process_weights_after_loading(self, layer: torch.nn.Module):
@@ -52,4 +52,4 @@ class CompressedTensorsScheme(ABC):
52
Called after weight loading is complete for any cleanup that
53
needs to occur.
54
"""
55
- raise NotImplementedError
55
+ raise NotImplementedError()
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
-10
@@ -8,16 +8,6 @@ from vllm import _custom_ops as ops
8
from vllm.platforms import current_platform
9
10
11
-def sparse_cutlass_supported() -> bool:
12
- if not current_platform.is_cuda():
13
- return False
14
-
15
- capability_tuple = current_platform.get_device_capability()
16
- capability = -1 if capability_tuple is None else capability_tuple.to_int()
17
-
18
- return ops.cutlass_sparse_scaled_mm_supported(capability)
19
-
20
-
11
def cutlass_fp8_supported() -> bool:
12
if not current_platform.is_cuda():
13
return False