@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 1f1b1bc03b

Fork 0

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 15, 08:30 UTC

[V1][Quantization] Add CUDA graph compatible v1 GGUF support (#18646)

Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Isotr0py <2037008807@qq.com>

Isotr0py committed May 27, 2025 at 12:40 UTC 1f1b1bc03bfa36dbd0d884cede6f4dcdf0937e25

5 files changed +188 -59

tests/kernels/quantization/test_gguf.py

+1 -3

@@ -8,7 +8,6 @@ from gguf import GGMLQuantizationType, GGUFReader, ReaderTensor, dequantize

8 from huggingface_hub import snapshot_download

10 import vllm._custom_ops as ops

11 -from vllm.model_executor.layers.activation import SiluAndMul

11 from vllm.model_executor.layers.fused_moe import fused_experts

12 from vllm.model_executor.layers.quantization.gguf import _fused_moe_gguf

13 from vllm.platforms import current_platform

@@ -176,12 +175,11 @@ def test_moe(num_tokens: int, hidden_size: int, dtype: torch.dtype,

175

176 w2_dequant = torch.tensor(dequantize(w2.data, quant_type),

177 device="cuda").to(dtype)

179 - act = SiluAndMul()

178

179 output = _fused_moe_gguf(x, torch.tensor(w13.data, device="cuda"),

180 torch.tensor(w2.data,

181 device="cuda"), topk_weights,

184 - topk_ids, quant_type, quant_type, act)

182 + topk_ids, quant_type, quant_type, "silu")

183

184 ref_output = fused_experts(x, w13_dequant, w2_dequant, topk_weights,

185 topk_ids).reshape(output.shape)

tests/models/quantization/test_gguf.py

+6 -2

@@ -78,8 +78,12 @@ DOLPHIN_CONFIG = GGUFTestConfig(

78 )

80 MODELS = [

81 - LLAMA_CONFIG, QWEN2_CONFIG, PHI3_CONFIG, GPT2_CONFIG, STABLELM_CONFIG,

82 - DOLPHIN_CONFIG

81 + LLAMA_CONFIG,

82 + QWEN2_CONFIG,

83 + PHI3_CONFIG,

84 + GPT2_CONFIG,

85 + # STABLELM_CONFIG, # enable this when v1 support head_size=80

86 + DOLPHIN_CONFIG,

87 # STARCODER_CONFIG, # broken

88 ]

vllm/engine/arg_utils.py

-8

@@ -1291,14 +1291,6 @@ class EngineArgs:

1291 recommend_to_remove=False)

1292 return False

1293

1294 - # Some quantization is not compatible with torch.compile.

1295 - V1_UNSUPPORTED_QUANT = ["gguf"]

1296 - if model_config.quantization in V1_UNSUPPORTED_QUANT:

1297 - _raise_or_fallback(

1298 - feature_name=f"--quantization {model_config.quantization}",

1299 - recommend_to_remove=False)

1300 - return False

1301 -

1294 # No Embedding Models so far.

1295 if model_config.task not in ["generate"]:

1296 _raise_or_fallback(feature_name=f"--task {model_config.task}",

vllm/model_executor/layers/linear.py

-4

@@ -587,8 +587,6 @@ class MergedColumnParallelLinear(ColumnParallelLinear):

587 param.shard_id.append(loaded_shard_id)

588 param.shard_id_map[loaded_shard_id] = len(param.data_container)

589 param.data_container.append(loaded_weight)

590 - if len(param.data_container) == 2:

591 - self.qweight = param.materialize_nested()

590 return

591

592 param_data = param.data

@@ -982,8 +980,6 @@ class QKVParallelLinear(ColumnParallelLinear):

980 param.shard_id.append(loaded_shard_id)

981 param.shard_id_map[loaded_shard_id] = len(param.data_container)

982 param.data_container.append(loaded_weight)

985 - if len(param.data_container) == 3:

986 - self.qweight = param.materialize_nested()

983 return

984

985 param_data = param.data

vllm/model_executor/layers/quantization/gguf.py

+181 -42

@@ -9,7 +9,6 @@ from torch.nn.parameter import Parameter, UninitializedParameter

10 from vllm import _custom_ops as ops

11 from vllm.logger import init_logger

12 -from vllm.model_executor.layers.activation import SiluAndMul

12 from vllm.model_executor.layers.fused_moe.layer import (FusedMoE,

13 FusedMoEMethodBase)

14 from vllm.model_executor.layers.linear import LinearBase, LinearMethodBase

@@ -19,6 +18,7 @@ from vllm.model_executor.layers.quantization.base_config import (

18 from vllm.model_executor.layers.vocab_parallel_embedding import (

19 VocabParallelEmbedding)

20 from vllm.model_executor.utils import set_weight_attrs

21 +from vllm.utils import direct_register_custom_op

23 logger = init_logger(__name__)

@@ -96,8 +96,8 @@ MMVQ_QUANT_TYPES = STANDARD_QUANT_TYPES | KQUANT_TYPES | IMATRIX_QUANT_TYPES

96 MMQ_QUANT_TYPES = STANDARD_QUANT_TYPES | KQUANT_TYPES

99 -def _fuse_mul_mat(x: torch.Tensor, qweight: torch.Tensor,

100 - qweight_type: int) -> torch.Tensor:

99 +def _fused_mul_mat_gguf(x: torch.Tensor, qweight: torch.Tensor,

100 + qweight_type: int) -> torch.Tensor:

101 # HACK: when doing chunked prefill we don't generate output tokens

102 # so input to logits generator is empty which causes invalid parameter

103 if x.shape[0] == 0:

@@ -130,6 +130,30 @@ def _fuse_mul_mat(x: torch.Tensor, qweight: torch.Tensor,

130 return y

131

132

133 +def _fused_mul_mat_gguf_fake(

134 + x: torch.Tensor,

135 + qweight: torch.Tensor,

136 + qweight_type: int,

137 +) -> torch.Tensor:

138 + return torch.empty(x.shape[0],

139 + qweight.shape[0],

140 + dtype=x.dtype,

141 + device=x.device)

142 +

143 +

144 +try:

145 + direct_register_custom_op(

146 + op_name="_fused_mul_mat_gguf",

147 + op_func=_fused_mul_mat_gguf,

148 + mutates_args=[],

149 + fake_impl=_fused_mul_mat_gguf_fake,

150 + )

151 + fused_mul_mat_gguf = torch.ops.vllm._fused_mul_mat_gguf

152 +

153 +except AttributeError as error:

154 + raise error

155 +

156 +

157 def _fused_moe_gguf(

158 x: torch.Tensor,

159 w1: torch.Tensor,

@@ -138,8 +162,21 @@ def _fused_moe_gguf(

162 topk_ids: torch.Tensor,

163 qweight_type: int,

164 qweight_type2: int,

141 - act,

165 + activation: str,

166 ) -> torch.Tensor:

167 +

168 + def act(x: torch.Tensor):

169 + d = x.shape[-1] // 2

170 + output_shape = (x.shape[:-1] + (d, ))

171 + out = torch.empty(output_shape, dtype=x.dtype, device=x.device)

172 + if activation == "silu":

173 + torch.ops._C.silu_and_mul(out, x)

174 + elif activation == "gelu":

175 + torch.ops._C.gelu_and_mul(out, x)

176 + else:

177 + raise ValueError(f"Unsupported activation: {activation}")

178 + return out

179 +

180 # lazy import to avoid triggering triton import in CPU backend

181 from vllm.model_executor.layers.fused_moe.fused_moe import (

182 moe_align_block_size)

@@ -189,12 +226,12 @@ def _fused_moe_gguf(

226 for ww, ii in zip(w, idx):

227 expert_up = w1[ii]

228

192 - out = _fuse_mul_mat(inp, expert_up, qweight_type)

229 + out = fused_mul_mat_gguf(inp, expert_up, qweight_type)

230 out = act(out)

231

232 expert_down = w2[ii]

196 - current_state = _fuse_mul_mat(out, expert_down,

197 - qweight_type2).mul_(ww)

233 + current_state = fused_mul_mat_gguf(out, expert_down,

234 + qweight_type2).mul_(ww)

235 if current_hidden_state is None:

236 current_hidden_state = current_state

237 else:

@@ -203,6 +240,78 @@ def _fused_moe_gguf(

240 return out_hidden_states

241

242

243 +def _fused_moe_gguf_fake(

244 + x: torch.Tensor,

245 + w1: torch.Tensor,

246 + w2: torch.Tensor,

247 + topk_weights: torch.Tensor,

248 + topk_ids: torch.Tensor,

249 + qweight_type: int,

250 + qweight_type2: int,

251 + activation: str,

252 +) -> torch.Tensor:

253 + return torch.empty_like(x)

254 +

255 +

256 +try:

257 + direct_register_custom_op(

258 + op_name="_fused_moe_gguf",

259 + op_func=_fused_moe_gguf,

260 + mutates_args=[],

261 + fake_impl=_fused_moe_gguf_fake,

262 + )

263 + fused_moe_gguf = torch.ops.vllm._fused_moe_gguf

264 +

265 +except AttributeError as error:

266 + raise error

267 +

268 +

269 +def _apply_gguf_embedding(

270 + x: torch.Tensor,

271 + qweight: torch.Tensor,

272 + qweight_type: int,

273 + hidden_size: int,

274 + dtype: Optional[torch.dtype] = None,

275 +) -> torch.Tensor:

276 + if qweight_type in UNQUANTIZED_TYPES:

277 + return torch.embedding(qweight, x)

278 + elif qweight_type in DEQUANT_TYPES:

279 + block_size, type_size = gguf.GGML_QUANT_SIZES[qweight_type]

280 + x_flat = x.flatten()

281 + assert (hidden_size == qweight.shape[1] // type_size * block_size)

282 + quant = torch.index_select(qweight, dim=0, index=x_flat)

283 + dequant = ops.ggml_dequantize(quant, qweight_type, hidden_size,

284 + x_flat.shape[0], dtype)

285 + return dequant.view(*x.shape, hidden_size)

286 + else:

287 + qweight_type = WeightType(qweight_type)

288 + raise NotImplementedError(

289 + f"Unsupported GGUF quantization type: {qweight_type}")

290 +

291 +

292 +def _apply_gguf_embedding_fake(

293 + x: torch.Tensor,

294 + qweight: torch.Tensor,

295 + qweight_type: int,

296 + hidden_size: int,

297 + dtype: Optional[torch.dtype] = None,

298 +) -> torch.Tensor:

299 + return torch.empty(x.shape[0], hidden_size, dtype=dtype, device=x.device)

300 +

301 +

302 +try:

303 + direct_register_custom_op(

304 + op_name="_apply_gguf_embedding",

305 + op_func=_apply_gguf_embedding,

306 + mutates_args=[],

307 + fake_impl=_apply_gguf_embedding_fake,

308 + )

309 + apply_gguf_embedding = torch.ops.vllm._apply_gguf_embedding

310 +

311 +except AttributeError as error:

312 + raise error

313 +

314 +

315 class GGUFLinearMethod(LinearMethodBase):

316 """Linear method for GGUF.

317

@@ -249,26 +358,76 @@ class GGUFLinearMethod(LinearMethodBase):

358 set_weight_attrs(qweight_type, extra_weight_attrs)

359 layer.register_parameter("qweight_type", qweight_type)

360

361 + def process_weights_after_loading(self, layer: torch.nn.Module):

362 + qweight_type = layer.qweight_type.weight_type

363 + if not (qweight_type in UNQUANTIZED_TYPES

364 + or qweight_type in DEQUANT_TYPES):

365 + qweight_type = WeightType(qweight_type)

366 + raise ValueError(

367 + f"Unsupported GGUF quantization type {qweight_type} in "

368 + f"layer {layer}.")

369 + # For MergedColumnParallelLinear and QKVParallelLinear, we need to

370 + # materialize the padded weight parameter for CUDA Graph compatibility.

371 + self._create_padded_weight_param(layer)

372 +

373 + def _create_padded_weight_param(self, layer: torch.nn.Module):

374 + """Create padded weight parameter for GGUF MergedLinear layer."""

375 + qweight = layer.qweight

376 + shard_id_map = qweight.shard_id_map

377 + shard_id = qweight.shard_id

378 + if len(data_container := qweight.data_container) > 1:

379 + dtype = {data.dtype for data in data_container}

380 + assert len(dtype) == 1, ValueError(

381 + f"Data container has mixed dtypes: {dtype}")

382 + dtype = next(iter(dtype))

383 + # concat dim0 and pad dim1

384 + padded_side = max(x.size(1) for x in data_container)

385 + concat_side = sum(x.size(0) for x in data_container)

386 + # Pad the quantized weights to dense tensor, and create a map

387 + # with the location of each shard in the padded tensor.

388 + padded_data = torch.zeros((concat_side, padded_side),

389 + dtype=dtype,

390 + device=qweight.device)

391 + # (dim0_start, dim0_end, dim1_size)

392 + shard_offset_map = dict[str, tuple[int, int, int]]()

393 + for idx in shard_id:

394 + id_in_container = shard_id_map[idx]

395 + start = sum(

396 + x.size(0) for x in data_container[:id_in_container])

397 + end = start + data_container[id_in_container].size(0)

398 + size = data_container[id_in_container].size(1)

399 + padded_data[start:end, :size] = data_container[id_in_container]

400 + shard_offset_map[idx] = (start, end, size)

401 + qweight.data_container.clear()

402 + padded_param = Parameter(padded_data, requires_grad=False)

403 + set_weight_attrs(padded_param, vars(qweight))

404 + set_weight_attrs(padded_param,

405 + {"shard_offset_map": shard_offset_map})

406 + layer.register_parameter("qweight", padded_param)

407 +

408 def apply(self,

409 layer: torch.nn.Module,

410 x: torch.Tensor,

411 bias: Optional[torch.Tensor] = None) -> torch.Tensor:

256 - shard_id = getattr(layer.qweight, "shard_id", None)

412 + shard_id = layer.qweight.shard_id

413

414 if shard_id:

415 # dequantize shard weights respectively

416 shard_id = ["q", "k", "v"] if "q" in shard_id else shard_id

261 - qweight = layer.qweight.unbind(0)

417 + qweight = layer.qweight

418 result = []

419 for idx in shard_id:

264 - q_idx = layer.qweight.shard_id_map[idx]

420 + start, end, offset = layer.qweight.shard_offset_map[idx]

421 qweight_type = layer.qweight_type.shard_weight_type[idx]

266 - result.append(_fuse_mul_mat(x, qweight[q_idx], qweight_type))

422 + result.append(

423 + fused_mul_mat_gguf(

424 + x, qweight[start:end, :offset].contiguous(),

425 + qweight_type))

426 out = torch.cat(result, axis=1)

427 else:

428 qweight = layer.qweight

429 qweight_type = layer.qweight_type.weight_type

271 - out = _fuse_mul_mat(x, qweight, qweight_type)

430 + out = fused_mul_mat_gguf(x, qweight, qweight_type)

431 if bias is not None:

432 out.add_(bias)

433 return out

@@ -338,7 +497,6 @@ class GGUFMoEMethod(FusedMoEMethodBase):

497

498 set_weight_attrs(w2_qweight_type, extra_weight_attrs)

499 layer.register_parameter("w2_qweight_type", w2_qweight_type)

341 - self.act = SiluAndMul()

500

501 def apply(

502 self,

@@ -375,10 +533,10 @@ class GGUFMoEMethod(FusedMoEMethodBase):

533 custom_routing_function=custom_routing_function,

534 scoring_func=scoring_func,

535 e_score_correction_bias=e_score_correction_bias)

378 - return _fused_moe_gguf(x, layer.w13_qweight, layer.w2_qweight,

379 - topk_weights, topk_ids,

380 - layer.w13_qweight_type.weight_type,

381 - layer.w2_qweight_type.weight_type, self.act)

536 + return fused_moe_gguf(x, layer.w13_qweight, layer.w2_qweight,

537 + topk_weights, topk_ids,

538 + layer.w13_qweight_type.weight_type,

539 + layer.w2_qweight_type.weight_type, activation)

540

541

542 class GGUFEmbeddingMethod(GGUFLinearMethod):

@@ -392,34 +550,15 @@ class GGUFEmbeddingMethod(GGUFLinearMethod):

550 x: torch.Tensor) -> torch.Tensor:

551 qweight = layer.qweight

552 qweight_type = layer.qweight_type.weight_type

553 + hidden_size = qweight.tensor_shape[1]

554

396 - block_size, type_size = gguf.GGML_QUANT_SIZES[qweight_type]

397 - hidden_size = qweight.shape[1] // type_size * block_size

398 - if qweight_type < 2:

399 - return torch.embedding(qweight, x)

400 - x_flat = x.flatten()

401 - quant = torch.index_select(qweight, dim=0, index=x_flat)

402 - dequant = ops.ggml_dequantize(quant, qweight_type, hidden_size,

403 - x_flat.shape[0], self.params_dtype)

404 - return dequant.view(*x.shape, hidden_size)

555 + return apply_gguf_embedding(x,

556 + qweight,

557 + qweight_type,

558 + hidden_size,

559 + dtype=self.params_dtype)

560

561

562 class GGUFUninitializedParameter(UninitializedParameter):

563 cls_to_become = Parameter

564 data_container: list[torch.Tensor]

410 -

411 - def materialize_nested(self) -> Parameter:

412 - dtype = {data.dtype for data in self.data_container}

413 - assert len(dtype) == 1, ValueError(

414 - f"Data container has mixed dtypes: {dtype}")

415 - dtype = next(iter(dtype))

416 - nested_data = torch.nested.nested_tensor(self.data_container,

417 - device=self.device,

418 - dtype=dtype)

419 - self.data_container.clear()

420 - param = torch.Tensor._make_subclass(self.cls_to_become,

421 - nested_data,

422 - require_grad=False)

423 - for k, v in self.__dict__.items():

424 - setattr(param, k, v)

425 - return param