Skip to content

vllm.model_executor.layers.fused_moe.oracle.w4a8_int8

Classes:

Functions:

W4A8Int8MoEKernelOracle

Bases: MoEKernelOracle[W4A8Int8MoeBackend]

W4A8 INT8 MoE backend selection for one deployment configuration.

Source code in vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py 
73
74
75
76
77
78
79
80
81
class W4A8Int8MoEKernelOracle(MoEKernelOracle[W4A8Int8MoeBackend]):
    """W4A8 INT8 MoE backend selection for one deployment configuration."""

    name = "W4A8 Int8"

    def backend_to_kernel_cls(
        self, backend: W4A8Int8MoeBackend
    ) -> list[type[mk.FusedMoEExperts]]:
        return backend_to_kernel_cls(backend)

_get_priority_backends(moe_config)

Get available backends in priority order based on platform and config.

Currently only CPU INT4 backend is available for W4A8 INT8 MoE.

Source code in vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py 
33
34
35
36
37
38
39
40
41
42
def _get_priority_backends(
    moe_config: FusedMoEConfig,
) -> list[W4A8Int8MoeBackend]:
    """Get available backends in priority order based on platform and config.

    Currently only CPU INT4 backend is available for W4A8 INT8 MoE.
    """
    if current_platform.is_cpu():
        return [W4A8Int8MoeBackend.CPU_INT4]
    return []

backend_to_kernel_cls(backend)

Map W4A8Int8MoeBackend to kernel class.

Source code in vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
def backend_to_kernel_cls(
    backend: W4A8Int8MoeBackend,
) -> list[type[mk.FusedMoEExperts]]:
    """Map W4A8Int8MoeBackend to kernel class."""
    if backend == W4A8Int8MoeBackend.CPU_INT4:
        from vllm.model_executor.layers.fused_moe.experts.cpu_int4_moe import (
            CPUExpertsInt4,
        )

        return [CPUExpertsInt4]

    else:
        raise ValueError(f"Unknown W4A8 Int8 MoE backend: {backend.value}")

convert_to_w4a8_int8_moe_format(w13_weight, w2_weight, w13_weight_scale, w2_weight_scale, group_size, w13_bias=None, w2_bias=None)

Pack INT4 MoE weights to KleidiAI format.

This function packs the INT4 weights (stored as int8 values) into the format expected by the KleidiAI dynamic_4bit_int_moe kernel.

Parameters:

  • w13_weight

    (Tensor) –

    [E, 2*IN, H] int8 tensor (int4 values in [-8,7])

  • w2_weight

    (Tensor) –

    [E, H, IN] int8 tensor (int4 values in [-8,7])

  • w13_weight_scale

    (Tensor) –

    [E, 2*IN, H/g or 1] scale tensor

  • w2_weight_scale

    (Tensor) –

    [E, H, IN/g or 1] scale tensor

  • group_size

    (int) –

    Quantization group size (-1 for channel-wise)

  • w13_bias

    (Tensor | None, default: None ) –

    Optional [E, 2*IN] bias tensor

  • w2_bias

    (Tensor | None, default: None ) –

    Optional [E, H] bias tensor

Returns:

Source code in vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
def convert_to_w4a8_int8_moe_format(
    w13_weight: torch.Tensor,
    w2_weight: torch.Tensor,
    w13_weight_scale: torch.Tensor,
    w2_weight_scale: torch.Tensor,
    group_size: int,
    w13_bias: torch.Tensor | None = None,
    w2_bias: torch.Tensor | None = None,
) -> tuple[
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor,
    torch.Tensor | None,
    torch.Tensor | None,
]:
    """Pack INT4 MoE weights to KleidiAI format.

    This function packs the INT4 weights (stored as int8 values) into
    the format expected by the KleidiAI dynamic_4bit_int_moe kernel.

    Args:
        w13_weight: [E, 2*IN, H] int8 tensor (int4 values in [-8,7])
        w2_weight: [E, H, IN] int8 tensor (int4 values in [-8,7])
        w13_weight_scale: [E, 2*IN, H/g or 1] scale tensor
        w2_weight_scale: [E, H, IN/g or 1] scale tensor
        group_size: Quantization group size (-1 for channel-wise)
        w13_bias: Optional [E, 2*IN] bias tensor
        w2_bias: Optional [E, H] bias tensor

    Returns:
        Tuple of (w13_packed, w2_packed) tensors

    """
    # Derive dimensions from tensor shapes
    E = w13_weight.shape[0]  # num_experts
    w13_out_features = w13_weight.shape[1]  # 2 * intermediate_size
    H = w13_weight.shape[2]  # w13_in_features (hidden_size)
    IN = w2_weight.shape[2]  # w2_in_features (intermediate_size)
    w2_out_features = w2_weight.shape[1]  # Should equal H

    # Pack per expert
    w13_packed_list = []
    w2_packed_list = []

    for e in range(E):
        w13_packed_list.append(
            pack_int4_weights_for_kleidi(
                w13_weight[e],  # [2I, H]
                w13_weight_scale[e],  # [2I, H/g or 1]
                w13_bias[e] if w13_bias is not None else None,  # [2I]
                H,
                w13_out_features,
                group_size,
            )
        )
        w2_packed_list.append(
            pack_int4_weights_for_kleidi(
                w2_weight[e],  # [H, IN]
                w2_weight_scale[e],  # [H, IN/g or 1]
                w2_bias[e] if w2_bias is not None else None,  # [H]
                IN,
                w2_out_features,  # in_features=IN, out_features=H
                group_size,
            )
        )

    # Stack all experts
    w13_packed = torch.stack(w13_packed_list, dim=0)
    w2_packed = torch.stack(w2_packed_list, dim=0)
    empty = torch.empty(0)

    return w13_packed, w2_packed, empty, empty, empty, empty

make_w4a8_int8_moe_kernel(moe_quant_config, moe_config, experts_cls, routing_tables=None)

Create FusedMoEKernel for W4A8 Int8 MoE.

Parameters:

Returns:

Source code in vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py 
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
def make_w4a8_int8_moe_kernel(
    moe_quant_config: FusedMoEQuantConfig,
    moe_config: FusedMoEConfig,
    experts_cls: type[mk.FusedMoEExperts],
    routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,
) -> mk.FusedMoEKernel:
    """Create FusedMoEKernel for W4A8 Int8 MoE.

    Args:
        moe_quant_config: Quantization configuration
        moe_config: MoE configuration
        experts_cls: Expert kernel class (should be CPUExpertsInt4)
        routing_tables: Optional routing tables for expert parallelism

    Returns:
        Configured FusedMoEKernel instance

    """
    # Create Experts.
    # W4A8 Int8 currently only supports monolithic interface
    if not issubclass(experts_cls, mk.FusedMoEExpertsMonolithic):
        raise ValueError(
            f"W4A8 Int8 MoE only supports monolithic experts, "
            f"but got {experts_cls.__name__}"
        )

    return make_moe_kernel(moe_quant_config, moe_config, experts_cls, routing_tables)

make_w4a8_int8_moe_quant_config(block_shape=None)

Create FusedMoEQuantConfig for W4A8 Int8 MoE.

Parameters:

  • block_shape

    (tuple[int, int] | None, default: None ) –

    Quantization block shape (row, col). For channel-wise: (-1, 1) or None For group-wise: (1, group_size)

Returns:

Source code in vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py 
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
def make_w4a8_int8_moe_quant_config(
    block_shape: tuple[int, int] | None = None,
) -> FusedMoEQuantConfig:
    """Create FusedMoEQuantConfig for W4A8 Int8 MoE.

    Args:
        block_shape: Quantization block shape (row, col).
                    For channel-wise: (-1, 1) or None
                    For group-wise: (1, group_size)

    Returns:
        FusedMoEQuantConfig with appropriate settings for W4A8 Int8

    """
    # W4A8 Int8 uses static weight quantization, dynamic activation quantization
    # Weights are 4-bit (stored as int8, packed to uint8),
    # activations are dynamically quantized to 8-bit in kernel

    group_shape = GroupShape(*block_shape) if block_shape is not None else None

    return FusedMoEQuantConfig(
        # Activations: unquantized (FP/BF16), dynamically quantized in kernel
        _a1=FusedMoEQuantDesc(shape=group_shape),
        _a2=FusedMoEQuantDesc(shape=group_shape),
        # Weights: INT8 (4-bit values), pre-packed with scales
        # dtype=None means already quantized/packed
        _w1=FusedMoEQuantDesc(dtype=None, shape=group_shape),
        _w2=FusedMoEQuantDesc(dtype=None, shape=group_shape),
    )

map_w4a8_int8_backend(runner_backend)

Map user's MoEBackend to W4A8Int8MoeBackend.

Source code in vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py 
60
61
62
63
64
65
66
67
68
69
70
def map_w4a8_int8_backend(runner_backend: MoEBackend) -> W4A8Int8MoeBackend:
    """Map user's MoEBackend to W4A8Int8MoeBackend."""
    mapping = {
        "cpu": W4A8Int8MoeBackend.CPU_INT4,
    }
    if backend := mapping.get(runner_backend):
        return backend
    raise ValueError(
        f"moe_backend='{runner_backend}' is not supported for W4A8 Int8 MoE. "
        f"Expected one of {list(mapping.keys())}."
    )

pack_int4_weights_for_kleidi(int4_as_int8, scales, bias, in_features, out_features, group_size)

Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format.

Parameters:

  • int4_as_int8

    (Tensor) –

    [out, in] int8 tensor with values in [-8, 7]

  • scales

    (Tensor) –

    [out, in//group_size] or [out, 1] for channel-wise

  • bias

    (Tensor | None) –

    [out] optional bias

  • in_features

    (int) –

    Input dimension

  • out_features

    (int) –

    Output dimension

  • group_size

    (int) –

    Quantization group size (-1 for channel-wise)

Returns:

  • Tensor –

    Packed weight tensor in KleidiAI format

Source code in vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py 
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
def pack_int4_weights_for_kleidi(
    int4_as_int8: torch.Tensor,
    scales: torch.Tensor,
    bias: torch.Tensor | None,
    in_features: int,
    out_features: int,
    group_size: int,
) -> torch.Tensor:
    """Pack INT4 weights (stored as int8 in [-8,7]) to KleidiAI format.

    Args:
        int4_as_int8: [out, in] int8 tensor with values in [-8, 7]
        scales: [out, in//group_size] or [out, 1] for channel-wise
        bias: [out] optional bias
        in_features: Input dimension
        out_features: Output dimension
        group_size: Quantization group size (-1 for channel-wise)

    Returns:
        Packed weight tensor in KleidiAI format

    """
    # Shift to unsigned nibble [0, 15]
    tmp = int4_as_int8.add(8)
    # Pack pairs along input dimension
    uint8_nibbles = ((tmp[:, 1::2] << 4) | tmp[:, ::2]).to(torch.uint8)

    # Determine scale dtype based on group_size
    # KleidiAI groupwise kernels accept bfloat16 scales
    # KleidiAI channelwise kernels accept float32 scales
    scale_dtype = torch.float32 if group_size == -1 else torch.bfloat16
    scales_typed = scales.to(scale_dtype)
    bias_typed = None if bias is None else bias.to(torch.float32)

    # Pack using KleidiAI op
    actual_group_size = in_features if group_size == -1 else group_size
    return torch.ops.aten._dyn_quant_pack_4bit_weight(
        uint8_nibbles,
        scales_typed,
        bias_typed,
        actual_group_size,
        in_features,
        out_features,
    )

select_w4a8_int8_moe_backend(config, weight_key, activation_key)

Select the primary W4A8 Int8 MoE backend.

Parameters:

  • config

    (FusedMoEConfig) –

    MoE configuration

  • weight_key

    (QuantKey | None) –

    Weight quantization key (should be one of kInt4W4A8Static*)

  • activation_key

    (QuantKey | None) –

    Activation quantization key (currently unused for W4A8)

Returns:

Source code in vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
def select_w4a8_int8_moe_backend(
    config: FusedMoEConfig,
    weight_key: QuantKey | None,
    activation_key: QuantKey | None,
) -> tuple[W4A8Int8MoeBackend, type[mk.FusedMoEExperts]]:
    """Select the primary W4A8 Int8 MoE backend.

    Args:
        config: MoE configuration
        weight_key: Weight quantization key (should be one of kInt4W4A8Static*)
        activation_key: Activation quantization key (currently unused for W4A8)

    Returns:
        Tuple of (backend, kernel_class)

    """
    AVAILABLE_BACKENDS = _get_priority_backends(config)

    if not AVAILABLE_BACKENDS:
        raise NotImplementedError("W4A8 Int8 MoE is only supported on CPU platforms")

    oracle = W4A8Int8MoEKernelOracle(config, weight_key, activation_key)

    # Handle explicit moe_backend from user.
    runner_backend = config.moe_backend
    if runner_backend != "auto":
        requested_backend = map_w4a8_int8_backend(runner_backend)
        return oracle.return_or_raise(requested_backend, AVAILABLE_BACKENDS)

    # Select kernels in order of backend.
    for backend in AVAILABLE_BACKENDS:
        kernel_cls, _ = oracle.try_backend(backend)
        if kernel_cls is not None:
            logger.info_once(oracle.make_log_backend(backend, AVAILABLE_BACKENDS))
            return backend, kernel_cls

    raise NotImplementedError(
        "No W4A8 Int8 MoE backend supports the deployment configuration."
    )