Skip to content

vllm.model_executor.kernels.linear.mxfp4.rdna3_hip

MXFP4 linear kernel for RDNA3/RDNA3.5 (gfx11) with an int8-activation decode path.

Routes on the number of activation rows M

1 <= M <= MAX_W4A8_BATCH_SIZE (decode, speculative verify): HIP W4A8 GEMV (_rocm_C.mxfp4_w4a8_gemv_rdna3). Activations are quantized to int8 per 32-element group and reduced with int8 dot instructions. otherwise (prefill, profiling): dequantize the weight and run a high-precision GEMM on the unquantized activations (weight-only).

Both paths read the checkpoint layout directly ([N, K/2] uint8 E2M1 codes, [N, K/32] uint8 E8M0 scales), so there is no weight repack.

Selected when int8 activations are requested, e.g. with --quantization-config.linear.activation int8_per_group_32. For W4A4 checkpoints this replaces the MXFP4 activation QDQ: decode uses int8 activations and prefill uses the model dtype, which kept accuracy where mixing MXFP4-QDQ prefill with int8 decode did not.

Classes:

Rdna3W4A8MxFp4LinearKernel

Bases: MxFp4LinearKernel

MXFP4 GEMM for RDNA3/RDNA3.5: W4A8 int8-dot GEMV for M <= 8, dequant + high-precision GEMM above.

Source code in vllm/model_executor/kernels/linear/mxfp4/rdna3_hip.py 
 81
 82
 83
 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
123
124
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
class Rdna3W4A8MxFp4LinearKernel(MxFp4LinearKernel):
    """MXFP4 GEMM for RDNA3/RDNA3.5: W4A8 int8-dot GEMV for M <= 8, dequant +
    high-precision GEMM above."""

    @classmethod
    def is_supported(
        cls, compute_capability: int | None = None
    ) -> tuple[bool, str | None]:
        if not current_platform.is_rocm():
            return False, "only targets ROCm"
        from vllm.platforms.rocm import on_gfx11

        if not on_gfx11():
            return False, "requires RDNA3/RDNA3.5 (gfx11)"
        if not (
            hasattr(torch.ops, "_rocm_C")
            and hasattr(torch.ops._rocm_C, "mxfp4_w4a8_gemv_rdna3")
        ):
            return False, "_rocm_C.mxfp4_w4a8_gemv_rdna3 is not built"
        return True, None

    @classmethod
    def _supports_quant_scheme(
        cls,
        weight_key: QuantKey | None,
        activation_key: QuantKey | None,
    ) -> bool:
        return weight_key == kMxfp4Static and activation_key == kInt8Dynamic32Sym

    @classmethod
    def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
        if config.activation_quant_key != kInt8Dynamic32Sym:
            return False, "only supports dynamic per-32-group int8 activations"
        if not has_quark():
            return False, "amd-quark package not available (needed for M > 8)"
        return True, None

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        weight = layer.weight.data
        weight_scale = layer.weight_scale.data
        if weight.dtype != torch.uint8 or weight_scale.dtype != torch.uint8:
            raise ValueError(
                "Rdna3W4A8MxFp4LinearKernel expects uint8 packed weights and "
                f"uint8 E8M0 scales, got {weight.dtype} and {weight_scale.dtype}"
            )
        N, k_half = weight.shape
        K = 2 * k_half
        if K % _MXFP4_GROUP_SIZE != 0 or weight_scale.shape != (
            N,
            K // _MXFP4_GROUP_SIZE,
        ):
            raise ValueError(
                "Rdna3W4A8MxFp4LinearKernel expects weight [N, K/2] and "
                f"weight_scale [N, K/32] with K % 32 == 0, got "
                f"{tuple(weight.shape)} and {tuple(weight_scale.shape)}"
            )
        layer.weight = Parameter(weight.contiguous(), requires_grad=False)
        layer.weight_scale = Parameter(weight_scale.contiguous(), requires_grad=False)

    def apply_weights(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        x_2d = x.reshape(-1, x.shape[-1])
        if not x_2d.is_contiguous():
            x_2d = x_2d.contiguous()
        out = torch.ops.vllm.rdna3_mxfp4_w4a8_apply(
            x_2d, layer.weight, layer.weight_scale, bias
        )
        return out.reshape(*x.shape[:-1], out.shape[-1])