@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9fd750f00d

[XPU][INC] Add int4 w4a8 (dynamic int8 activation) backend for INC linear layers (#50501)

Signed-off-by: Thanaji Rao Thakkalapelli <thanaji.rao.thakkalapelli@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>

Thanaji Rao Thakkalapelli committed Aug 21, 2026 at 01:38 UTC 9fd750f00de5ab8fa172c34e219b3345411e9a91
4 files changed +787
tests/quantization/test_auto_round.py
+618
@@ -33,6 +33,9 @@ from vllm.model_executor.layers.quantization.inc.schemes.inc_mxfp8_linear import
33 from vllm.model_executor.layers.quantization.inc.schemes.inc_scheme import (
34 INCLinearScheme,
35 )
36 +from vllm.model_executor.layers.quantization.inc.schemes.inc_w4a8_linear import (
37 + INCXPUW4A8LinearMethod,
38 +)
39 from vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear import (
40 INCARKLinearMethod,
41 INCWNA16LinearScheme,
@@ -1482,3 +1485,618 @@ class TestGetLayerConfigFusedQKV:
1485 DummyLayer(), "model.layers.0.self_attn.qkv_proj"
1486 )
1487 assert bits == 8
1488 +
1489 +
1490 +# ---------------------------------------------------------------------------
1491 +# INC XPU w4a8: int4 weights with dynamic per-token int8 activations
1492 +#
1493 +# The w4a8 path keeps int4 weights but dynamically quantizes activations to
1494 +# per-token symmetric int8, so the GEMM stays on the int8 datapath instead of
1495 +# upconverting weights to the activation dtype. It is opt-in through
1496 +# VLLM_XPU_INC_WNA16_BACKEND because which kernel is fastest depends on the
1497 +# device: ARK can use XMX int8 on some XPUs, and the default ("auto")
1498 +# preference order is deliberately left unchanged.
1499 +#
1500 +# The backend-selection tests stub out ARK and the platform; the weight-creation
1501 +# tests only need the TP-rank stubs. Both run anywhere. The apply_weights tests
1502 +# stub out int4_gemm_w4a8 — what needs guarding there is the calling convention
1503 +# (scale dtypes, argument order, weight reuse, batch-size gate), which is where
1504 +# this path is easy to get silently wrong, since the kernel accepts bf16 scales
1505 +# and returns plausible-looking garbage rather than raising. The ones that cross
1506 +# the token threshold reach ``vllm._xpu_ops`` for the activation quantizer, so
1507 +# they are XPU-image only; the small-batch fallback test returns before that
1508 +# import and must stay unmarked to keep covering the gate off-XPU. The
1509 +# end-to-end tests load a real int4 checkpoint and additionally need the op.
1510 +# ---------------------------------------------------------------------------
1511 +
1512 +_BACKEND_ENV = "VLLM_XPU_INC_WNA16_BACKEND"
1513 +_ARK_STATE = (
1514 + "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state"
1515 +)
1516 +_QUANT_REF = "vllm._xpu_ops.xpu_ops.dynamic_per_token_int8_quant_ref"
1517 +
1518 +IN_FEATURES = 128
1519 +OUT_FEATURES = 64
1520 +
1521 +# auto-round int4/sym/g128 checkpoints, which is what this backend serves. The
1522 +# GPTQ- and AWQ-packed variants take different repack branches in
1523 +# ``process_weights_after_loading``, so both are worth loading.
1524 +E2E_MODELS = [
1525 + pytest.param("OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc", id="auto_round:auto_gptq"),
1526 + pytest.param(
1527 + "Intel/Qwen2-0.5B-Instruct-int4-sym-AutoRound", id="auto_round:auto_awq"
1528 + ),
1529 +]
1530 +
1531 +
1532 +def _has_xpu_ops() -> bool:
1533 + """True when ``vllm._xpu_ops`` is importable.
1534 +
1535 + ``apply_weights`` imports it for the activation quantizer, and the stubs
1536 + below patch through it, so tests that exercise that path cannot even be set
1537 + up on a build without vllm-xpu-kernels installed.
1538 + """
1539 + if not current_platform.is_xpu():
1540 + return False
1541 + try:
1542 + import vllm._xpu_ops # noqa: F401
1543 + except ImportError:
1544 + return False
1545 + return True
1546 +
1547 +
1548 +def _has_w4a8_kernel() -> bool:
1549 + """True when this build can actually run the w4a8 GEMM.
1550 +
1551 + The op is registered as a side effect of importing ``vllm._xpu_ops``, so it
1552 + is not visible on ``torch.ops._xpu_C`` until that import happens.
1553 + """
1554 + return _has_xpu_ops() and hasattr(torch.ops._xpu_C, "int4_gemm_w4a8")
1555 +
1556 +
1557 +requires_xpu_ops = pytest.mark.skipif(
1558 + not _has_xpu_ops(),
1559 + reason="needs an XPU image with vllm._xpu_ops available",
1560 +)
1561 +
1562 +requires_w4a8_kernel = pytest.mark.skipif(
1563 + not _has_w4a8_kernel(),
1564 + reason="needs an XPU with the int4_gemm_w4a8 op",
1565 +)
1566 +
1567 +
1568 +def _fake_xpu(monkeypatch, ark_available: bool = True) -> None:
1569 + class DummyQuantLinear:
1570 + pass
1571 +
1572 + monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
1573 + monkeypatch.setattr(current_platform, "is_cpu", lambda: False)
1574 + monkeypatch.setattr(
1575 + _ARK_STATE,
1576 + lambda: (
1577 + (True, None, object(), DummyQuantLinear)
1578 + if ark_available
1579 + else (False, "missing", None, None)
1580 + ),
1581 + )
1582 +
1583 +
1584 +def _with_w4a8_kernel(monkeypatch) -> None:
1585 + """Present the op so the w4a8 support check passes off-XPU."""
1586 +
1587 + class FakeOps:
1588 + int4_gemm_w4a8 = object()
1589 +
1590 + monkeypatch.setattr(torch.ops, "_xpu_C", FakeOps, raising=False)
1591 +
1592 +
1593 +def _dispatch(layer_config=None):
1594 + return INCWna16Scheme().get_linear_method(
1595 + object(), object(), "layer", layer_config or make_layer_config()
1596 + )
1597 +
1598 +
1599 +@pytest.fixture
1600 +def single_tp_rank(monkeypatch):
1601 + """Parameter creation reads the TP rank/size; there is no distributed init."""
1602 + monkeypatch.setattr(
1603 + "vllm.model_executor.parameter.get_tensor_model_parallel_rank",
1604 + lambda: 0,
1605 + )
1606 + monkeypatch.setattr(
1607 + "vllm.model_executor.parameter.get_tensor_model_parallel_world_size",
1608 + lambda: 1,
1609 + )
1610 +
1611 +
1612 +def _create_weights(
1613 + method,
1614 + input_size_per_partition: int = IN_FEATURES,
1615 + output_partition_sizes: list[int] | None = None,
1616 +) -> torch.nn.Module:
1617 + class DummyLayer(torch.nn.Module):
1618 + pass
1619 +
1620 + layer = DummyLayer()
1621 + output_partition_sizes = (
1622 + [OUT_FEATURES] if output_partition_sizes is None else output_partition_sizes
1623 + )
1624 + method.create_weights(
1625 + layer=layer,
1626 + input_size_per_partition=input_size_per_partition,
1627 + output_partition_sizes=output_partition_sizes,
1628 + input_size=input_size_per_partition,
1629 + output_size=sum(output_partition_sizes),
1630 + params_dtype=torch.bfloat16,
1631 + weight_loader=lambda *args, **kwargs: None,
1632 + )
1633 + return layer
1634 +
1635 +
1636 +@pytest.fixture
1637 +def w4a8_layer(single_tp_rank):
1638 + """A w4a8 method plus a layer whose weights are created but not processed."""
1639 + method = INCXPUW4A8LinearMethod(make_layer_config())
1640 + return method, _create_weights(method)
1641 +
1642 +
1643 +def _stub_quant(monkeypatch, captured=None):
1644 + def fake_quant(x, use_sym, bits):
1645 + if captured is not None:
1646 + captured["quant_args"] = (tuple(x.shape), use_sym, bits)
1647 + m, k = x.shape
1648 + return (
1649 + torch.zeros(m, k, dtype=torch.int8),
1650 + torch.ones(m, 1, dtype=x.dtype),
1651 + torch.zeros(m, 1, dtype=torch.int32),
1652 + )
1653 +
1654 + monkeypatch.setattr(_QUANT_REF, fake_quant, raising=False)
1655 +
1656 +
1657 +def _stub_gemm(monkeypatch, captured=None):
1658 + def fake_gemm(*args):
1659 + if captured is not None:
1660 + captured["gemm_args"] = args
1661 + # The real kernel always emits float16, whatever the model dtype.
1662 + return torch.zeros(args[0].shape[0], OUT_FEATURES, dtype=torch.float16)
1663 +
1664 + monkeypatch.setattr(torch.ops._xpu_C, "int4_gemm_w4a8", fake_gemm, raising=False)
1665 +
1666 +
1667 +# ---------------------------------------------------------------------------
1668 +# Backend selection
1669 +# ---------------------------------------------------------------------------
1670 +
1671 +
1672 +def test_default_backend_keeps_ark_preference(monkeypatch) -> None:
1673 + """The "auto" default must not change upstream behaviour.
1674 +
1675 + ARK still wins whenever it is importable.
1676 + """
1677 + monkeypatch.delenv(_BACKEND_ENV, raising=False)
1678 + _fake_xpu(monkeypatch, ark_available=True)
1679 + _with_w4a8_kernel(monkeypatch)
1680 +
1681 + method = _dispatch()
1682 +
1683 + assert isinstance(method, INCLinearMethod)
1684 + assert isinstance(method.scheme, INCARKLinearMethod)
1685 +
1686 +
1687 +def test_w4a8_backend_overrides_available_ark(monkeypatch) -> None:
1688 + """An explicit request wins even though ARK would otherwise be chosen."""
1689 + monkeypatch.setenv(_BACKEND_ENV, "w4a8")
1690 + _fake_xpu(monkeypatch, ark_available=True)
1691 + _with_w4a8_kernel(monkeypatch)
1692 +
1693 + method = _dispatch()
1694 +
1695 + assert isinstance(method.scheme, INCXPUW4A8LinearMethod)
1696 + assert not isinstance(method.scheme, INCARKLinearMethod)
1697 + # w4a8 subclasses w4a16 so it can fall back for small batches.
1698 + assert isinstance(method.scheme, INCXPULinearMethod)
1699 +
1700 +
1701 +def test_w4a16_backend_overrides_available_ark(monkeypatch) -> None:
1702 + monkeypatch.setenv(_BACKEND_ENV, "w4a16")
1703 + _fake_xpu(monkeypatch, ark_available=True)
1704 +
1705 + method = _dispatch()
1706 +
1707 + assert isinstance(method.scheme, INCXPULinearMethod)
1708 + assert not isinstance(method.scheme, INCXPUW4A8LinearMethod)
1709 + assert not isinstance(method.scheme, INCARKLinearMethod)
1710 +
1711 +
1712 +def test_ark_backend_requested_but_unavailable_raises(monkeypatch) -> None:
1713 + """Explicitly asking for ARK must fail loudly, not fall back silently."""
1714 + monkeypatch.setenv(_BACKEND_ENV, "ark")
1715 + _fake_xpu(monkeypatch, ark_available=False)
1716 +
1717 + with pytest.raises(NotImplementedError, match="auto_round_kernel is unavailable"):
1718 + _dispatch()
1719 +
1720 +
1721 +@pytest.mark.parametrize("backend", ["w4a16", "w4a8"])
1722 +def test_onednn_backends_reject_int2(monkeypatch, backend) -> None:
1723 + """The oneDNN int4 GEMMs cannot serve int2; ARK is the only int2 path."""
1724 + monkeypatch.setenv(_BACKEND_ENV, backend)
1725 + _fake_xpu(monkeypatch, ark_available=True)
1726 + _with_w4a8_kernel(monkeypatch)
1727 +
1728 + with pytest.raises(NotImplementedError, match="only supports int4"):
1729 + _dispatch(make_layer_config(bits=2))
1730 +
1731 +
1732 +@pytest.mark.parametrize("group_size", [48, -1, 0])
1733 +def test_w4a8_rejects_unaligned_group_size(monkeypatch, group_size) -> None:
1734 + """The kernel needs a group size that is a positive multiple of 32."""
1735 + monkeypatch.setenv(_BACKEND_ENV, "w4a8")
1736 + _fake_xpu(monkeypatch, ark_available=True)
1737 + _with_w4a8_kernel(monkeypatch)
1738 +
1739 + with pytest.raises(NotImplementedError, match="multiple of 32"):
1740 + _dispatch(make_layer_config(group_size=group_size))
1741 +
1742 +
1743 +@pytest.mark.parametrize("group_size", [128, 64, 32])
1744 +def test_w4a8_accepts_aligned_group_size(monkeypatch, group_size) -> None:
1745 + monkeypatch.setenv(_BACKEND_ENV, "w4a8")
1746 + _fake_xpu(monkeypatch, ark_available=True)
1747 + _with_w4a8_kernel(monkeypatch)
1748 +
1749 + method = _dispatch(make_layer_config(group_size=group_size))
1750 +
1751 + assert isinstance(method.scheme, INCXPUW4A8LinearMethod)
1752 +
1753 +
1754 +def test_w4a8_requires_the_kernel(monkeypatch) -> None:
1755 + """Older vllm-xpu-kernels builds lack the op entirely."""
1756 + monkeypatch.setenv(_BACKEND_ENV, "w4a8")
1757 + _fake_xpu(monkeypatch, ark_available=True)
1758 +
1759 + class FakeOpsNoW4A8:
1760 + pass
1761 +
1762 + monkeypatch.setattr(torch.ops, "_xpu_C", FakeOpsNoW4A8, raising=False)
1763 +
1764 + with pytest.raises(NotImplementedError, match="int4_gemm_w4a8"):
1765 + _dispatch()
1766 +
1767 +
1768 +def test_invalid_backend_value_raises(monkeypatch) -> None:
1769 + import vllm.envs as envs
1770 +
1771 + monkeypatch.setenv(_BACKEND_ENV, "onednn")
1772 + with pytest.raises(ValueError, match=_BACKEND_ENV):
1773 + _ = envs.VLLM_XPU_INC_WNA16_BACKEND
1774 +
1775 +
1776 +# ---------------------------------------------------------------------------
1777 +# Partition shape
1778 +# ---------------------------------------------------------------------------
1779 +#
1780 +# ``XPUW4A8IntLinearKernel.can_implement`` rejects partition shapes whose in/out
1781 +# dims are not multiples of 8, but that check is unreachable from the INC path:
1782 +# the shapes only exist once ``create_weights`` is called, well after backend
1783 +# selection. Unaligned dims must not reach the kernel — the int4 packing divides
1784 +# both dims by 8, so an unaligned size is silently truncated into a weight tensor
1785 +# that is too small rather than rejected.
1786 +
1787 +
1788 +def test_rejects_unaligned_input_partition(single_tp_rank) -> None:
1789 + method = INCXPUW4A8LinearMethod(make_layer_config(group_size=32))
1790 +
1791 + with pytest.raises(NotImplementedError, match=r"multiples of 8.*input=132"):
1792 + _create_weights(method, input_size_per_partition=132)
1793 +
1794 +
1795 +def test_rejects_unaligned_output_partition(single_tp_rank) -> None:
1796 + method = INCXPUW4A8LinearMethod(make_layer_config())
1797 +
1798 + with pytest.raises(NotImplementedError, match=r"multiples of 8.*output=60"):
1799 + _create_weights(method, output_partition_sizes=[60])
1800 +
1801 +
1802 +def test_rejects_unaligned_sum_of_output_partitions(single_tp_rank) -> None:
1803 + """A fused QKV/MLP layer is only as aligned as the sum of its shards."""
1804 + method = INCXPUW4A8LinearMethod(make_layer_config())
1805 +
1806 + with pytest.raises(NotImplementedError, match=r"multiples of 8.*output=76"):
1807 + _create_weights(method, output_partition_sizes=[OUT_FEATURES, 12])
1808 +
1809 +
1810 +def test_reports_both_unaligned_dims(single_tp_rank) -> None:
1811 + """The message names every offending dim, not just the first."""
1812 + method = INCXPUW4A8LinearMethod(make_layer_config(group_size=32))
1813 +
1814 + with pytest.raises(NotImplementedError) as excinfo:
1815 + _create_weights(
1816 + method, input_size_per_partition=132, output_partition_sizes=[60]
1817 + )
1818 +
1819 + assert "input=132" in str(excinfo.value)
1820 + assert "output=60" in str(excinfo.value)
1821 +
1822 +
1823 +def test_accepts_aligned_partition_shape(single_tp_rank) -> None:
1824 + """Multiples of 8 that are not powers of two are still fine."""
1825 + method = INCXPUW4A8LinearMethod(make_layer_config(group_size=32))
1826 +
1827 + layer = _create_weights(
1828 + method, input_size_per_partition=96, output_partition_sizes=[24, 48]
1829 + )
1830 +
1831 + assert layer.qweight.shape == (96 // 8, 72)
1832 +
1833 +
1834 +def test_w4a16_accepts_unaligned_partition_shape(single_tp_rank) -> None:
1835 + """The constraint is a w4a8 kernel requirement; w4a16 must stay unaffected."""
1836 + method = INCXPULinearMethod(make_layer_config(group_size=32))
1837 +
1838 + layer = _create_weights(method, output_partition_sizes=[60])
1839 +
1840 + assert layer.scales.shape[1] == 60
1841 +
1842 +
1843 +# ---------------------------------------------------------------------------
1844 +# Weight processing
1845 +# ---------------------------------------------------------------------------
1846 +
1847 +
1848 +def test_keeps_fp16_scales(w4a8_layer) -> None:
1849 + """The kernel reads scales as fp16 regardless of activation dtype.
1850 +
1851 + Passing bf16 scales does not raise — it silently returns wrong results — so
1852 + an fp16 copy must exist, while the bf16 ``scales`` survive for the
1853 + small-batch w4a16 fallback.
1854 + """
1855 + method, layer = w4a8_layer
1856 + layer.scales.data.fill_(0.125) # exactly representable in both dtypes
1857 +
1858 + method.process_weights_after_loading(layer)
1859 +
1860 + assert layer.scales_fp16.dtype is torch.float16
1861 + assert layer.scales_fp16.shape == layer.scales.shape
1862 + assert torch.equal(layer.scales_fp16.data.float(), layer.scales.data.float())
1863 + # The w4a16 fallback still needs the original activation-dtype scales.
1864 + assert layer.scales.dtype is torch.bfloat16
1865 +
1866 +
1867 +def test_reuses_w4a16_weight_layout(w4a8_layer) -> None:
1868 + """w4a8 and w4a16 accept bit-identical weights, so there is no repacking."""
1869 + method, layer = w4a8_layer
1870 + qweight_before = layer.qweight.data.clone()
1871 +
1872 + method.process_weights_after_loading(layer)
1873 +
1874 + assert layer.qweight.dtype is torch.int32
1875 + assert layer.qweight.shape == (IN_FEATURES // 8, OUT_FEATURES)
1876 + assert torch.equal(layer.qweight.data, qweight_before)
1877 + # Symmetric int4: the zero point is the constant 8, as for w4a16.
1878 + assert layer.qzeros.dtype is torch.int8
1879 + assert layer.qzeros.tolist() == [8]
1880 +
1881 +
1882 +# ---------------------------------------------------------------------------
1883 +# apply_weights
1884 +# ---------------------------------------------------------------------------
1885 +
1886 +
1887 +def test_falls_back_to_w4a16_for_small_batches(monkeypatch, w4a8_layer) -> None:
1888 + """Below the token threshold, per-token quant overhead dominates."""
1889 + method, layer = w4a8_layer
1890 + method.process_weights_after_loading(layer)
1891 +
1892 + calls = []
1893 +
1894 + def record_fallback(self, layer, x, bias=None):
1895 + calls.append(tuple(x.shape))
1896 + return torch.zeros(x.shape[:-1] + (OUT_FEATURES,), dtype=x.dtype)
1897 +
1898 + monkeypatch.setattr(INCXPULinearMethod, "apply_weights", record_fallback)
1899 + monkeypatch.setattr(
1900 + torch.ops._xpu_C,
1901 + "int4_gemm_w4a8",
1902 + lambda *args: pytest.fail("small batches must not call the w4a8 kernel"),
1903 + raising=False,
1904 + )
1905 +
1906 + tokens = method._MIN_TOKENS_FOR_INT8 - 1
1907 + x = torch.zeros(tokens, IN_FEATURES, dtype=torch.bfloat16)
1908 + out = method.apply_weights(layer, x)
1909 +
1910 + assert calls == [(tokens, IN_FEATURES)]
1911 + assert out.shape == (tokens, OUT_FEATURES)
1912 +
1913 +
1914 +@requires_xpu_ops
1915 +def test_calls_kernel_at_threshold(monkeypatch, w4a8_layer) -> None:
1916 + """At/above the threshold, activations are int8-quantized per token."""
1917 + method, layer = w4a8_layer
1918 + method.process_weights_after_loading(layer)
1919 +
1920 + monkeypatch.setattr(
1921 + INCXPULinearMethod,
1922 + "apply_weights",
1923 + lambda self, layer, x, bias=None: pytest.fail(
1924 + "large batches must not take the w4a16 fallback"
1925 + ),
1926 + )
1927 + captured: dict[str, tuple] = {}
1928 + _stub_quant(monkeypatch, captured)
1929 + _stub_gemm(monkeypatch, captured)
1930 +
1931 + tokens = method._MIN_TOKENS_FOR_INT8
1932 + x = torch.zeros(tokens, IN_FEATURES, dtype=torch.bfloat16)
1933 + out = method.apply_weights(layer, x)
1934 +
1935 + # Symmetric 8-bit per-token quantization over the flattened activations.
1936 + assert captured["quant_args"] == ((tokens, IN_FEATURES), True, 8)
1937 +
1938 + (
1939 + quant_x,
1940 + x_scale,
1941 + x_zero,
1942 + qweight,
1943 + w_scale,
1944 + w_zp,
1945 + group_size,
1946 + g_idx,
1947 + bias,
1948 + ) = captured["gemm_args"]
1949 + assert quant_x.dtype is torch.int8
1950 + # Both scales must reach the kernel as fp16, whatever the model dtype.
1951 + assert x_scale.dtype is torch.float16
1952 + assert w_scale.dtype is torch.float16
1953 + assert w_scale is layer.scales_fp16
1954 + assert x_zero.dtype is torch.int32
1955 + assert qweight is layer.qweight
1956 + assert w_zp is layer.qzeros
1957 + assert group_size == 128
1958 + assert g_idx is None
1959 + assert bias is None
1960 +
1961 + # The kernel emits fp16; the result is cast back to the activation dtype.
1962 + assert out.dtype is torch.bfloat16
1963 + assert out.shape == (tokens, OUT_FEATURES)
1964 +
1965 +
1966 +@requires_xpu_ops
1967 +@pytest.mark.parametrize("act_dtype", [torch.bfloat16, torch.float16])
1968 +def test_scales_are_fp16_for_any_activation_dtype(
1969 + monkeypatch, w4a8_layer, act_dtype
1970 +) -> None:
1971 + """fp16 scales are required for bf16 *and* fp16 activations alike."""
1972 + method, layer = w4a8_layer
1973 + layer.scales.data = layer.scales.data.to(act_dtype)
1974 + method.process_weights_after_loading(layer)
1975 +
1976 + captured: dict[str, tuple] = {}
1977 + _stub_quant(monkeypatch, captured)
1978 + _stub_gemm(monkeypatch, captured)
1979 +
1980 + x = torch.zeros(method._MIN_TOKENS_FOR_INT8, IN_FEATURES, dtype=act_dtype)
1981 + out = method.apply_weights(layer, x)
1982 +
1983 + assert captured["gemm_args"][1].dtype is torch.float16
1984 + assert captured["gemm_args"][4].dtype is torch.float16
1985 + assert out.dtype is act_dtype
1986 +
1987 +
1988 +@requires_xpu_ops
1989 +def test_forwards_bias_to_kernel(monkeypatch, w4a8_layer) -> None:
1990 + """Bias is applied by the kernel, not added afterwards."""
1991 + method, layer = w4a8_layer
1992 + method.process_weights_after_loading(layer)
1993 +
1994 + captured: dict[str, tuple] = {}
1995 + _stub_quant(monkeypatch)
1996 + _stub_gemm(monkeypatch, captured)
1997 +
1998 + bias = torch.zeros(OUT_FEATURES, dtype=torch.bfloat16)
1999 + x = torch.zeros(method._MIN_TOKENS_FOR_INT8, IN_FEATURES, dtype=torch.bfloat16)
2000 + method.apply_weights(layer, x, bias)
2001 +
2002 + assert captured["gemm_args"][8] is bias
2003 +
2004 +
2005 +@requires_xpu_ops
2006 +def test_preserves_leading_dims(monkeypatch, w4a8_layer) -> None:
2007 + """3D activations must round-trip through the flatten/reshape."""
2008 + method, layer = w4a8_layer
2009 + method.process_weights_after_loading(layer)
2010 +
2011 + _stub_quant(monkeypatch)
2012 + _stub_gemm(monkeypatch)
2013 +
2014 + tokens = method._MIN_TOKENS_FOR_INT8
2015 + x = torch.zeros(4, tokens, IN_FEATURES, dtype=torch.bfloat16)
2016 + out = method.apply_weights(layer, x)
2017 +
2018 + assert out.shape == (4, tokens, OUT_FEATURES)
2019 + assert out.dtype is torch.bfloat16
2020 +
2021 +
2022 +@requires_xpu_ops
2023 +def test_batch_gate_uses_flattened_token_count(monkeypatch, w4a8_layer) -> None:
2024 + """The gate counts total tokens, not the leading dimension."""
2025 + method, layer = w4a8_layer
2026 + method.process_weights_after_loading(layer)
2027 +
2028 + monkeypatch.setattr(
2029 + INCXPULinearMethod,
2030 + "apply_weights",
2031 + lambda self, layer, x, bias=None: pytest.fail(
2032 + "flattened token count is at the threshold; must use w4a8"
2033 + ),
2034 + )
2035 + _stub_quant(monkeypatch)
2036 + _stub_gemm(monkeypatch)
2037 +
2038 + tokens = method._MIN_TOKENS_FOR_INT8
2039 + # Each leading slice is below the threshold, but 4 * (tokens // 4) is not.
2040 + x = torch.zeros(4, tokens // 4, IN_FEATURES, dtype=torch.bfloat16)
2041 + out = method.apply_weights(layer, x)
2042 +
2043 + assert out.shape == (4, tokens // 4, OUT_FEATURES)
2044 +
2045 +
2046 +# ---------------------------------------------------------------------------
2047 +# End to end on real weights
2048 +# ---------------------------------------------------------------------------
2049 +
2050 +
2051 +@requires_w4a8_kernel
2052 +@pytest.mark.parametrize("model", E2E_MODELS)
2053 +def test_e2e_generates_with_real_weights(vllm_runner, monkeypatch, model) -> None:
2054 + """The whole path — real int4 checkpoint, real kernel, real output.
2055 +
2056 + The stubbed tests above cannot catch a wrong-but-plausible calling
2057 + convention, because the fake kernel accepts anything. Greedy decoding of a
2058 + factual prompt does: wrongly set scales or a bad weight layout turn the
2059 + answer into noise rather than raising.
2060 + """
2061 + monkeypatch.setenv(_BACKEND_ENV, "w4a8")
2062 +
2063 + with vllm_runner(
2064 + model, max_model_len=512, enforce_eager=True, gpu_memory_utilization=0.55
2065 + ) as vllm_model:
2066 + out = vllm_model.generate_greedy(["The capital of France is"], max_tokens=8)
2067 +
2068 + assert "Paris" in out[0][1]
2069 +
2070 +
2071 +@requires_w4a8_kernel
2072 +def test_e2e_selects_w4a8_for_every_linear(vllm_runner, monkeypatch) -> None:
2073 + """Guard against the validation silently becoming dead code.
2074 +
2075 + If backend selection stopped routing here, the generation test above would
2076 + still pass on the w4a16 path, so assert the method is actually in use and
2077 + that every partition shape it accepted honours the alignment rule.
2078 + """
2079 + monkeypatch.setenv(_BACKEND_ENV, "w4a8")
2080 + # apply_model pickles the closure below.
2081 + monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
2082 +
2083 + with vllm_runner(
2084 + E2E_MODELS[0].values[0],
2085 + max_model_len=512,
2086 + enforce_eager=True,
2087 + gpu_memory_utilization=0.55,
2088 + ) as vllm_model:
2089 +
2090 + def check_model(model):
2091 + shapes = []
2092 + for module in model.modules():
2093 + scheme = getattr(getattr(module, "quant_method", None), "scheme", None)
2094 + if isinstance(scheme, INCXPUW4A8LinearMethod):
2095 + # qweight is the packed [in // 8, out] layout.
2096 + packed_in, out = module.qweight.shape
2097 + shapes.append((packed_in * 8, out))
2098 +
2099 + assert shapes, "no linear layer used the w4a8 method"
2100 + assert all(i % 8 == 0 and o % 8 == 0 for i, o in shapes), shapes
2101 +
2102 + vllm_model.apply_model(check_model)
vllm/envs.py
+11
@@ -313,6 +313,7 @@ if TYPE_CHECKING:
313 VLLM_NIXL_EP_MAX_NUM_RANKS: int = 32
314 VLLM_XPU_ENABLE_XPU_GRAPH: bool = False
315 VLLM_XPU_USE_SAMPLER_KERNEL: bool = True
316 + VLLM_XPU_INC_WNA16_BACKEND: Literal["auto", "ark", "w4a16", "w4a8"] = "auto"
317 VLLM_LORA_ENABLE_DUAL_STREAM: bool = False
318 VLLM_GPU_NIC_PCIE_MAPPING: str = ""
319 VLLM_NIC_SELECTION_VARS: str = ""
@@ -2111,6 +2112,16 @@ environment_variables: dict[str, Callable[[], Any]] = {
2112 "VLLM_XPU_USE_SAMPLER_KERNEL": lambda: bool(
2113 int(os.getenv("VLLM_XPU_USE_SAMPLER_KERNEL", "1"))
2114 ),
2115 + # Kernel backend for INC weight-only intN (WNA16) linear layers on XPU.
2116 + # "auto" keeps the default preference order (ARK when importable, else the
2117 + # oneDNN w4a16 path). "ark" forces the auto_round_kernel backend, "w4a16"
2118 + # forces oneDNN int4_gemm_w4a16, and "w4a8" additionally quantizes
2119 + # activations to per-token int8 (int4_gemm_w4a8) for large token counts.
2120 + # The two oneDNN backends are int4-only; ARK also serves int2.
2121 + # Which one is fastest is device-dependent, so this is left as an opt-in.
2122 + "VLLM_XPU_INC_WNA16_BACKEND": env_with_choices(
2123 + "VLLM_XPU_INC_WNA16_BACKEND", "auto", ["auto", "ark", "w4a16", "w4a8"]
2124 + ),
2125 # Enable simple KV offload.
2126 "VLLM_USE_SIMPLE_KV_OFFLOAD": lambda: bool(
2127 int(os.getenv("VLLM_USE_SIMPLE_KV_OFFLOAD", "0"))
vllm/model_executor/layers/quantization/inc/schemes/inc_w4a8_linear.py new
+112
@@ -0,0 +1,112 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +from torch.nn.parameter import Parameter
6 +
7 +from .inc_wna16_linear import INCXPULinearMethod
8 +
9 +
10 +class INCXPUW4A8LinearMethod(INCXPULinearMethod):
11 + """XPU linear method for INC int4 weights with dynamic int8 activations.
12 +
13 + Uses the same GPTQ-packed "NT" qweight layout as ``INCXPULinearMethod`` —
14 + ``int4_gemm_w4a8`` and ``int4_gemm_w4a16`` accept bit-identical weights, so
15 + no extra repacking is needed. Activations are dynamically quantized
16 + per-token to symmetric int8, which keeps the GEMM on the int8 datapath
17 + instead of upconverting the weights to the activation dtype. That is a win
18 + for compute-bound shapes (large token counts, e.g. diffusion), where w4a16
19 + is dominated by dequantization rather than by weight bandwidth.
20 +
21 + Both scale tensors must be float16: the kernel reads them as fp16
22 + regardless of the activation dtype, so bf16 scales are silently
23 + misinterpreted. The activations themselves may be bf16 or fp16.
24 +
25 + Opt in with ``VLLM_XPU_INC_WNA16_BACKEND=w4a8``. Whether this beats ARK is
26 + device-dependent, so it is not the default: it was measured faster on B70
27 + (Xe2), where ARK cannot use its XMX int8 path and falls back to fp16.
28 + """
29 +
30 + # int8 activations only pay off once the GEMM is compute-bound. Below this
31 + # token count the per-token quantization overhead dominates and w4a16 is
32 + # faster, so fall back to it per-call.
33 + _MIN_TOKENS_FOR_INT8 = 512
34 +
35 + # int4_gemm_w4a8 requires both GEMM dims to be multiples of 8. Tensor
36 + # parallelism shards these, so an aligned model dimension can still yield an
37 + # unaligned partition.
38 + _DIM_ALIGNMENT = 8
39 +
40 + def create_weights(
41 + self,
42 + layer: torch.nn.Module,
43 + input_size_per_partition: int,
44 + output_partition_sizes: list[int],
45 + input_size: int,
46 + output_size: int,
47 + params_dtype: torch.dtype,
48 + **extra_weight_attrs,
49 + ) -> None:
50 + output_size_per_partition = sum(output_partition_sizes)
51 + unaligned = [
52 + (name, size)
53 + for name, size in (
54 + ("input", input_size_per_partition),
55 + ("output", output_size_per_partition),
56 + )
57 + if size % self._DIM_ALIGNMENT != 0
58 + ]
59 + if unaligned:
60 + raise NotImplementedError(
61 + "VLLM_XPU_INC_WNA16_BACKEND=w4a8 requires partitioned in/out "
62 + f"sizes that are multiples of {self._DIM_ALIGNMENT}, got "
63 + + ", ".join(f"{name}={size}" for name, size in unaligned)
64 + + f". Partition shape: ({input_size_per_partition}, "
65 + f"{output_size_per_partition})."
66 + )
67 + super().create_weights(
68 + layer=layer,
69 + input_size_per_partition=input_size_per_partition,
70 + output_partition_sizes=output_partition_sizes,
71 + input_size=input_size,
72 + output_size=output_size,
73 + params_dtype=params_dtype,
74 + **extra_weight_attrs,
75 + )
76 +
77 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
78 + super().process_weights_after_loading(layer)
79 + # The kernel reads both scale tensors as fp16; keep an fp16 copy for the
80 + # w4a8 path and leave ``scales`` untouched for the w4a16 fallback.
81 + layer.scales_fp16 = Parameter(
82 + layer.scales.data.to(torch.float16), requires_grad=False
83 + )
84 +
85 + def apply_weights(
86 + self,
87 + layer: torch.nn.Module,
88 + x: torch.Tensor,
89 + bias: torch.Tensor | None = None,
90 + ) -> torch.Tensor:
91 + reshaped_x = x.reshape(-1, x.shape[-1])
92 + if reshaped_x.shape[0] < self._MIN_TOKENS_FOR_INT8:
93 + return super().apply_weights(layer, x, bias)
94 +
95 + from vllm._xpu_ops import xpu_ops as ops
96 +
97 + out_shape = x.shape[:-1] + (layer.qweight.shape[1],)
98 + quant_x, x_scale, x_zero = ops.dynamic_per_token_int8_quant_ref(
99 + reshaped_x, True, 8
100 + )
101 + out = torch.ops._xpu_C.int4_gemm_w4a8(
102 + quant_x,
103 + x_scale.to(torch.float16),
104 + x_zero,
105 + layer.qweight,
106 + layer.scales_fp16,
107 + layer.qzeros,
108 + self.group_size,
109 + None,
110 + bias,
111 + )
112 + return out.to(x.dtype).reshape(out_shape)
vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py
+46
@@ -3,6 +3,7 @@
3
4 from typing import TYPE_CHECKING
5
6 +from vllm import envs
7 from vllm.logger import init_logger
8 from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
9 from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
@@ -21,6 +22,32 @@ logger = init_logger(__name__)
22
23 XPU_WNA16_SUPPORTED_BITS = {2, 4}
24
25 +# Backends selectable through VLLM_XPU_INC_WNA16_BACKEND that are served by the
26 +# oneDNN int4 GEMMs rather than by ARK. These only cover int4.
27 +XPU_ONEDNN_BACKENDS = ("w4a16", "w4a8")
28 +
29 +
30 +def _check_xpu_w4a8_supported(layer_config: "INCLayerConfig", prefix: str) -> None:
31 + """Raise unless ``int4_gemm_w4a8`` can serve this layer.
32 +
33 + The backend is requested explicitly, so an unusable configuration is an
34 + error rather than something to silently fall back from.
35 + """
36 + import torch
37 +
38 + if not hasattr(torch.ops._xpu_C, "int4_gemm_w4a8"):
39 + raise NotImplementedError(
40 + "VLLM_XPU_INC_WNA16_BACKEND=w4a8 requires the int4_gemm_w4a8 op, "
41 + "which this build of vllm-xpu-kernels does not provide. "
42 + f"Layer: {prefix}."
43 + )
44 + if layer_config.group_size <= 0 or layer_config.group_size % 32 != 0:
45 + raise NotImplementedError(
46 + "VLLM_XPU_INC_WNA16_BACKEND=w4a8 requires a group size that is a "
47 + f"positive multiple of 32, got {layer_config.group_size}. "
48 + f"Layer: {prefix}."
49 + )
50 +
51
52 class INCWna16Scheme(INCScheme):
53 @staticmethod
@@ -38,12 +65,31 @@ class INCWna16Scheme(INCScheme):
65 if current_platform.is_xpu():
66 if layer_config.bits in XPU_WNA16_SUPPORTED_BITS and layer_config.sym:
67 from .inc_ark_ops import get_ark_state
68 + from .inc_w4a8_linear import INCXPUW4A8LinearMethod
69 from .inc_wna16_linear import (
70 INCARKLinearMethod,
71 INCXPULinearMethod,
72 )
73
74 + backend = envs.VLLM_XPU_INC_WNA16_BACKEND
75 + if backend in XPU_ONEDNN_BACKENDS:
76 + if layer_config.bits != 4:
77 + raise NotImplementedError(
78 + f"VLLM_XPU_INC_WNA16_BACKEND={backend} only supports "
79 + f"int4, got int{layer_config.bits}. Layer: {prefix}."
80 + )
81 + if backend == "w4a8":
82 + _check_xpu_w4a8_supported(layer_config, prefix)
83 + return INCLinearMethod(INCXPUW4A8LinearMethod(layer_config))
84 + return INCLinearMethod(INCXPULinearMethod(layer_config))
85 +
86 is_ark_available, ark_error, _, _ = get_ark_state()
87 + if backend == "ark" and not is_ark_available:
88 + raise NotImplementedError(
89 + "VLLM_XPU_INC_WNA16_BACKEND=ark was requested but "
90 + f"auto_round_kernel is unavailable: "
91 + f"{ark_error or 'unknown error'}. Layer: {prefix}."
92 + )
93 if is_ark_available:
94 return INCLinearMethod(INCARKLinearMethod(layer_config))
95 elif layer_config.bits == 2: