[INC][ARK] Direct Register Custom Op for ARK (#46361)
Signed-off-by: Zhenzhong1 <zhenzhong.xu@intel.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Zhenzhong Xu committed
Jul 6, 2026 at 13:45 UTC
d9c1767cd47f2744bd855597f3145d944572e329
6 files changed
+152
-43
docs/features/quantization/inc.md
+1
-4
@@ -75,14 +75,11 @@ vllm serve Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound \
75
--max-model-len 4096
76
```
77
78
-!!! note
79
- To deploy `wNa16` models on Intel GPU/CPU, please add `--enforce-eager` for now.
80
-
78
## Evaluating the Quantized Model with vLLM
79
80
```bash
81
lm_eval --model vllm \
85
- --model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048,enforce_eager=True" \
82
+ --model_args pretrained="Intel/DeepSeek-R1-0528-Qwen3-8B-int4-AutoRound,max_model_len=8192,max_num_batched_tokens=32768,max_num_seqs=128,gpu_memory_utilization=0.8,dtype=bfloat16,max_gen_toks=2048" \
83
--tasks gsm8k \
84
--num_fewshot 5 \
85
--batch_size 128
requirements/xpu.txt
+1
-1
@@ -16,5 +16,5 @@ torch==2.12.0
16
torchaudio
17
torchvision
18
19
-auto_round_lib>=0.13.3
19
+auto_round_lib>=0.14.0
20
vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.10.1/vllm_xpu_kernels-0.1.10.1-cp38-abi3-manylinux_2_28_x86_64.whl
tests/quantization/test_auto_round.py
+5
-5
@@ -61,7 +61,7 @@ MODELS = [
61
)
62
@pytest.mark.parametrize("model", MODELS)
63
def test_auto_round_model(vllm_runner, model):
64
- with vllm_runner(model, enforce_eager=True) as llm:
64
+ with vllm_runner(model) as llm:
65
output = llm.generate_greedy(["The capital of France is"], max_tokens=8)
66
67
assert output
@@ -336,7 +336,7 @@ def test_wna16_xpu_prefers_ark_when_available(monkeypatch) -> None:
336
monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
337
monkeypatch.setattr(current_platform, "is_cpu", lambda: False)
338
monkeypatch.setattr(
339
- "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state",
339
+ "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state",
340
lambda: (True, None, object(), DummyQuantLinear),
341
)
342
@@ -355,7 +355,7 @@ def test_wna16_xpu_falls_back_when_ark_unavailable(monkeypatch) -> None:
355
monkeypatch.setattr(current_platform, "is_xpu", lambda: True)
356
monkeypatch.setattr(current_platform, "is_cpu", lambda: False)
357
monkeypatch.setattr(
358
- "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state",
358
+ "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state",
359
lambda: (False, "missing", None, None),
360
)
361
@@ -377,7 +377,7 @@ def test_wna16_cpu_gptq_prefers_ark_when_available(monkeypatch) -> None:
377
monkeypatch.setattr(current_platform, "is_xpu", lambda: False)
378
monkeypatch.setattr(current_platform, "is_cpu", lambda: True)
379
monkeypatch.setattr(
380
- "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state",
380
+ "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state",
381
lambda: (True, None, object(), DummyQuantLinear),
382
)
383
@@ -398,7 +398,7 @@ def test_wna16_cpu_gptq_raises_when_ark_and_marlin_unavailable(
398
monkeypatch.setattr(current_platform, "is_xpu", lambda: False)
399
monkeypatch.setattr(current_platform, "is_cpu", lambda: True)
400
monkeypatch.setattr(
401
- "vllm.model_executor.layers.quantization.inc.schemes.inc_wna16_linear.get_ark_state",
401
+ "vllm.model_executor.layers.quantization.inc.schemes.inc_ark_ops.get_ark_state",
402
lambda: (False, "missing", None, None),
403
)
404
monkeypatch.setattr(
vllm/model_executor/layers/quantization/inc/schemes/inc_ark_ops.py
new
+123
@@ -0,0 +1,123 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from functools import lru_cache
5
+from typing import Any
6
+
7
+import torch
8
+
9
+from vllm.logger import init_logger
10
+from vllm.platforms import current_platform
11
+from vllm.utils.torch_utils import direct_register_custom_op
12
+
13
+logger = init_logger(__name__)
14
+
15
+_OPS_REGISTERED = False
16
+
17
+
18
+@lru_cache(maxsize=1)
19
+def get_ark_state() -> tuple[bool, str | None, Any | None, Any | None]:
20
+ """Return ARK availability, error details, cached module, and QuantLinear."""
21
+ try:
22
+ import auto_round_kernel as ark
23
+ from auto_round_kernel.qlinear import QuantLinear
24
+
25
+ logger.info("Successfully imported auto_round_kernel.")
26
+ except ImportError as error:
27
+ return False, str(error), None, None
28
+
29
+ if getattr(ark, "cpu_lib", None) is None and getattr(ark, "xpu_lib", None) is None:
30
+ return (
31
+ False,
32
+ "No ARK backend library is available.",
33
+ None,
34
+ None,
35
+ )
36
+
37
+ logger.info("Successfully loaded auto_round_kernel backend library.")
38
+ return True, None, ark, QuantLinear
39
+
40
+
41
+def _inc_ark_woq_linear_impl(
42
+ x: torch.Tensor,
43
+ qweight: torch.Tensor,
44
+ bias: torch.Tensor | None,
45
+ out_features: int,
46
+ in_features: int,
47
+ group_size: int,
48
+ compute_type: str,
49
+ weight_type: str,
50
+ scale_type: str,
51
+ asym: bool,
52
+) -> torch.Tensor:
53
+ ark = get_ark_state()[2]
54
+ assert ark is not None
55
+
56
+ return ark.woqgemm_linear(
57
+ x,
58
+ qweight,
59
+ bias,
60
+ out_features,
61
+ in_features,
62
+ group_size,
63
+ compute_type,
64
+ weight_type,
65
+ scale_type,
66
+ asym,
67
+ )
68
+
69
+
70
+def _inc_ark_woq_linear_fake(
71
+ x: torch.Tensor,
72
+ qweight: torch.Tensor,
73
+ bias: torch.Tensor | None,
74
+ out_features: int,
75
+ in_features: int,
76
+ group_size: int,
77
+ compute_type: str,
78
+ weight_type: str,
79
+ scale_type: str,
80
+ asym: bool,
81
+) -> torch.Tensor:
82
+ del qweight
83
+ del bias
84
+ del in_features
85
+ del group_size
86
+ del compute_type
87
+ del weight_type
88
+ del scale_type
89
+ del asym
90
+ return torch.empty(
91
+ (*x.shape[:-1], out_features),
92
+ dtype=x.dtype,
93
+ device=x.device,
94
+ )
95
+
96
+
97
+class ark_ops:
98
+ @staticmethod
99
+ def register_ops_once() -> None:
100
+ global _OPS_REGISTERED
101
+ if _OPS_REGISTERED:
102
+ return
103
+
104
+ is_available, error_str, _, _ = get_ark_state()
105
+ if not is_available:
106
+ logger.debug(
107
+ "Skip registering ark op because ARK is unavailable: %s",
108
+ error_str or "unknown error",
109
+ )
110
+ return
111
+
112
+ direct_register_custom_op(
113
+ op_name="inc_ark_woq_linear",
114
+ op_func=_inc_ark_woq_linear_impl,
115
+ fake_impl=_inc_ark_woq_linear_fake,
116
+ dispatch_key=current_platform.dispatch_key,
117
+ )
118
+ _OPS_REGISTERED = True
119
+
120
+
121
+ark_ops.register_ops_once()
122
+
123
+__all__ = ["get_ark_state"]
vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_linear.py
+20
-31
@@ -1,13 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-from functools import lru_cache
4
from typing import TYPE_CHECKING, Any
5
6
import torch
7
from torch.nn.parameter import Parameter
8
10
-from vllm.logger import init_logger
9
from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig
10
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQConfig
11
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
@@ -22,35 +20,10 @@ from vllm.scalar_type import scalar_types
20
21
from .inc_scheme import INCLinearScheme
22
25
-logger = init_logger(__name__)
26
-
23
if TYPE_CHECKING:
24
from ..config_parser import INCLayerConfig
25
26
31
-@lru_cache(maxsize=1)
32
-def get_ark_state() -> tuple[bool, str | None, Any | None, Any | None]:
33
- """Return ARK availability, error details, cached module, and QuantLinear."""
34
- try:
35
- import auto_round_kernel as ark
36
- from auto_round_kernel.qlinear import QuantLinear
37
-
38
- logger.info("Successfully imported auto_round_kernel.")
39
- except ImportError as error:
40
- return False, str(error), None, None
41
-
42
- if getattr(ark, "cpu_lib", None) is None and getattr(ark, "xpu_lib", None) is None:
43
- return (
44
- False,
45
- "No ARK backend library is available.",
46
- None,
47
- None,
48
- )
49
- logger.info("Successfully loaded auto_round_kernel backend library.")
50
-
51
- return True, None, ark, QuantLinear
52
-
53
-
27
class INCWNA16LinearScheme(INCLinearScheme):
28
def __init__(self, layer_config: "INCLayerConfig") -> None:
29
self.layer_config = layer_config
@@ -380,6 +353,8 @@ class INCARKLinearMethod(INCXPULinearBase):
353
def __init__(self, layer_config: "INCLayerConfig") -> None:
354
super().__init__(layer_config)
355
356
+ from .inc_ark_ops import get_ark_state
357
+
358
is_available, error_str, _, quant_linear_cls = get_ark_state()
359
if not is_available or quant_linear_cls is None:
360
reason = error_str or "unknown error"
@@ -453,9 +428,13 @@ class INCARKLinearMethod(INCXPULinearBase):
428
ark_linear.bias.copy_(layer.bias.detach())
429
430
ark_linear.post_init()
456
- layer.ark_linear = ark_linear
431
458
- del layer.qweight
432
+ layer.qweight = Parameter(ark_linear.qweight.detach(), requires_grad=False)
433
+ layer.ark_bias = ark_linear.bias
434
+ layer.ark_compute_type = ark_linear.cdt
435
+ layer.ark_weight_type = ark_linear.wdt
436
+ layer.ark_scale_type = ark_linear.sdt
437
+
438
if hasattr(layer, "qzeros"):
439
del layer.qzeros
440
del layer.scales
@@ -466,8 +445,18 @@ class INCARKLinearMethod(INCXPULinearBase):
445
x: torch.Tensor,
446
bias: torch.Tensor | None = None,
447
) -> torch.Tensor:
469
- del bias
470
- return layer.ark_linear.forward(x)
448
+ return torch.ops.vllm.inc_ark_woq_linear.default(
449
+ x,
450
+ layer.qweight,
451
+ layer.ark_bias,
452
+ layer.out_features,
453
+ layer.in_features,
454
+ self.group_size,
455
+ layer.ark_compute_type,
456
+ layer.ark_weight_type,
457
+ layer.ark_scale_type,
458
+ not self.sym,
459
+ )
460
461
462
class INCXPUW4A16LinearScheme(INCXPULinearMethod):
vllm/model_executor/layers/quantization/inc/schemes/inc_wna16_scheme.py
+2
-2
@@ -36,10 +36,10 @@ class INCWna16Scheme(INCScheme):
36
del config, layer
37
if current_platform.is_xpu():
38
if layer_config.bits == 4 and layer_config.sym:
39
+ from .inc_ark_ops import get_ark_state
40
from .inc_wna16_linear import (
41
INCARKLinearMethod,
42
INCXPULinearMethod,
42
- get_ark_state,
43
)
44
45
is_ark_available, ark_error, _, _ = get_ark_state()
@@ -57,10 +57,10 @@ class INCWna16Scheme(INCScheme):
57
58
if current_platform.is_cpu() and layer_config.is_gptq:
59
if layer_config.bits == 4 and layer_config.sym:
60
+ from .inc_ark_ops import get_ark_state
61
from .inc_wna16_linear import (
62
INCARKLinearMethod,
63
INCWNA16LinearScheme,
63
- get_ark_state,
64
)
65
66
is_ark_available, ark_error, _, _ = get_ark_state()