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

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 3, 16:30 UTC

[Quantization][Autoround][Toolkit] Add W4A16 Support (#39778)

Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com> Signed-off-by: Zhenzhong Xu <zhenzhong.xu@intel.com>

Zhenzhong Xu committed May 14, 2026 at 19:18 UTC 1ea94013644558f83e134f0c6f53d6b261f7c824

2 files changed +209 -31

requirements/xpu.txt

@@ -15,4 +15,5 @@ torch==2.11.0+xpu

15 torchaudio

16 torchvision

18 +auto_round_lib>=0.13.0

19 vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.8/vllm_xpu_kernels-0.1.8-cp38-abi3-manylinux_2_28_x86_64.whl

vllm/model_executor/layers/quantization/inc.py

+208 -31

@@ -2,6 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from fractions import Fraction

5 +from functools import lru_cache

6 from typing import TYPE_CHECKING, Any

8 import regex as re

@@ -429,7 +430,23 @@ class INCConfig(QuantizationConfig):

430 raise NotImplementedError(

431 "INC W4A16 on XPU only supports symmetric quantization for now."

432 )

433 +

434 if isinstance(layer, (LinearBase, ParallelLMHead)):

435 + is_ark_available, ark_error, _, _ = get_ark_state()

436 + if is_ark_available:

437 + return INCARKLinearMethod(

438 + weight_bits=weight_bits,

439 + group_size=group_size,

440 + sym=sym,

441 + )

442 +

443 + logger.debug(

444 + "ARK backend is unavailable for layer %s; "

445 + "falling back to the default XPU INC path. Error: %s",

446 + prefix,

447 + ark_error or "unknown error",

448 + )

449 +

450 return INCXPULinearMethod(

451 weight_bits=weight_bits,

452 group_size=group_size,

@@ -455,6 +472,21 @@ class INCConfig(QuantizationConfig):

472 "INC W4A16 on CPU only supports symmetric quantization for now."

473 )

474 if isinstance(layer, (LinearBase, ParallelLMHead)):

475 + is_ark_available, ark_error, _, _ = get_ark_state()

476 + if is_ark_available:

477 + return INCARKLinearMethod(

478 + weight_bits=weight_bits,

479 + group_size=group_size,

480 + sym=sym,

481 + )

482 +

483 + logger.debug(

484 + "ARK backend is unavailable for layer %s; "

485 + "falling back to the default CPU INC path. Error: %s",

486 + prefix,

487 + ark_error or "unknown error",

488 + )

489 +

490 return self.apply_gptq_quant_layer(layer, prefix)

491 return None

492

@@ -465,6 +497,7 @@ class INCConfig(QuantizationConfig):

497 layer_name == prefix or layer_name == f"model.{prefix}"

498 ) and self.extra_config[layer_name].get("bits", 16) >= 16:

499 return UnquantizedLinearMethod()

500 +

501 if current_platform.is_xpu():

502 return self.apply_xpu_w4a16_quant_layer(layer, prefix)

503 is_gptq = "gptq" in self.packing_format or "gptq" in self.backend

@@ -493,54 +526,39 @@ class INCConfig(QuantizationConfig):

526 return None

527

528

496 -class INCXPULinearMethod(LinearMethodBase):

497 - """XPU linear method for INC w4a16 GPTQ quantization (symmetric only).

498 -

499 - Repacks GPTQ weights from [in_packed, out] to oneDNN [out, in_packed]

500 - layout and calls torch.ops._xpu_C.int4_gemm_w4a16.

501 -

502 - GPTQ format: qweight [in_packed, out] with sequential nibble order.

503 -

504 - Note: Asymmetric quantization (sym=false) is not for now.

505 -

506 - FIXME(yiliu30): Refine the implementation to reuse XPUwNa16LinearKernel.

507 - """

508 -

529 +class INCXPULinearBase(LinearMethodBase):

530 def __init__(self, weight_bits: int, group_size: int, sym: bool):

531 self.weight_bits = weight_bits

532 self.group_size = group_size

533 self.sym = sym

534 self.pack_factor = 32 // weight_bits

535

515 - def create_weights(

536 + def _create_inc_weights(

537 self,

538 layer: torch.nn.Module,

539 input_size_per_partition: int,

540 output_partition_sizes: list[int],

520 - input_size: int,

521 - output_size: int,

541 params_dtype: torch.dtype,

523 - **extra_weight_attrs,

524 - ):

525 - del output_size # Unused.

542 + weight_loader: Any,

543 + group_size: int,

544 + pack_factor: int,

545 + ) -> None:

546 output_size_per_partition = sum(output_partition_sizes)

527 - weight_loader = extra_weight_attrs.get("weight_loader")

528 - scales_and_zp_size = input_size_per_partition // self.group_size

547 + scales_and_zp_size = input_size_per_partition // group_size

548

530 - # GPTQ: qweight [in // pack_factor, out] packed along input dim

549 qweight = PackedvLLMParameter(

550 data=torch.empty(

533 - input_size_per_partition // self.pack_factor,

551 + input_size_per_partition // pack_factor,

552 output_size_per_partition,

553 dtype=torch.int32,

554 ),

555 input_dim=0,

556 output_dim=1,

557 packed_dim=0,

540 - packed_factor=self.pack_factor,

558 + packed_factor=pack_factor,

559 weight_loader=weight_loader,

560 )

543 - # scales: [num_groups, out] params_dtype

561 +

562 scales = GroupQuantScaleParameter(

563 data=torch.empty(

564 scales_and_zp_size,

@@ -551,17 +569,17 @@ class INCXPULinearMethod(LinearMethodBase):

569 output_dim=1,

570 weight_loader=weight_loader,

571 )

554 - # qzeros: [num_groups, out // pack_factor] int32

572 +

573 qzeros = PackedvLLMParameter(

574 data=torch.empty(

575 scales_and_zp_size,

558 - output_size_per_partition // self.pack_factor,

576 + output_size_per_partition // pack_factor,

577 dtype=torch.int32,

578 ),

579 input_dim=0,

580 output_dim=1,

581 packed_dim=1,

564 - packed_factor=self.pack_factor,

582 + packed_factor=pack_factor,

583 weight_loader=weight_loader,

584 )

585

@@ -569,11 +587,9 @@ class INCXPULinearMethod(LinearMethodBase):

587 layer.register_parameter("scales", scales)

588 layer.register_parameter("qzeros", qzeros)

589

572 - # GPTQ checkpoints may include g_idx for activation reordering.

573 - # Register it so the weight loader doesn't error on unexpected keys.

590 g_idx = RowvLLMParameter(

591 data=torch.tensor(

576 - [i // self.group_size for i in range(input_size_per_partition)],

592 + [i // group_size for i in range(input_size_per_partition)],

593 dtype=torch.int32,

594 ),

595 input_dim=0,

@@ -581,6 +597,66 @@ class INCXPULinearMethod(LinearMethodBase):

597 )

598 layer.register_parameter("g_idx", g_idx)

599

600 + def create_weights(

601 + self,

602 + layer: torch.nn.Module,

603 + input_size_per_partition: int,

604 + output_partition_sizes: list[int],

605 + input_size: int,

606 + output_size: int,

607 + params_dtype: torch.dtype,

608 + **extra_weight_attrs,

609 + ):

610 + self._create_inc_weights(

611 + layer=layer,

612 + input_size_per_partition=input_size_per_partition,

613 + output_partition_sizes=output_partition_sizes,

614 + params_dtype=params_dtype,

615 + weight_loader=extra_weight_attrs.get("weight_loader"),

616 + group_size=self.group_size,

617 + pack_factor=self.pack_factor,

618 + )

619 +

620 +

621 +@lru_cache(maxsize=1)

622 +def get_ark_state() -> tuple[bool, str | None, Any | None, Any | None]:

623 + """Return ARK availability, error details, cached instance, and QuantLinear."""

624 + try:

625 + import auto_round_kernel

626 + from auto_round_kernel.qlinear import QuantLinear

627 +

628 + logger.info("Successfully imported auto_round_kernel.")

629 + except ImportError as error:

630 + return False, str(error), None, None

631 +

632 + ark_loader = getattr(auto_round_kernel, "_ark_instance", None)

633 + if not callable(ark_loader):

634 + return False, "auto_round_kernel does not expose _ark_instance().", None, None

635 +

636 + try:

637 + ark_instance = ark_loader()

638 + except Exception as error:

639 + return False, str(error), None, None

640 +

641 + if ark_instance is None:

642 + return False, "auto_round_kernel._ark_instance() returned None.", None, None

643 +

644 + return True, None, ark_instance, QuantLinear

645 +

646 +

647 +class INCXPULinearMethod(INCXPULinearBase):

648 + """XPU linear method for INC w4a16 GPTQ quantization (symmetric only).

649 +

650 + Repacks GPTQ weights from [in_packed, out] to oneDNN [out, in_packed]

651 + layout and calls torch.ops._xpu_C.int4_gemm_w4a16.

652 +

653 + GPTQ format: qweight [in_packed, out] with sequential nibble order.

654 +

655 + Note: Asymmetric quantization (sym=false) is not for now.

656 +

657 + FIXME(yiliu30): Refine the implementation to reuse XPUwNa16LinearKernel.

658 + """

659 +

660 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

661 """Repack GPTQ weights into kernel-ready NT layout."""

662 device = layer.qweight.data.device

@@ -624,3 +700,104 @@ class INCXPULinearMethod(LinearMethodBase):

700 None, # g_idx not needed: desc_act is always False for INC models

701 )

702 return out.reshape(out_shape)

703 +

704 +

705 +class INCARKLinearMethod(INCXPULinearBase):

706 + """XPU & CPU w4a16 linear method for INC quantization utilizing the ARK backend.

707 +

708 + See: https://github.com/intel/auto-round/blob/main/auto_round_extension/ark/README.md

709 +

710 + Repacks GPTQ/INC weights into ARK's layout.

711 + """

712 +

713 + def __init__(self, weight_bits: int, group_size: int, sym: bool):

714 + super().__init__(weight_bits=weight_bits, group_size=group_size, sym=sym)

715 +

716 + is_available, error_str, _, quant_linear_cls = get_ark_state()

717 + if not is_available or quant_linear_cls is None:

718 + reason = error_str or "unknown error"

719 + raise ImportError(f"Failed to import auto_round_kernel. {reason}")

720 +

721 + self.QuantLinear = quant_linear_cls

722 +

723 + def create_weights(

724 + self,

725 + layer: torch.nn.Module,

726 + input_size_per_partition: int,

727 + output_partition_sizes: list[int],

728 + input_size: int,

729 + output_size: int,

730 + params_dtype: torch.dtype,

731 + **extra_weight_attrs,

732 + ):

733 + super().create_weights(

734 + layer=layer,

735 + input_size_per_partition=input_size_per_partition,

736 + output_partition_sizes=output_partition_sizes,

737 + input_size=input_size,

738 + output_size=output_size,

739 + params_dtype=params_dtype,

740 + **extra_weight_attrs,

741 + )

742 + layer.in_features = input_size_per_partition

743 + layer.out_features = sum(output_partition_sizes)

744 + layer.params_dtype = params_dtype

745 +

746 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:

747 + if hasattr(layer, "input_size_per_partition"):

748 + in_features = layer.input_size_per_partition

749 + elif hasattr(layer, "input_size"):

750 + in_features = layer.input_size

751 + else:

752 + raise AttributeError("Cannot determine in_features for layer.")

753 +

754 + if hasattr(layer, "output_partition_sizes"):

755 + out_features = sum(layer.output_partition_sizes)

756 + elif hasattr(layer, "output_size_per_partition"):

757 + out_features = layer.output_size_per_partition

758 + elif hasattr(layer, "output_size"):

759 + out_features = layer.output_size

760 + else:

761 + out_features = layer.scales.shape[-1]

762 +

763 + ark_linear = self.QuantLinear(

764 + bits=self.weight_bits,

765 + group_size=self.group_size,

766 + sym=self.sym,

767 + in_features=in_features,

768 + out_features=out_features,

769 + bias=layer.bias is not None,

770 + weight_dtype=layer.params_dtype,

771 + )

772 +

773 + ark_linear.to(layer.qweight.device)

774 +

775 + with torch.no_grad():

776 + ark_linear.qweight.copy_(layer.qweight.detach())

777 +

778 + if hasattr(layer, "qzeros") and layer.qzeros is not None:

779 + ark_linear.qzeros.copy_(layer.qzeros.detach())

780 + else:

781 + ark_linear.qzeros = None

782 +

783 + ark_linear.scales.copy_(layer.scales.detach())

784 +

785 + if hasattr(layer, "bias") and layer.bias is not None:

786 + ark_linear.bias.copy_(layer.bias.detach())

787 +

788 + ark_linear.post_init()

789 +

790 + layer.ark_linear = ark_linear

791 +

792 + del layer.qweight

793 + if hasattr(layer, "qzeros"):

794 + del layer.qzeros

795 + del layer.scales

796 +

797 + def apply(

798 + self,

799 + layer: torch.nn.Module,

800 + x: torch.Tensor,

801 + bias: torch.Tensor | None = None,

802 + ) -> torch.Tensor:

803 + return layer.ark_linear.forward(x)