chore: import upstream snapshot with attribution

vllm/model_executor/kernels/attention/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/model_executor/kernels/attention/dsa/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/model_executor/kernels/attention/dsa/dcp_indexer_cutedsl.py

@@ -0,0 +1,420 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from functools import cache
+
+import cutlass
+import cutlass.cute as cute
+import torch
+from cuda.bindings.driver import CUstream
+from cutlass import Float32, Int32, Uint32, Uint64
+from quack.compile_utils import make_fake_tensor
+
+from vllm.cute_utils import recast_val
+from vllm.triton_utils import tl, triton
+
+
+def stable_topk_from_gathered_candidates_cutedsl(
+    gathered: torch.Tensor,
+    topk: int,
+    out: torch.Tensor | None = None,
+) -> torch.Tensor:
+    if out is None:
+        out = torch.empty(
+            (gathered.shape[0], topk),
+            dtype=torch.int32,
+            device=gathered.device,
+        )
+    StableTopKFromGatheredCandidatesKernel.compile(topk, gathered.shape[1])(
+        gathered, out
+    )
+    return out
+
+
+def pack_dcp_topk_candidates_cutedsl(
+    logits: torch.Tensor,
+    topk_indices: torch.Tensor,
+    packed: torch.Tensor,
+    dcp_rank: int,
+    dcp_world_size: int,
+    cp_interleave: int,
+    row_starts: torch.Tensor | None,
+) -> None:
+    topk = topk_indices.shape[1]
+    grid = (topk_indices.shape[0], triton.cdiv(topk, 512))
+    row_starts_arg = row_starts if row_starts is not None else topk_indices
+    _pack_dcp_topk_candidates_triton_kernel[grid](
+        logits,
+        topk_indices,
+        packed,
+        row_starts_arg,
+        logits.stride(0),
+        logits.stride(1),
+        topk_indices.stride(0),
+        topk_indices.stride(1),
+        packed.stride(0),
+        packed.stride(1),
+        packed.stride(2),
+        logits.shape[1],
+        DCP_RANK=dcp_rank,
+        DCP_WORLD_SIZE=dcp_world_size,
+        CP_INTERLEAVE=cp_interleave,
+        HAS_ROW_STARTS=row_starts is not None,
+        TOPK=topk,
+        BLOCK_SIZE=512,
+        num_warps=8,
+    )
+
+
+@triton.jit
+def _pack_dcp_topk_candidates_triton_kernel(
+    logits,
+    topk_indices,
+    packed,
+    row_starts,
+    logits_stride0: tl.constexpr,
+    logits_stride1: tl.constexpr,
+    topk_stride0: tl.constexpr,
+    topk_stride1: tl.constexpr,
+    packed_stride0: tl.constexpr,
+    packed_stride1: tl.constexpr,
+    packed_stride2: tl.constexpr,
+    num_cols,
+    DCP_RANK: tl.constexpr,
+    DCP_WORLD_SIZE: tl.constexpr,
+    CP_INTERLEAVE: tl.constexpr,
+    HAS_ROW_STARTS: tl.constexpr,
+    TOPK: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    row = tl.program_id(0)
+    tile = tl.program_id(1)
+    cols = tile * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = cols < TOPK
+
+    local_idx = tl.load(
+        topk_indices + row * topk_stride0 + cols * topk_stride1,
+        mask=mask,
+        other=-1,
+    )
+    valid = local_idx >= 0
+    safe_local_idx = tl.maximum(local_idx, 0)
+
+    row_start = 0
+    if HAS_ROW_STARTS:
+        row_start = tl.load(row_starts + row)
+
+    score_col = safe_local_idx + row_start
+    score_col = tl.minimum(score_col, tl.maximum(num_cols - 1, 0))
+    score = tl.load(
+        logits + row * logits_stride0 + score_col * logits_stride1,
+        mask=mask & valid,
+        other=-float("inf"),
+    )
+
+    global_id = (
+        (safe_local_idx // CP_INTERLEAVE) * (DCP_WORLD_SIZE * CP_INTERLEAVE)
+        + DCP_RANK * CP_INTERLEAVE
+        + safe_local_idx % CP_INTERLEAVE
+    )
+    global_id = tl.where(valid, global_id, -1)
+
+    packed_base = packed + row * packed_stride0 + cols * packed_stride1
+    tl.store(packed_base, score, mask=mask)
+    tl.store(packed_base + packed_stride2, global_id.to(tl.float32), mask=mask)
+
+
+@cute.jit
+def _warp_scan_inclusive_i32(val: Int32, lane: Int32) -> Int32:
+    for i in cutlass.range_constexpr(cute.arch.WARP_SIZE.bit_length() - 1):
+        offset = 1 << i
+        partial = cute.arch.shuffle_sync_up(val, offset=offset, mask_and_clamp=0)
+        if lane >= offset:
+            val += partial
+    return val
+
+
+@cute.jit
+def _block_scan_inclusive_i32(
+    val: Int32,
+    lane: Int32,
+    warp_id: Int32,
+    warp_scratch: cute.Tensor,
+    warps_per_block: int,
+) -> Int32:
+    prefix = _warp_scan_inclusive_i32(val, lane)
+    if lane == Int32(cute.arch.WARP_SIZE - 1):
+        warp_scratch[0, warp_id] = prefix
+    cute.arch.sync_threads()
+
+    if warp_id == Int32(0):
+        warp_total = Int32(0)
+        if lane < Int32(warps_per_block):
+            warp_total = warp_scratch[0, lane]
+        warp_prefix = _warp_scan_inclusive_i32(warp_total, lane)
+        if lane < Int32(warps_per_block):
+            warp_scratch[0, lane] = warp_prefix - warp_total
+    cute.arch.sync_threads()
+
+    return prefix + warp_scratch[0, warp_id]
+
+
+class StableTopKFromGatheredCandidatesKernel:
+    tb_size = 512
+    hist_bins = 2048
+    radix_bits = (hist_bins - 1).bit_length()
+    assert hist_bins == 1 << radix_bits
+    key_bits = Uint64.width
+    radix_passes = (key_bits + radix_bits - 1) // radix_bits
+    final_radix_bits = key_bits - radix_bits * (radix_passes - 1)
+    hist_chunks = (hist_bins + tb_size - 1) // tb_size
+    warps_per_block = tb_size // cute.arch.WARP_SIZE
+
+    def __init__(self, topk: int, num_candidates: int):
+        assert num_candidates % self.tb_size == 0, (
+            "StableTopKFromGatheredCandidatesKernel requires candidate count "
+            f"to be a multiple of {self.tb_size}, got {num_candidates}"
+        )
+        self.topk = topk
+        self.keys_per_thread = num_candidates // self.tb_size
+
+        @cute.struct
+        class SharedStorage:
+            hist: cute.struct.MemRange[Int32, self.hist_bins]
+            committed_count: cute.struct.MemRange[Int32, 1]
+            running_count: cute.struct.MemRange[Int32, 1]
+            threshold_bin: cute.struct.MemRange[Int32, 1]
+            threshold_found: cute.struct.MemRange[Int32, 1]
+            include_threshold_bin: cute.struct.MemRange[Int32, 1]
+            prefix_s: cute.struct.Align[cute.struct.MemRange[Uint64, 1], 8]
+            warp_totals: cute.struct.MemRange[Int32, self.warps_per_block]
+
+        self.shared_storage = SharedStorage
+
+    @cute.jit
+    def __call__(
+        self,
+        gathered: cute.Tensor,
+        out: cute.Tensor,
+        stream: CUstream,
+    ):
+        grid = (gathered.shape[0], 1, 1)
+        self.kernel(gathered, out).launch(
+            grid=grid,
+            block=(self.tb_size, 1, 1),
+            stream=stream,
+        )
+
+    @cute.jit
+    def _stable_key(self, score: Float32, token_id: Int32) -> Uint64:
+        bits = recast_val(score, Uint32)
+        mask = Uint32(0x80000000)
+        if (bits & Uint32(0x80000000)) != Uint32(0):
+            mask = Uint32(0xFFFFFFFF)
+        score_key = Uint64(bits ^ mask) << Uint64(32)
+        id_key = Uint64(~Uint32(token_id))
+        key = score_key | id_key
+        if token_id < Int32(0):
+            key = Uint64(0)
+        return key
+
+    @cute.jit
+    def _prefix_matches(
+        self,
+        key: Uint64,
+        prefix: Uint64,
+        prefix_bits: Int32,
+    ):
+        matches = prefix_bits == Int32(0)
+        if prefix_bits != Int32(0):
+            shift = Int32(self.key_bits) - prefix_bits
+            matches = (key >> Uint64(shift)) == (prefix >> Uint64(shift))
+        return matches
+
+    @cute.jit
+    def _radix_pass(
+        self,
+        keys: cute.Tensor,
+        output: cute.Tensor,
+        storage,
+        tid: Int32,
+        step: Int32,
+        bits: int,
+        is_final_pass: bool,
+    ):
+        hist_smem = storage.hist.get_tensor(cute.make_layout((self.hist_bins,)))
+        committed_count_smem = storage.committed_count.data_ptr()
+        running_count_smem = storage.running_count.data_ptr()
+        threshold_bin_smem = storage.threshold_bin.data_ptr()
+        threshold_found_smem = storage.threshold_found.data_ptr()
+        include_threshold_bin_smem = storage.include_threshold_bin.data_ptr()
+        prefix_smem = storage.prefix_s.data_ptr()
+        warp_totals_smem = storage.warp_totals.get_tensor(
+            cute.make_layout((1, self.warps_per_block))
+        )
+
+        prefix_bits = step * Int32(self.radix_bits)
+        num_bins = 1 << bits
+        block_scan_iterations = (num_bins + self.tb_size - 1) // self.tb_size
+        shift = Int32(self.key_bits) - prefix_bits - Int32(bits)
+        bin_mask = Uint64(num_bins - 1)
+        prefix = prefix_smem.load()
+
+        for chunk in cutlass.range_constexpr(self.hist_chunks):
+            hist_smem[tid + Int32(chunk * self.tb_size)] = Int32(0)
+        if tid == Int32(0):
+            running_count_smem.store(committed_count_smem.load())
+            include_threshold_bin_smem.store(Int32(0))
+            threshold_found_smem.store(Int32(0))
+        cute.arch.sync_threads()
+
+        for key_idx in cutlass.range_constexpr(self.keys_per_thread):
+            key = keys[key_idx]
+            if self._prefix_matches(key, prefix, prefix_bits):
+                bin_idx = Int32((key >> Uint64(shift)) & bin_mask)
+                cute.arch.atomic_add(
+                    hist_smem.iterator + bin_idx,
+                    Int32(1),
+                    sem="relaxed",
+                    scope="cta",
+                )
+        cute.arch.sync_threads()
+
+        lane = cute.arch.lane_idx()
+        warp_id = cute.arch.warp_idx()
+        # Each iteration scans one tb_size-wide slice of bins, high to low.
+        iter = Int32(0)
+        threshold_found = threshold_found_smem.load()
+        while threshold_found == Int32(0) and iter < Int32(block_scan_iterations):
+            bin_idx = Int32(num_bins - 1) - (iter * Int32(self.tb_size) + tid)
+            count = hist_smem[bin_idx]
+            chunk_inclusive = _block_scan_inclusive_i32(
+                count,
+                lane,
+                warp_id,
+                warp_totals_smem,
+                self.warps_per_block,
+            )
+            running_count = running_count_smem.load()
+            prior_in_scan_slice = chunk_inclusive - count
+            remaining = Int32(self.topk) - running_count - prior_in_scan_slice
+            if count > Int32(0) and remaining > Int32(0) and remaining <= count:
+                threshold_bin_smem.store(bin_idx)
+                if count <= remaining or cutlass.const_expr(is_final_pass):
+                    include_threshold_bin_smem.store(Int32(1))
+                threshold_found_smem.store(Int32(1))
+            # Barrier: every thread must finish reading running_count for this
+            # slice before tb_size-1 advances it, else a warp racing ahead to
+            # the store makes a lagging thread double-count the slice total
+            # (-> remaining too small -> threshold too high -> under-fill).
+            cute.arch.sync_threads()
+            if tid == Int32(self.tb_size - 1):
+                running_count_smem.store(running_count + chunk_inclusive)
+            cute.arch.sync_threads()
+
+            threshold_found = threshold_found_smem.load()
+            iter += Int32(1)
+
+        threshold = threshold_bin_smem.load()
+        should_include_threshold = include_threshold_bin_smem.load() != Int32(0)
+        for key_idx in cutlass.range_constexpr(self.keys_per_thread):
+            key = keys[key_idx]
+            if self._prefix_matches(key, prefix, prefix_bits):
+                bin_idx = Int32((key >> Uint64(shift)) & bin_mask)
+                selected = bin_idx > threshold
+                if should_include_threshold:
+                    selected = selected or bin_idx == threshold
+                if selected:
+                    dst = cute.arch.atomic_add(
+                        committed_count_smem,
+                        Int32(1),
+                        sem="relaxed",
+                        scope="cta",
+                    )
+                    if dst < Int32(self.topk):
+                        output[dst] = recast_val(~Uint32(key), Int32)
+        cute.arch.sync_threads()
+
+        pass_finished = include_threshold_bin_smem.load()
+        if tid == Int32(0) and pass_finished == Int32(0):
+            prefix_smem.store(prefix | (Uint64(threshold) << Uint64(shift)))
+        cute.arch.sync_threads()
+        return pass_finished
+
+    @cute.kernel
+    def kernel(
+        self,
+        input: cute.Tensor,
+        out: cute.Tensor,
+    ):
+        row, _, _ = cute.arch.block_idx()
+        tid, _, _ = cute.arch.thread_idx()
+        input_row = input[row, None, None]
+        output_row = out[row, None]
+        keys = cute.make_rmem_tensor((self.keys_per_thread,), Uint64)
+
+        smem = cutlass.utils.SmemAllocator()
+        storage = smem.allocate(self.shared_storage, 8)
+        committed_count_smem = storage.committed_count.data_ptr()
+        prefix_smem = storage.prefix_s.data_ptr()
+        for i in range(tid, self.topk, self.tb_size):
+            output_row[i] = Int32(-1)
+
+        for key_idx in cutlass.range_constexpr(self.keys_per_thread):
+            col = tid + Int32(key_idx * self.tb_size)
+            score = Float32(input_row[col, 0])
+            token_id = Int32(input_row[col, 1])
+            keys[key_idx] = self._stable_key(score, token_id)
+
+        if tid == Int32(0):
+            committed_count_smem.store(Int32(0))
+            prefix_smem.store(Uint64(0))
+        cute.arch.sync_threads()
+
+        step = Int32(0)
+        finished = Int32(0)
+        while finished == Int32(0) and step < Int32(self.radix_passes - 1):
+            finished = self._radix_pass(
+                keys,
+                output_row,
+                storage,
+                tid,
+                step,
+                self.radix_bits,
+                False,
+            )
+            step += Int32(1)
+
+        if finished == Int32(0):
+            self._radix_pass(
+                keys,
+                output_row,
+                storage,
+                tid,
+                Int32(self.radix_passes - 1),
+                self.final_radix_bits,
+                True,
+            )
+
+    @cache
+    @staticmethod
+    def compile(topk: int, num_candidates: int):
+        num_rows = cute.sym_int()
+
+        gathered = cute.runtime.make_fake_tensor(
+            Float32,
+            (num_rows, num_candidates, 2),
+            stride=(cute.sym_int64(divisibility=2), 2, 1),
+            assumed_align=8,
+        )
+        out = make_fake_tensor(Int32, (num_rows, topk), divisibility=1)
+
+        kernel = StableTopKFromGatheredCandidatesKernel(topk, num_candidates)
+        stream = cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=True)
+        return cute.compile(
+            kernel,
+            gathered,
+            out,
+            stream,
+            options="--enable-tvm-ffi",
+        )

vllm/model_executor/kernels/linear/base.py

@@ -0,0 +1,332 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any, ClassVar, Generic, TypeVar
+
+import torch
+from typing_extensions import Self
+
+from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
+
+
+@dataclass
+class MMLinearLayerConfig: ...
+
+
+@dataclass
+class Params:
+    """Base class for quantized layer parameters.
+
+    This class provides a typed interface for accessing quantized weights and scales
+    from layer modules. It serves as a parameter container that can be extracted from
+    layers and passed to kernel implementations.
+
+    Attributes:
+        weight: The quantized weight tensor
+        weight_scale: weight scaling factors
+        input_scale: Optional input scaling factors
+
+    Class Variables:
+        WEIGHT: Attribute name for weight tensor on the layer module
+        WEIGHT_SCALE: Attribute name for weight scale tensor on the layer module
+        INPUT_SCALE: Attribute name for input scale tensor on the layer module
+
+    Important:
+        The string values of WEIGHT, WEIGHT_SCALE, and INPUT_SCALE class variables
+        MUST match the attribute names used in the corresponding quantization method's
+        create_weights() implementation.
+        For example, if FP8LinearMethod.create_weights()
+        sets layer.weight and layer.weight_scale,
+        then WEIGHT="weight" and
+        WEIGHT_SCALE="weight_scale" must be used here.
+
+    Usage:
+        ```python
+        # Extract parameters from a quantized layer
+        params = Params.from_layer(layer)
+
+        # Access typed parameters
+        output = func(input, params.weight, params.weight_scale)
+        ```
+    """
+
+    weight: torch.Tensor
+    weight_scale: torch.Tensor
+    input_scale: torch.Tensor | None
+
+    # Attribute names on the layer
+    WEIGHT: ClassVar[str] = "weight"
+    WEIGHT_SCALE: ClassVar[str] = "weight_scale"
+    INPUT_SCALE: ClassVar[str] = "input_scale"
+
+    @classmethod
+    def from_layer(cls, layer: torch.nn.Module) -> Self:
+        return cls(
+            weight=getattr(layer, cls.WEIGHT),
+            weight_scale=getattr(layer, cls.WEIGHT_SCALE),
+            input_scale=getattr(layer, cls.INPUT_SCALE, None),
+        )
+
+
+@dataclass
+class FP8Params(Params):
+    """FP8 layer parameters with typed fields"""
+
+    input_scale_ub: torch.Tensor | None
+
+    INPUT_SCALE_UB: ClassVar[str] = "input_scale_ub"
+
+    @classmethod
+    def from_layer(cls, layer: torch.nn.Module) -> "FP8Params":
+        """Extract parameters from layer"""
+        return cls(
+            weight=getattr(layer, cls.WEIGHT),
+            weight_scale=getattr(layer, cls.WEIGHT_SCALE),
+            input_scale=getattr(layer, cls.INPUT_SCALE, None),
+            input_scale_ub=getattr(layer, cls.INPUT_SCALE_UB, None),
+        )
+
+
+@dataclass
+class Int8Params(Params):
+    """Int8 layer parameters with typed fields"""
+
+    input_zero_point: torch.Tensor | None
+    azp_adj: torch.Tensor | None
+
+    INPUT_ZERO_POINT: ClassVar[str] = "input_zero_point"
+    AZP_ADJ: ClassVar[str] = "azp_adj"
+
+    @classmethod
+    def from_layer(cls, layer: torch.nn.Module) -> "Int8Params":
+        """Extract parameters from layer"""
+        return cls(
+            weight=getattr(layer, cls.WEIGHT),
+            weight_scale=getattr(layer, cls.WEIGHT_SCALE),
+            input_scale=getattr(layer, cls.INPUT_SCALE, None),
+            input_zero_point=getattr(layer, cls.INPUT_ZERO_POINT, None),
+            azp_adj=getattr(layer, cls.AZP_ADJ, None),
+        )
+
+
+_ParamsT = TypeVar("_ParamsT", bound=Params)
+_ConfigT = TypeVar("_ConfigT", bound=MMLinearLayerConfig)
+
+
+class MMLinearKernel(ABC, Generic[_ConfigT, _ParamsT]):
+    """Abstract base class for quantized matrix multiplication kernels.
+
+    This class provides the interface for implementing custom quantized linear layer
+    kernels in vLLM. Subclasses should implement specific quantization strategies
+    (e.g., FP8, INT8) and their corresponding compute kernels.
+
+    Generic Type Parameters:
+        _ConfigT: Configuration type for the kernel (subclass of MMLinearLayerConfig).
+                  Contains kernel-specific settings like quantization keys, dtypes, etc.
+        _ParamsT: Parameter type for the kernel (subclass of Params).
+                  Defines the quantized weights and scales needed by the kernel.
+
+    Typical Usage:
+        1. Define a config dataclass inheriting from MMLinearLayerConfig
+        2. Define a params dataclass inheriting from Params (or FP8Params/Int8Params)
+        3. Subclass MMLinearKernel with your config and params types
+        4. Implement all abstract methods
+        5. Register the kernel with the quantization method
+
+    Example:
+        ```python
+        @dataclass
+        class MyKernelConfig(MMLinearLayerConfig):
+            static: bool
+            output_dtype: torch.dtype
+
+
+        @dataclass
+        class MyKernelParams(FP8Params):
+            custom_scale: torch.Tensor
+            CUSTOM_SCALE: ClassVar[str] = "custom_scale"
+
+
+        class MyKernel(MMLinearKernel[MyKernelConfig, MyKernelParams]):
+            @classmethod
+            def is_supported(cls, compute_capability=None):
+                if compute_capability and compute_capability < 90:
+                    return False, "Requires compute capability >= 9.0"
+                return True, None
+
+            @classmethod
+            def can_implement(cls, config):
+                if not config.static:
+                    return False, "Only static quantization supported"
+                return True, None
+
+            def process_weights_after_loading(self, layer):
+                # Preprocess weights for the kernel
+                params = self._get_layer_params(layer)
+                processed = preprocess_weights(params.weight)
+                replace_parameter(layer, params.WEIGHT, processed)
+
+            def _get_layer_params(self, layer, **kwargs):
+                return MyKernelParams.from_layer(layer)
+
+            def apply_weights(self, layer, x, bias=None, **kwargs):
+                params = self._get_layer_params(layer)
+                # Call your custom kernel
+                output = my_custom_kernel(x, params.weight, params.weight_scale)
+                if bias is not None:
+                    output += bias
+                return output
+        ```
+
+    Lifecycle:
+        1. Kernel selection: is_supported() and can_implement() check compatibility
+        2. Initialization: __init__() creates kernel instance with config
+        3. Weight loading: process_weights_after_loading() preprocesses weights
+        4. Inference: apply_weights() executes the quantized matmul
+    """
+
+    @classmethod
+    @abstractmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        """Check if this kernel is supported on the current hardware.
+
+        This method checks hardware-level compatibility (e.g., GPU architecture,
+        compute capability, available instructions). It's called during kernel
+        selection to filter out kernels that cannot run on the current device.
+
+        Args:
+            compute_capability: GPU compute capability (e.g., 80 for A100, 90 for H100).
+                               If None, should check the current device.
+
+        Returns:
+            A tuple of (is_supported, reason):
+                - is_supported: True if the kernel can run on this hardware
+                - reason: If not supported, a string explaining why; otherwise None
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def can_implement(cls, config: _ConfigT) -> tuple[bool, str | None]:
+        """Check if this kernel can implement the given configuration.
+
+        This method checks configuration-level compatibility (e.g., quantization
+        scheme, group sizes, static vs dynamic quantization). It's called after
+        is_supported() to determine if this kernel can handle the specific
+        quantization configuration.
+
+        Args:
+            config: The kernel configuration to check
+
+        Returns:
+            A tuple of (can_implement, reason):
+                - can_implement: True if this kernel supports the config
+                - reason: If not supported, a string explaining why; otherwise None
+            ```
+        """
+        raise NotImplementedError
+
+    def __init__(self, config: _ConfigT) -> None:
+        """Initialize the kernel with the given configuration.
+
+        Args:
+            config: Kernel-specific configuration containing settings like
+                   quantization keys, output dtypes, etc.
+        """
+        self.config = config
+
+    def input_quant_key(self) -> QuantKey | None:
+        """Return the input quantization key supported by this kernel. If the kernel
+        does not support input quantization outside of the kernel, return None.
+        """
+        return None
+
+    @abstractmethod
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Process and transform weights after loading from checkpoint.
+
+        This method is called once after weights are loaded but before inference.
+        Use it to preprocess weights into the format required by your kernel
+        (e.g., reordering, padding, format conversion).
+
+        Modifications should be done in-place using replace_parameter() to ensure
+        the layer's parameters are properly updated.
+
+        Args:
+            layer: The layer module containing the weights to process
+
+        Example:
+            ```python
+            def process_weights_after_loading(self, layer):
+                params = self._get_layer_params(layer)
+                # Reorder weights for better memory access
+                weight_reordered = reorder_weights(params.weight)
+                replace_parameter(layer, params.WEIGHT, weight_reordered)
+            ```
+        """
+        raise NotImplementedError
+
+    # return a covariant type in the subclass
+    @abstractmethod
+    def _get_layer_params(self, layer: torch.nn.Module, **kwargs: Any) -> _ParamsT:
+        """Extract typed parameters from the layer module.
+
+        This internal method retrieves the quantized weights and scales from
+        the layer as a typed parameter object. Subclasses should typically
+        delegate to ParamsClass.from_layer().
+
+        Args:
+            layer: The layer module containing the parameters
+            **kwargs: Additional arguments
+
+        Returns:
+            A typed parameter object containing weights, scales, and other
+            quantization parameters
+
+        Example:
+            ```python
+            def _get_layer_params(self, layer, **kwargs):
+                return MyKernelParams.from_layer(layer)
+            ```
+        """
+        raise NotImplementedError
+
+    def get_output_padding(self) -> int | None:
+        """Get the number of output tokens to pad for this kernel.
+
+        Some kernels require input padding for optimal performance.
+        Override this method to specify padding requirements.
+
+        Returns:
+            Number of tokens to pad, or None for no padding (default)
+        """
+        return None
+
+    @abstractmethod
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        """Apply the quantized weights to the input tensor.
+
+        This is the main inference method that performs the quantized matrix
+        multiplication. It should handle input quantization (if needed), call
+        the underlying kernel, and apply bias.
+
+        Args:
+            layer: The layer module containing the quantized weights
+            x: Input tensor of shape [..., in_features]
+            bias: Optional bias tensor of shape [out_features]
+            **kwargs: Additional kernel-specific arguments
+
+        Returns:
+            Output tensor of shape [..., out_features]
+        """
+        raise NotImplementedError

vllm/model_executor/kernels/linear/mixed_precision/MPLinearKernel.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.scalar_type import ScalarType
+
+
+@dataclass
+class MPLinearLayerConfig:
+    full_weight_shape: tuple[int, int]  # [in, out]
+    partition_weight_shape: tuple[int, int]
+    weight_type: ScalarType
+    act_type: torch.dtype
+    group_size: int
+    zero_points: bool
+    has_g_idx: bool
+    out_type: torch.dtype | None = None
+
+
+class MPLinearKernel(ABC):
+    @classmethod
+    @abstractmethod
+    def get_min_capability(cls) -> int:
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        raise NotImplementedError
+
+    def __init__(
+        self,
+        c: MPLinearLayerConfig,
+        w_q_param_name: str,
+        w_s_param_name: str,
+        w_zp_param_name: str | None = None,
+        w_gidx_param_name: str | None = None,
+    ) -> None:
+        assert self.can_implement(c)
+        self.config = c
+        self.w_q_name = w_q_param_name
+        self.w_s_name = w_s_param_name
+        if c.zero_points:
+            assert w_zp_param_name is not None
+        if c.has_g_idx:
+            assert w_gidx_param_name is not None
+        self.w_zp_name: str | None = w_zp_param_name
+        self.w_gidx_name: str | None = w_gidx_param_name
+
+    @abstractmethod
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def _transform_param(
+        self, layer: torch.nn.Module, name: str | None, fn: Callable
+    ) -> None:
+        if name is not None and getattr(layer, name, None) is not None:
+            old_param = getattr(layer, name)
+            new_param = fn(old_param)
+            # replace the parameter with torch.nn.Parameter for TorchDynamo
+            # compatibility
+            replace_parameter(
+                layer, name, torch.nn.Parameter(new_param.data, requires_grad=False)
+            )
+
+    def _get_weight_params(
+        self, layer: torch.nn.Module
+    ) -> tuple[
+        torch.Tensor,  # w_q
+        torch.Tensor,  # w_s
+        torch.Tensor | None,  # w_zp,
+        torch.Tensor | None,  # w_gidx
+    ]:
+        return (
+            getattr(layer, self.w_q_name),
+            getattr(layer, self.w_s_name),
+            getattr(layer, self.w_zp_name or "", None),
+            getattr(layer, self.w_gidx_name or "", None),
+        )

vllm/model_executor/kernels/linear/mixed_precision/__init__.py

@@ -0,0 +1,66 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.model_executor.kernels.linear.mixed_precision.allspark import (
+    AllSparkLinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.conch import (
+    ConchLinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.cpu import (
+    CPUWNA16LinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.cutlass import (
+    CutlassW4A8LinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.dynamic_4bit import (
+    Dynamic4bitLinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.exllama import (
+    ExllamaLinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.humming import (
+    HummingLinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.machete import (
+    MacheteLinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.marlin import (
+    MarlinLinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
+    MPLinearKernel,
+    MPLinearLayerConfig,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import (
+    RDNA3W4A16LinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16 import (
+    TritonW4A16LinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
+    XPUW4A8IntLinearKernel,
+    XPUwNa16LinearKernel,
+)
+from vllm.model_executor.kernels.linear.mixed_precision.zentorch import (
+    ZentorchWNA16LinearKernel,
+)
+
+__all__ = [
+    "MPLinearKernel",
+    "MPLinearLayerConfig",
+    "AllSparkLinearKernel",
+    "ConchLinearKernel",
+    "CPUWNA16LinearKernel",
+    "CutlassW4A8LinearKernel",
+    "Dynamic4bitLinearKernel",
+    "ExllamaLinearKernel",
+    "HummingLinearKernel",
+    "MacheteLinearKernel",
+    "MarlinLinearKernel",
+    "RDNA3W4A16LinearKernel",
+    "TritonW4A16LinearKernel",
+    "XPUW4A8IntLinearKernel",
+    "XPUwNa16LinearKernel",
+    "ZentorchWNA16LinearKernel",
+]

vllm/model_executor/kernels/linear/mixed_precision/allspark.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.model_executor.layers.quantization.utils.allspark_utils import (
+    ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
+    check_allspark_supported_dtype_shape,
+)
+from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
+from vllm.utils.platform_utils import num_compute_units
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+
+class AllSparkLinearKernel(MPLinearKernel):
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 80
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        if c.has_g_idx:
+            return False, "Act reordering currently not supported by AllSpark"
+
+        if c.zero_points:
+            return False, "Zero points currently not supported by AllSpark"
+
+        return check_allspark_supported_dtype_shape(
+            c.partition_weight_shape[0],  # in_features
+            c.partition_weight_shape[1],  # out_features
+            c.group_size,
+            c.weight_type,
+            c.act_type,
+        )
+
+    # note assumes that
+    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
+    #  `weight_scale` is: {input_dim = 0, output_dim = 1}
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        device = getattr(layer, self.w_q_name).device
+        c = self.config
+
+        # prepare the parameters required for the kernel
+        properties = torch.cuda.get_device_properties(device.index)
+        sm_count = num_compute_units(device.index)
+        sm_version = properties.major * 10 + properties.minor
+        gemm_args = {}
+        gemm_args["sm_count"] = sm_count
+        gemm_args["sm_version"] = sm_version
+
+        self.gemm_args = gemm_args
+
+        # transform param weight, scale
+        old_weight_param = getattr(layer, self.w_q_name)
+        old_scale_param = getattr(layer, self.w_s_name)
+
+        assert isinstance(old_weight_param, BasevLLMParameter)
+        permute_param_layout_(old_weight_param, input_dim=0, output_dim=1, packed_dim=0)
+
+        assert isinstance(old_scale_param, BasevLLMParameter)
+        permute_param_layout_(old_scale_param, input_dim=0, output_dim=1)
+
+        # unpack weight from K / 4 x N int32 to K x N uint8
+        new_weight_param = torch.nn.Parameter(
+            old_weight_param.data, requires_grad=False
+        )
+        new_weight_param.data = (
+            new_weight_param.data.t().contiguous().view(dtype=torch.uint8)
+        )
+        new_weight_param.data = new_weight_param.data.t().contiguous()
+
+        new_scale_param = torch.nn.Parameter(old_scale_param.data, requires_grad=False)
+
+        # reorder K x N weight as N32K16 format for Ampere W8A16
+        new_weight_param.data, new_scale_param.data, _ = ops.allspark_repack_weight(
+            new_weight_param.data, new_scale_param.data, None, c.zero_points
+        )
+
+        replace_parameter(layer, self.w_q_name, new_weight_param.data)
+        replace_parameter(layer, self.w_s_name, new_scale_param.data)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        c = self.config
+        gemm_args = self.gemm_args
+        w_q, w_s, _, _ = self._get_weight_params(layer)
+
+        reshaped_x = x.reshape(-1, x.shape[-1])
+        out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
+
+        output = ops.allspark_w8a16_gemm(
+            a=reshaped_x,
+            b_qweight=w_q,
+            b_scales=w_s,
+            b_qzeros=None,
+            n=c.partition_weight_shape[1],
+            group_size=c.group_size,
+            sm_count=gemm_args["sm_count"],
+            sm_version=gemm_args["sm_version"],
+            CUBLAS_M_THRESHOLD=ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
+            has_zp=c.zero_points,
+            n32k16_reorder=True,
+        )
+
+        if bias is not None:
+            output.add_(bias)  # In-place add
+
+        return output.reshape(out_shape)

vllm/model_executor/kernels/linear/mixed_precision/conch.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from importlib.util import find_spec
+from typing import Final
+
+import torch
+
+from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
+from vllm.scalar_type import scalar_types
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+_CONCH_SUPPORTED_WEIGHT_TYPES: Final = [
+    scalar_types.uint4,
+    scalar_types.uint8,
+    scalar_types.uint4b8,
+    scalar_types.uint8b128,
+]
+_CONCH_SUPPORTED_GROUP_SIZES: Final = [-1, 128]
+
+
+class ConchLinearKernel(MPLinearKernel):
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 80
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        if c.weight_type not in _CONCH_SUPPORTED_WEIGHT_TYPES:
+            error_msg = (
+                f"Weight type ({c.weight_type}) not supported by "
+                "ConchLinearKernel, supported types are: "
+                f"{_CONCH_SUPPORTED_WEIGHT_TYPES}"
+            )
+            return False, error_msg
+
+        if c.group_size not in _CONCH_SUPPORTED_GROUP_SIZES:
+            error_msg = (
+                f"Group size ({c.group_size}) not supported by "
+                "ConchLinearKernel, supported group sizes are: "
+                f"{_CONCH_SUPPORTED_GROUP_SIZES}"
+            )
+            return False, error_msg
+
+        if c.has_g_idx:
+            return (
+                False,
+                "Activation reordering (g_idx) is not supported by ConchLinearKernel",
+            )
+
+        if find_spec("conch") is None:
+            error_msg = (
+                "conch-triton-kernels is not installed, please "
+                "install it via `pip install conch-triton-kernels` "
+                "and try again!"
+            )
+            return False, error_msg
+
+        return True, None
+
+    # note assumes that
+    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
+    #  `weight_scale` is: {input_dim = 0, output_dim = 1}
+    #  `weight_zero_point` is: {input_dim = 1, output_dim = 0, packed_dim = 0}
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        def transform_w_q(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
+            x.data = x.data.contiguous()
+            return x
+
+        def transform_w_s(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1)
+            x.data = x.data.contiguous()
+            return x
+
+        def transform_w_zp(x):
+            # Zero points are stored PACKED as [N//pack_factor, K//G]
+            # The Conch kernel expects UNPACKED zeros: [K//G, N]
+            # We need to unpack and reorder
+            assert isinstance(x, BasevLLMParameter)
+            packed = x.data  # shape: [N//pack_factor, K//G], dtype: int32
+
+            # Determine packing based on weight bit width
+            size_bits = self.config.weight_type.size_bits
+            pack_factor = 32 // size_bits  # 8 for 4-bit, 4 for 8-bit
+            mask = (1 << size_bits) - 1  # 0xF for 4-bit, 0xFF for 8-bit
+
+            n_packed, k_groups = packed.shape
+            n_full = n_packed * pack_factor
+
+            # Unpack using vectorized bitwise ops
+            # shifts = [0, size_bits, 2*size_bits, ...] for each packed position
+            shifts = torch.arange(
+                0, 32, size_bits, dtype=torch.int32, device=packed.device
+            )
+            # packed: [N//pack_factor, K//G] -> [N//pack_factor, K//G, 1]
+            # shifts: [pack_factor] -> [1, 1, pack_factor]
+            # Result: [N//pack_factor, K//G, pack_factor]
+            unpacked = (packed.unsqueeze(-1) >> shifts) & mask
+
+            # Permute to [K//G, N//pack_factor, pack_factor] then reshape to [K//G, N]
+            unpacked = unpacked.permute(1, 0, 2).reshape(k_groups, n_full)
+
+            x.data = unpacked.to(torch.uint8).contiguous()
+
+            # Update metadata - zeros are no longer packed
+            if hasattr(x, "_input_dim"):
+                x._input_dim = 0
+            if hasattr(x, "_output_dim"):
+                x._output_dim = 1
+            if hasattr(x, "_packed_factor"):
+                x._packed_factor = 1
+            return x
+
+        self._transform_param(layer, self.w_q_name, transform_w_q)
+        self._transform_param(layer, self.w_s_name, transform_w_s)
+        if self.config.zero_points:
+            self._transform_param(layer, self.w_zp_name, transform_w_zp)
+        elif self.w_zp_name is not None:
+            layer.register_parameter(self.w_zp_name, None)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        from conch.ops.quantization.gemm import mixed_precision_gemm
+
+        w_q, w_s, w_zp, _ = self._get_weight_params(layer)
+
+        # Map channelwise group_size=-1 to the actual input dimension K.
+        # The conch kernel computes stride_mul = block_k / group_size;
+        # passing -1 produces a negative stride that reads out-of-bounds
+        # scale values for all K-blocks after the first.
+        group_size = self.config.group_size
+        if group_size == -1:
+            group_size = x.shape[-1]
+
+        x_2d = x.reshape(-1, x.shape[-1])
+        out_shape = x.shape[:-1] + (self.config.partition_weight_shape[1],)
+
+        output = mixed_precision_gemm(
+            x=x_2d,
+            w_q_packed=w_q.data,
+            w_s=w_s.data,
+            w_zp=w_zp.data if w_zp is not None else None,
+            weight_size_bits=self.config.weight_type.size_bits,
+            weight_bias=self.config.weight_type.bias,
+            group_size=group_size,
+        )
+
+        if bias is not None:
+            output.add_(bias)  # In-place add
+
+        return output.reshape(out_shape)

vllm/model_executor/kernels/linear/mixed_precision/cpu.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm import envs
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    pack_quantized_values_into_int32,
+    unpack_quantized_values_into_int32,
+)
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.scalar_type import scalar_types
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+_CPUWNA16_SUPPORTED_QUANT_TYPES = (scalar_types.uint4, scalar_types.uint4b8)
+
+
+class CPUWNA16LinearKernel(MPLinearKernel):
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return -1
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        if not current_platform.is_cpu():
+            return False, "CPUWNA16 only supported on CPU"
+
+        if c.weight_type not in _CPUWNA16_SUPPORTED_QUANT_TYPES:
+            return (
+                False,
+                f"Quant type ({c.weight_type}) not supported by "
+                "CPUWNA16, supported types are: "
+                f"{_CPUWNA16_SUPPORTED_QUANT_TYPES}",
+            )
+
+        if c.group_size != -1 and c.group_size % 2 != 0:
+            return (
+                False,
+                f"Group size ({c.group_size}) not supported by "
+                "CPUWNA16, supported group sizes are multiples of 2",
+            )
+
+        if c.partition_weight_shape[0] % 32 != 0:
+            return (
+                False,
+                f"Input size ({c.partition_weight_shape[0]}) not supported by "
+                "CPUWNA16, supported sizes are multiples of 32",
+            )
+
+        if c.partition_weight_shape[1] % 32 != 0:
+            return (
+                False,
+                f"Output size ({c.partition_weight_shape[1]}) not supported by "
+                "CPUWNA16, supported sizes are multiples of 32",
+            )
+
+        return True, None
+
+    # note assumes that
+    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
+    #  `weight_scale`  is: {input_dim = 0, output_dim = 1}
+    #  `weight_zp`     is: {input_dim = 0, output_dim = 1, packed_dim = 1}
+    def _process_gptq_weights_w4a16(self, layer: torch.nn.Module):
+        packed_weight = getattr(layer, self.w_q_name)
+        bits = self.config.weight_type.mantissa
+        pack_factor = 32 // bits
+        p_w_k, _ = packed_weight.size()
+        input_size = p_w_k * pack_factor
+        isa_hint = _get_isa_hint(getattr(layer, self.w_s_name).dtype)
+        layer.isa_hint = isa_hint
+
+        # convert input dim packed to output dim packed
+        weight = unpack_quantized_values_into_int32(
+            packed_weight, self.config.weight_type, 0
+        )
+        weight = pack_quantized_values_into_int32(weight, self.config.weight_type, 1)
+        # make 16 output channel as a block and transpose to the make
+        # the block contiguous
+        weight = (
+            weight.view(input_size, -1, 16 // pack_factor)
+            .permute(1, 0, 2)
+            .reshape(-1, input_size * 16 // pack_factor)
+            .contiguous()
+        )
+        getattr(layer, self.w_q_name).data = weight
+
+    # note assumes that
+    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
+    #  `weight_scale`  is: {input_dim = 0, output_dim = 1}
+    #  `weight_zp`     is: {input_dim = 0, output_dim = 1, packed_dim = 1}
+    def _process_gptq_weights_w4a8(self, layer: torch.nn.Module):
+        packed_weight = getattr(layer, self.w_q_name)
+        scales = getattr(layer, self.w_s_name)
+        group_num = scales.data.size(0)
+        zp_output_size = scales.data.size(1) // 8
+        if self.config.zero_points:
+            assert self.w_zp_name
+            packed_zp = getattr(layer, self.w_zp_name)
+        else:
+            # w4a8 kernel always requires zp, allocate a fake zp
+            assert self.w_zp_name
+            packed_zp = torch.nn.Parameter(
+                torch.ones(group_num, zp_output_size, dtype=torch.int32) * -2004318072,
+                requires_grad=False,
+            )
+            setattr(layer, self.w_zp_name, packed_zp)
+
+        # FIXME: some bugs in convert_weight_packed_scale_zp with GPTQ format,
+        # repack to AWQ weight
+        weight = unpack_quantized_values_into_int32(
+            packed_weight, self.config.weight_type, 0
+        )
+        input_size, output_size = weight.size()
+        weight = weight.view(input_size, output_size // 8, 8)
+        weight = weight[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(input_size, output_size)
+        weight = pack_quantized_values_into_int32(
+            weight, self.config.weight_type, 1
+        ).contiguous()
+
+        zp = unpack_quantized_values_into_int32(packed_zp, self.config.weight_type, 1)
+        zp = zp.view(group_num, output_size // 8, 8)
+        zp = zp[:, :, (0, 2, 4, 6, 1, 3, 5, 7)].reshape(group_num, output_size)
+        zp = pack_quantized_values_into_int32(
+            zp, self.config.weight_type, 1
+        ).contiguous()
+
+        blocked_w, blocked_zp, blocked_s = ops.convert_weight_packed_scale_zp(
+            weight,
+            zp,
+            scales.data,
+            ops.CPUQuantAlgo.AWQ,
+        )
+
+        if layer.bias is not None:
+            layer.bias.data = layer.bias.float()
+
+        packed_weight.data = blocked_w
+        scales.data = blocked_s
+        packed_zp.data = blocked_zp
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        if (not self.config.zero_points) and (self.w_zp_name is not None):
+            setattr(layer, self.w_zp_name, None)
+
+        if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
+            setattr(layer, self.w_gidx_name, None)
+
+        weights = getattr(layer, self.w_q_name)
+        # Require GPTQ pack format
+        assert weights.input_dim == weights.packed_dim
+
+        # Weights in CT format is [output_size, input_size]
+        if weights.input_dim == 1:
+            weights.data = weights.t()
+
+        # Scales in CT format is [output_size, group_num]
+        scales = getattr(layer, self.w_s_name)
+        if scales.output_dim == 0:
+            scales.data = scales.t().contiguous()
+
+        # Zero points in CT format is [output_size, group_num]
+        # Zero points in awq_marlin format is [output_size, group_num]
+        if self.config.zero_points:
+            assert self.w_zp_name
+            zp = getattr(layer, self.w_zp_name)
+            if zp.output_dim == 0:
+                zp.data = zp.t().contiguous()
+
+        supports_amx = torch.cpu._is_amx_tile_supported()
+        supports_riscv = current_platform.get_cpu_architecture() == CpuArchEnum.RISCV
+        layer.use_w4a8 = (
+            envs.VLLM_CPU_INT4_W4A8
+            and not self.config.has_g_idx
+            and self.config.act_type == torch.bfloat16
+            and (supports_amx or supports_riscv)
+        )
+        # layer.use_w4a8 = False
+        # AWQ format will be converted to GPTQ format in `AutoAWQMarlinLinearMethod`
+        if layer.use_w4a8:
+            self._process_gptq_weights_w4a8(layer)
+        else:
+            self._process_gptq_weights_w4a16(layer)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
+        if layer.use_w4a8:
+            x = ops.int4_scaled_mm_cpu(
+                x=x,
+                w=w_q,
+                w_zeros=w_zp,
+                w_scales=w_s,
+                bias=bias,
+            )
+        else:
+            x = ops.cpu_gemm_wna16(
+                input=x,
+                q_weight=w_q,
+                scales=w_s,
+                zeros=w_zp,
+                g_idx=w_gidx,
+                bias=bias,
+                pack_factor=8,  # 32 // 4
+                isa_hint=layer.isa_hint,
+            )
+        return x
+
+
+def _get_isa_hint(dtype: torch.dtype) -> str:
+    supports_amx = torch.cpu._is_amx_tile_supported()
+    if supports_amx and dtype in (torch.bfloat16,):
+        return "amx"
+    elif current_platform.get_cpu_architecture() == CpuArchEnum.RISCV:
+        return "rvv"
+    else:
+        return "vec"

vllm/model_executor/kernels/linear/mixed_precision/cutlass.py

@@ -0,0 +1,131 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    convert_bf16_scales_to_fp8,
+    convert_packed_uint4b8_to_signed_int4_inplace,
+)
+from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
+from vllm.platforms import current_platform
+from vllm.scalar_type import scalar_types
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+
+class CutlassW4A8LinearKernel(MPLinearKernel):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # dynamic per-tok fp8 activation quantization
+        self.quant_fp8 = QuantFP8(static=False, group_shape=GroupShape.PER_TOKEN)
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 90
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "CUTLASS only supported on CUDA"
+
+        if not current_platform.is_device_capability(90):
+            return False, "CUTLASS W4A8 requires compute capability of 90 (Hopper)"
+
+        if c.act_type != torch.float8_e4m3fn:
+            return False, "CUTLASS W4A8 only supports FP8 (e4m3) activations"
+
+        if c.has_g_idx:
+            return False, "Act reordering not supported by CUTLASS W4A8"
+
+        if c.zero_points:
+            return False, "Zero points not supported by CUTLASS W4A8"
+
+        if c.weight_type != scalar_types.int4:
+            return (
+                False,
+                f"Quant type ({c.weight_type}) not supported by "
+                "CUTLASS W4A8, only supported int4",
+            )
+
+        if c.group_size != 128:
+            return False, "Only group_size 128 is supported"
+
+        in_features, out_features = c.partition_weight_shape
+        if in_features % 128 or out_features % 128:
+            return (
+                False,
+                f"K and N must be divisible by 128, got {c.partition_weight_shape}",
+            )
+
+        if c.out_type != torch.bfloat16:
+            return (
+                False,
+                f"Only bfloat16 output type currently supportedgot {c.out_type=}",
+            )
+
+        return True, None
+
+    # note assumes that
+    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
+    #  `weight_scale`  is: {input_dim = 0, output_dim = 1}
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        def transform_w_q(x):
+            assert isinstance(x, BasevLLMParameter)
+            convert_packed_uint4b8_to_signed_int4_inplace(x.data)
+            torch.accelerator.synchronize()
+            permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
+            x.data = ops.cutlass_encode_and_reorder_int4b(x.data.t().contiguous().t())
+            return x
+
+        def transform_w_s(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1)
+            x.data = x.data.contiguous().to(torch.float8_e4m3fn)
+            x.data = ops.cutlass_pack_scale_fp8(x.data)
+            return x
+
+        w_s = getattr(layer, self.w_s_name)
+        fp8_scales, chan_scales = convert_bf16_scales_to_fp8(self.quant_fp8, w_s.data)
+        w_s.data = fp8_scales
+
+        # register per-channel scales
+        layer.register_parameter(
+            "weight_chan_scale", torch.nn.Parameter(chan_scales, requires_grad=False)
+        )
+
+        # Encode/reorder weights and pack scales
+        self._transform_param(layer, self.w_q_name, transform_w_q)
+        self._transform_param(layer, self.w_s_name, transform_w_s)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        c = self.config
+        w_q, w_s, _, _ = self._get_weight_params(layer)
+        w_ch_s = layer.weight_chan_scale
+
+        x_2d = x.reshape(-1, x.shape[-1])
+        out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
+
+        x_2d, act_scales = self.quant_fp8(x_2d)
+        output = ops.cutlass_w4a8_mm(
+            a=x_2d,
+            b_q=w_q,
+            b_group_scales=w_s,
+            b_group_size=c.group_size,
+            a_token_scales=act_scales,
+            b_channel_scales=w_ch_s,
+        )
+
+        if bias is not None:
+            output.add_(bias)  # In-place add
+
+        return output.reshape(out_shape)

vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.platforms import CpuArchEnum, current_platform
+from vllm.scalar_type import scalar_types
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+
+# This implementation is for the KleidiAI-accelerated w4a8int quantization
+# scheme on Arm CPUs:
+# torch.ops.aten._dyn_quant_matmul_4bit performs dynamic quantized matmul
+#   it takes:
+#       - int4 weights packed along with bias/scales by
+#         torch.ops.aten._dyn_quant_pack_4bit_weight
+#       - float32/bfloat16 activations
+#   then it leverages KleidiAI ukernels that:
+#       - dynamically quantize the activations to int8
+#       - unpack the int4 weights to int8
+#       - perform int8 x int8 -> int32 matmul
+#       - dequantize the int32 output to float32/bfloat16 outputs
+class Dynamic4bitLinearKernel(MPLinearKernel):
+    SUPPORTED_QUANT_TYPES = [scalar_types.int4]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 1
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        if not current_platform.is_cpu():
+            return False, "Only CPU is supported"
+        if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
+            return False, f"Unsupported quant type {c.weight_type}"
+        if (
+            current_platform.get_cpu_architecture() == CpuArchEnum.ARM
+            and c.act_type
+            not in [
+                torch.float32,
+                torch.bfloat16,
+                torch.float16,
+            ]
+        ):
+            return (
+                False,
+                "Dynamic4bitLinearKernel on Arm requires Float32 or"
+                " BFloat16 or Float16 activations",
+            )
+        if c.full_weight_shape[0] % c.group_size != 0:
+            return (
+                False,
+                f"Group size ({c.group_size}) does not evenly divide"
+                " the number of input features "
+                f"({c.full_weight_shape[0]})",
+            )
+        if current_platform.get_cpu_architecture() == CpuArchEnum.ARM:
+            try:
+                # Attempt to retrieve the operation
+                _ = torch.ops.aten._dyn_quant_matmul_4bit
+            except AttributeError:
+                return (
+                    False,
+                    f"PyTorch {torch.__version__} does not support"
+                    " _dyn_quant_matmul_4bit. Install a newer version",
+                )
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        c = self.config
+        packed_weight = getattr(layer, self.w_q_name)
+        packed_weight = packed_weight.add(8)
+        uint8_packed = (packed_weight[::, 1::2] << 4 | packed_weight[::, ::2]).to(
+            torch.uint8
+        )
+
+        scales = getattr(layer, self.w_s_name)
+        block_size = c.group_size
+
+        # Handle scaling factors for partitioned weights
+        if block_size == c.partition_weight_shape[0]:
+            scales = scales.to(
+                torch.float32
+            )  # Float32 & Bfloat16 variants requires float32 scales
+            scales = scales.view(-1, 1)  # Channel-wise scales
+            if layer.bias is not None:
+                # Float32 & Bfloat16 variants requires float32 bias
+                replace_parameter(
+                    layer,
+                    "bias",
+                    torch.nn.Parameter(
+                        layer.bias.to(torch.float32), requires_grad=False
+                    ),
+                )
+        else:
+            # KleidiAI kernel requires bfloat16 scales with groupwise scheme
+            scales = scales.to(torch.bfloat16)
+
+        # Repack weights as per kernel requirement
+        w = torch.ops.aten._dyn_quant_pack_4bit_weight(
+            uint8_packed,
+            scales,
+            layer.bias,
+            block_size,
+            c.partition_weight_shape[0],
+            c.partition_weight_shape[1],
+        )
+        replace_parameter(
+            layer, self.w_q_name, torch.nn.Parameter(w, requires_grad=False)
+        )
+        setattr(layer, self.w_s_name, None)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        # PyTorch / KleidiAI kernels natively support the following configs:
+        # - channelwise with bfloat16 / float32 activations
+        # - groupwise with float32 activations
+        # To support:
+        # - groupwise with bfloat16/float16 activations: we need to upcast
+        #   activations to float32 before matmul and downcast back to bfloat16/float16
+        # - channelwise with float16 activations, we need to upcast activations to
+        #   float32 before matmul and downcast back to float16
+        # Note: these activations will be dynamically quantized to int8 by the kernel.
+
+        c = self.config
+        is_groupwise = c.group_size != c.partition_weight_shape[0]
+        # dtype of activations before they get dynamically quantized to int8
+        original_pre_quant_act_dtype = x.dtype
+        pre_quant_act_dtype = original_pre_quant_act_dtype
+        if (
+            is_groupwise and pre_quant_act_dtype == torch.bfloat16
+        ) or pre_quant_act_dtype == torch.float16:
+            pre_quant_act_dtype = torch.float32
+
+        x_2d = x.reshape(-1, x.shape[-1])
+        if pre_quant_act_dtype != original_pre_quant_act_dtype:
+            x_2d = x_2d.to(pre_quant_act_dtype)
+
+        out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
+
+        w_q = getattr(layer, self.w_q_name)
+        output = torch.ops.aten._dyn_quant_matmul_4bit(
+            x_2d,
+            w_q,
+            c.group_size,
+            c.partition_weight_shape[0],
+            c.partition_weight_shape[1],
+        ).reshape(out_shape)
+
+        if pre_quant_act_dtype != original_pre_quant_act_dtype:
+            output = output.to(original_pre_quant_act_dtype)
+        return output

vllm/model_executor/kernels/linear/mixed_precision/exllama.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    pack_quantized_values_into_int32,
+)
+from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
+from vllm.platforms import current_platform
+from vllm.scalar_type import scalar_types
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+
+class ExllamaLinearKernel(MPLinearKernel):
+    SUPPORTED_QUANT_TYPES = [scalar_types.uint4b8, scalar_types.uint8b128]
+    # In theory supports `scalar_types.uint2b2, scalar_types.uint3b4` too but
+    # currently untested so not added to the list
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 60
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda_alike():
+            return (
+                False,
+                "Exllama is only supported on CUDA and ROCm",
+            )
+
+        if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]:
+            return (
+                False,
+                "Act reordering currently not supported by Exllama, "
+                "when the input features are partitioned across "
+                "devices",
+            )
+
+        if c.partition_weight_shape[1] % (32 // c.weight_type.size_bits) != 0:
+            return (
+                False,
+                "Output features must be a multiple of the pack "
+                "factor (32 / num_bits) so that we can correctly "
+                "pack the zero points",
+            )
+
+        if c.act_type != torch.float16:
+            return False, "Exllama only supports float16 activations"
+
+        if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
+            return (
+                False,
+                f"Quant type ({c.weight_type}) not supported by "
+                "Exllama, supported types are: "
+                f"{cls.SUPPORTED_QUANT_TYPES}",
+            )
+
+        if c.group_size <= 0:
+            return (
+                False,
+                f"Group size ({c.group_size}) must be positive, "
+                "Exllama does not support channelwise quantization",
+            )
+
+        if c.full_weight_shape[0] % c.group_size != 0:
+            return (
+                False,
+                f"Group size ({c.group_size}) does not evenly divide"
+                " the number of input features "
+                f"({c.full_weight_shape[0]})",
+            )
+
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        c = self.config
+
+        # For Exllama, we need to set a zero-point tensor if there is not one
+        if not c.zero_points:
+            self.w_zp_name = "qzeros"
+            device = getattr(layer, self.w_q_name).device
+            groups = c.partition_weight_shape[0] // c.group_size
+            out_features = c.partition_weight_shape[1]
+
+            if c.weight_type.has_bias():
+                # if the type has a bias we have to create a zeros tensor that
+                # contains the bias values repeated for each group (-1 due to
+                # a bug in the original GPTQ checkpoint format leading to
+                # exllama kernel adding 1 to the zero points during inference)
+                # Documentation of the bug can be found here:
+                #  https://garden.danieldk.eu/GPTQ-Checkpoint-Format
+                zeros = torch.full(
+                    (groups, out_features),
+                    c.weight_type.bias - 1,
+                    dtype=torch.int32,
+                    device=device,
+                )
+            else:
+                raise NotImplementedError(
+                    "A 0 zero-point is not supported by Exllama due to "
+                    "a bug in the original GPTQ checkpoint format leading to "
+                    "exllama kernel adding 1 to the zero points during "
+                    "inference"
+                )
+            zeros = pack_quantized_values_into_int32(zeros, c.weight_type, packed_dim=1)
+            setattr(
+                layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False)
+            )
+
+        if c.has_g_idx:
+
+            def transform_w_g_idx(x):
+                # Exllama wants the permutation array instead of the group
+                # indices
+                return torch.argsort(x).to(torch.int)
+
+            self._transform_param(layer, self.w_gidx_name, transform_w_g_idx)  # type: ignore
+        else:
+            self.w_gidx_name = "g_idx"
+            empty_g_idx = torch.nn.Parameter(
+                torch.empty((0,), dtype=torch.int, device=device), requires_grad=False
+            )
+            setattr(layer, self.w_gidx_name, empty_g_idx)
+
+        def transform_w_q(x):
+            assert isinstance(x, BasevLLMParameter)
+            assert self.w_gidx_name is not None
+            g_idx = getattr(layer, self.w_gidx_name)
+
+            permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
+            x_cont = x.data.contiguous()
+            ops.gptq_shuffle(x_cont, g_idx, c.weight_type.size_bits)
+            return x_cont
+
+        def transform_w_s(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1)
+            x.data = x.data.contiguous()
+            return x.to(dtype=c.act_type)
+
+        # Repack weights and scales for Machete
+        self._transform_param(layer, self.w_q_name, transform_w_q)
+        self._transform_param(layer, self.w_s_name, transform_w_s)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        c = self.config
+
+        x_2d = x.reshape(-1, x.shape[-1])
+        out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
+
+        w_q, w_s, w_zp, w_g_idx = self._get_weight_params(layer)
+
+        # gptq_gemm supports GPTQv2 format by passing use_v2_format=True.
+        # However, the MPLinearLayerConfig doesn't contain format info.
+        # So hardcode GPTQv1 format here, to keep its behavior unchanged.
+        use_v2_format = False
+
+        assert w_zp is not None, "Zero points are required by Exllama"
+        assert w_g_idx is not None, "Group index is required by Exllama"
+        output = ops.gptq_gemm(
+            x_2d, w_q, w_zp, w_s, w_g_idx, True, use_v2_format, c.weight_type.size_bits
+        )
+
+        if bias is not None:
+            output.add_(bias)
+        return output.reshape(out_shape)

vllm/model_executor/kernels/linear/mixed_precision/humming.py

@@ -0,0 +1,64 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Humming GEMM as a mixed-precision WNA16Int linear kernel."""
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_humming
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+
+class HummingLinearKernel(MPLinearKernel):
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 75
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "Humming is only supported on CUDA"
+        if not has_humming():
+            return False, "Humming is not installed"
+        if c.has_g_idx:
+            return False, "Humming does not support act-order (g_idx)"
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        from vllm.model_executor.layers.quantization.utils.humming_utils import (
+            convert_linear_layer_to_humming_standard,
+            prepare_humming_layer,
+        )
+
+        name_map = {"weight": self.w_q_name, "weight_scale": self.w_s_name}
+        group_size = self.config.group_size
+        quant_config = {
+            "quant_method": "humming",
+            "dtype": "int" + str(self.config.weight_type.size_bits),
+            "group_size": 0 if group_size == -1 else group_size,
+        }
+
+        if self.config.zero_points:
+            assert self.w_zp_name is not None
+            name_map["zero_point"] = self.w_zp_name
+            quant_config["has_zero_point"] = True
+
+        convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
+        prepare_humming_layer(layer, quant_config)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        from vllm.utils.humming import HummingMethod
+
+        flatten_inputs = x.view(-1, x.size(-1))
+        output = HummingMethod.forward_layer(
+            layer=layer,
+            inputs=flatten_inputs,
+            compute_config=layer.compute_config,
+        )
+        return output.view(*x.shape[:-1], output.size(-1))

vllm/model_executor/kernels/linear/mixed_precision/machete.py

@@ -0,0 +1,159 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from functools import partial
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.utils.machete_utils import (
+    check_machete_supports_shape,
+    query_machete_supported_group_sizes,
+    query_machete_supported_quant_types,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    pack_quantized_values_into_int32,
+    unpack_quantized_values_into_int32,
+)
+from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
+from vllm.platforms import current_platform
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+
+class MacheteLinearKernel(MPLinearKernel):
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 90
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        # Machete uses CUTLASS, so it can only be compatible with Nvidia
+        if not current_platform.is_cuda():
+            return False, "Machete only supported on CUDA"
+
+        if not current_platform.is_device_capability(90):
+            return False, "Machete requires compute capability of 90 (Hopper)"
+
+        if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]:
+            return (
+                False,
+                "Act reordering currently not supported by Machete, "
+                "when the input features are partitioned across "
+                "devices",
+            )
+
+        if c.weight_type not in query_machete_supported_quant_types(c.zero_points):
+            return (
+                False,
+                f"Quant type ({c.weight_type}) not supported by "
+                "Machete, supported types are: "
+                f"{query_machete_supported_quant_types(c.zero_points)}",
+            )
+
+        if c.group_size not in query_machete_supported_group_sizes(c.act_type):
+            return (
+                False,
+                f"Group size ({c.group_size}) not supported by "
+                "Machete, supported group sizes are: "
+                f"{query_machete_supported_group_sizes(c.act_type)}",
+            )
+
+        return check_machete_supports_shape(
+            c.partition_weight_shape[0], c.partition_weight_shape[1]
+        )
+
+    # note assumes that
+    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
+    #  `weight_scale`  is: {input_dim = 0, output_dim = 1}
+    #  `weight_zp`     is: {input_dim = 0, output_dim = 1, packed_dim = 1}
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        c = self.config
+
+        if c.has_g_idx:
+            assert self.w_gidx_name is not None
+            perm = torch.argsort(getattr(layer, self.w_gidx_name)).to(torch.int)
+
+            self.act_perm = lambda x: x[:, perm]
+            # use `ops.permute_cols` if possible
+            if (
+                c.act_type in [torch.float16, torch.bfloat16]
+                and c.partition_weight_shape[0] % 8 == 0
+            ):
+                self.act_perm = partial(ops.permute_cols, perm=perm)
+
+        def transform_w_q(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
+            if c.has_g_idx:
+                x_unpacked = unpack_quantized_values_into_int32(
+                    x.data, c.weight_type, packed_dim=0
+                )
+                x_perm = x_unpacked[perm, :]
+                x.data = pack_quantized_values_into_int32(
+                    x_perm, c.weight_type, packed_dim=0
+                )
+            x.data = ops.machete_prepack_B(
+                x.data.t().contiguous().t(),
+                a_type=c.act_type,
+                b_type=c.weight_type,
+                group_scales_type=c.act_type,
+            )
+            return x
+
+        def transform_w_s(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1)
+            x.data = x.data.contiguous()
+            return x
+
+        def transform_w_zp(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=1)
+            x_unpacked = unpack_quantized_values_into_int32(
+                x.data, c.weight_type, packed_dim=1
+            )
+            w_s = getattr(layer, self.w_s_name).data
+            # pre-apply scales to zero-points
+            x.data = (-1.0 * w_s * (x_unpacked.to(w_s.dtype))).contiguous()
+            return x
+
+        # Repack weights and scales for Machete
+        self._transform_param(layer, self.w_q_name, transform_w_q)
+        self._transform_param(layer, self.w_s_name, transform_w_s)
+        if c.zero_points:
+            self._transform_param(layer, self.w_zp_name, transform_w_zp)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        c = self.config
+        w_q, w_s, w_zp, _ = self._get_weight_params(layer)
+
+        x_2d = x.reshape(-1, x.shape[-1])
+        out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
+
+        if c.has_g_idx:
+            x_2d = self.act_perm(x_2d)
+
+        if c.zero_points:
+            assert w_zp is not None
+        else:
+            w_zp = None
+
+        output = ops.machete_mm(
+            a=x_2d,
+            b_q=w_q,
+            b_type=c.weight_type,
+            b_group_zeros=w_zp,
+            b_group_scales=w_s,
+            b_group_size=c.group_size,
+        )
+
+        if bias is not None:
+            output.add_(bias)  # In-place add
+
+        return output.reshape(out_shape)

vllm/model_executor/kernels/linear/mixed_precision/marlin.py

@@ -0,0 +1,245 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.utils.marlin_utils import (
+    MARLIN_SUPPORTED_GROUP_SIZES,
+    apply_gptq_marlin_linear,
+    check_marlin_supports_shape,
+    marlin_act_int8_process_scales,
+    marlin_is_k_full,
+    marlin_make_empty_g_idx,
+    marlin_make_workspace_new,
+    marlin_pad_dim,
+    marlin_pad_qweight,
+    marlin_pad_scales,
+    marlin_padded_nk,
+    marlin_permute_bias,
+    marlin_permute_scales,
+    marlin_sort_g_idx,
+    marlin_zero_points,
+    query_marlin_supported_quant_types,
+    unpack_cols,
+)
+from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
+from vllm.platforms import current_platform
+from vllm.scalar_type import scalar_types
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+
+class MarlinLinearKernel(MPLinearKernel):
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return 75
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        # Marlin uses inline PTX, so it can only be compatible with Nvidia
+        if not current_platform.is_cuda():
+            return False, "Marlin only supported on CUDA"
+
+        quant_types = query_marlin_supported_quant_types(c.zero_points)
+        if c.weight_type not in quant_types:
+            return (
+                False,
+                f"Quant type ({c.weight_type}) not supported by"
+                f"  Marlin, supported types are: {quant_types}",
+            )
+
+        if c.group_size not in MARLIN_SUPPORTED_GROUP_SIZES:
+            return (
+                False,
+                f"Group size ({c.group_size}) not supported by "
+                "Marlin, supported group sizes are: "
+                f"{MARLIN_SUPPORTED_GROUP_SIZES}",
+            )
+
+        if c.has_g_idx:
+            # Act-order couples K to the full-model group layout, so tile
+            # padding is not supported; keep the strict shape check.
+            return check_marlin_supports_shape(
+                c.partition_weight_shape[1],  # out_features
+                c.partition_weight_shape[0],  # in_features
+                c.full_weight_shape[0],  # in_features
+                c.group_size,
+            )
+
+        # A group straddling TP ranks cannot be fixed by padding.
+        if (
+            c.group_size != -1
+            and c.group_size < c.full_weight_shape[0]
+            and c.partition_weight_shape[0] % c.group_size != 0
+        ):
+            return False, (
+                f"in_features per partition {c.partition_weight_shape[0]} is "
+                f"not divisible by group_size = {c.group_size}."
+            )
+
+        # Tile misalignment is fixed by zero-padding at weight prep.
+        return True, None
+
+    # note assumes that
+    #  `weight_packed` is: {input_dim = 0, output_dim = 1, packed_dim = 0}
+    #  `weight_scale` is: {input_dim = 0, output_dim = 1}
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        device = getattr(layer, self.w_q_name).device
+        c = self.config
+        is_a_8bit = c.act_type is not None and c.act_type.itemsize == 1
+
+        if is_a_8bit:
+            assert c.weight_type == scalar_types.uint4b8, (
+                "W8A8 is not supported by marlin kernel."
+            )
+
+        if c.act_type == torch.float8_e4m3fn:
+            ops.marlin_int4_fp8_preprocess(getattr(layer, self.w_q_name), inplace=True)
+            getattr(layer, self.w_s_name).data = (
+                getattr(layer, self.w_s_name).data * 512
+            )
+
+        row_parallel = c.partition_weight_shape[0] != c.full_weight_shape[0]
+        self.is_k_full = marlin_is_k_full(c.has_g_idx, row_parallel)
+
+        size_k, size_n = c.partition_weight_shape
+        if c.has_g_idx:
+            # Act-order shapes were strictly validated in can_implement.
+            padded_n, padded_k = size_n, size_k
+        else:
+            padded_n, padded_k = marlin_padded_nk(size_n, size_k, c.group_size)
+
+        # Allocate marlin workspace.
+        self.workspace = marlin_make_workspace_new(device)
+
+        # Default names since marlin requires empty parameters for these,
+        # TODO: remove this requirement from marlin (allow optional tensors)
+        if self.w_gidx_name is None:
+            self.w_gidx_name = "g_idx"
+        if self.w_zp_name is None:
+            self.w_zp_name = "w_zp"
+
+        def transform_w_q(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
+            x.data = ops.gptq_marlin_repack(
+                marlin_pad_qweight(
+                    x.data.contiguous(), size_n, size_k, padded_n, padded_k
+                ),
+                perm=layer.g_idx_sort_indices,
+                size_k=padded_k,
+                size_n=padded_n,
+                num_bits=c.weight_type.size_bits,
+                is_a_8bit=is_a_8bit,
+            )
+            return x
+
+        def transform_w_s(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1)
+            x.data = marlin_permute_scales(
+                marlin_pad_scales(
+                    x.data.contiguous(),
+                    size_n,
+                    size_k,
+                    padded_n,
+                    padded_k,
+                    c.group_size,
+                ),
+                size_k=padded_k,
+                size_n=padded_n,
+                group_size=c.group_size,
+                is_a_8bit=is_a_8bit,
+            )
+
+            if c.group_size == -1:
+                num_groups = 1
+            else:
+                num_groups = c.partition_weight_shape[0] // c.group_size
+
+            if c.act_type == torch.int8 and num_groups > 1:
+                x.data, input_global_scale = marlin_act_int8_process_scales(x.data)
+                layer.register_parameter(
+                    "input_global_scale",
+                    torch.nn.Parameter(input_global_scale, requires_grad=False),
+                )
+            else:
+                layer.input_global_scale = None
+            return x
+
+        if c.has_g_idx:
+            g_idx, g_idx_sort_indices = marlin_sort_g_idx(
+                getattr(layer, self.w_gidx_name)
+            )
+            self._transform_param(layer, self.w_gidx_name, lambda _: g_idx)
+            layer.g_idx_sort_indices = g_idx_sort_indices
+        else:
+            setattr(layer, self.w_gidx_name, marlin_make_empty_g_idx(device))
+            layer.g_idx_sort_indices = marlin_make_empty_g_idx(device)
+
+        if c.zero_points:
+            grouped_k = size_k // c.group_size if c.group_size != -1 else 1
+            padded_grouped_k = padded_k // c.group_size if c.group_size != -1 else 1
+            self._transform_param(
+                layer,
+                self.w_zp_name,
+                lambda x: marlin_zero_points(
+                    marlin_pad_scales(
+                        unpack_cols(
+                            x.t(),
+                            c.weight_type.size_bits,
+                            grouped_k,
+                            size_n,
+                        ),
+                        size_n,
+                        size_k,
+                        padded_n,
+                        padded_k,
+                        c.group_size,
+                    ),
+                    size_k=padded_grouped_k,
+                    size_n=padded_n,
+                    num_bits=c.weight_type.size_bits,
+                    is_a_8bit=is_a_8bit,
+                ),
+            )
+        else:
+            setattr(layer, self.w_zp_name, marlin_make_empty_g_idx(device))
+        self._transform_param(layer, self.w_q_name, transform_w_q)
+        self._transform_param(layer, self.w_s_name, transform_w_s)
+
+        if hasattr(layer, "bias") and layer.bias is not None:
+            layer.bias.data = marlin_permute_bias(
+                marlin_pad_dim(layer.bias, size_n, padded_n)
+            )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        c = self.config
+        w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
+
+        # `process_weights_after_loading` will ensure w_zp and w_gidx are not
+        #  None for marlin
+
+        return apply_gptq_marlin_linear(
+            input=x,
+            weight=w_q,
+            weight_scale=w_s,
+            weight_zp=w_zp,  # type: ignore
+            g_idx=w_gidx,  # type: ignore
+            g_idx_sort_indices=layer.g_idx_sort_indices,
+            workspace=self.workspace,
+            wtype=c.weight_type,
+            input_size_per_partition=c.partition_weight_shape[0],
+            output_size_per_partition=c.partition_weight_shape[1],
+            is_k_full=self.is_k_full,
+            input_global_scale=getattr(layer, "input_global_scale", None),
+            bias=bias,
+            input_dtype=c.act_type,
+        )

vllm/model_executor/kernels/linear/mixed_precision/rdna3_w4a16.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""W4A16 GPTQ kernel for AMD RDNA3 (gfx1100) — fp16 + bf16.
+
+Drop-in replacement for ExllamaLinearKernel on RDNA3 that adds native bf16
+support. The HIP kernel lives in ``csrc/rocm/q_gemm_rdna3.cu``
+and is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3``.
+
+Registered ahead of TritonW4A16LinearKernel for the ROCm-RDNA3 path; falls
+through to the Triton kernel on non-RDNA3 ROCm devices (e.g. CDNA/MI300).
+"""
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    pack_quantized_values_into_int32,
+)
+from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
+from vllm.platforms import current_platform
+from vllm.scalar_type import scalar_types
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+
+class RDNA3W4A16LinearKernel(MPLinearKernel):
+    SUPPORTED_QUANT_TYPES = [scalar_types.uint4b8]
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        # ROCm gates via on_gfx1100() in can_implement.
+        return 60
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        if not current_platform.is_rocm():
+            return False, "RDNA3 W4A16 kernel is ROCm-only"
+
+        from vllm.platforms.rocm import on_gfx1100
+
+        if not on_gfx1100():
+            return False, "RDNA3 W4A16 kernel requires gfx1100"
+
+        # The HIP op is registered by the C++ extension; if a user is running
+        # against a vLLM build that doesn't include it (e.g. partial rebuild),
+        # fall through gracefully to the next kernel in the registry.
+        if not (
+            hasattr(torch.ops, "_rocm_C")
+            and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
+        ):
+            return (
+                False,
+                "torch.ops._rocm_C.gptq_gemm_rdna3 missing — rebuild C++ extension",
+            )
+
+        if c.act_type not in (torch.float16, torch.bfloat16):
+            return False, "RDNA3 W4A16 kernel only supports fp16 and bf16"
+
+        if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
+            return (
+                False,
+                f"Quant type ({c.weight_type}) not supported by "
+                f"RDNA3 W4A16 kernel; supported: {cls.SUPPORTED_QUANT_TYPES}",
+            )
+
+        if c.group_size <= 0:
+            return (
+                False,
+                "RDNA3 W4A16 kernel does not support channelwise quantization",
+            )
+
+        if c.full_weight_shape[0] % c.group_size != 0:
+            return (
+                False,
+                f"Group size ({c.group_size}) does not evenly divide K "
+                f"({c.full_weight_shape[0]})",
+            )
+
+        # Output features must be a multiple of the pack factor (8 nibbles per
+        # int32) and of 8 so that qzeros (packed 4-bit per col) align cleanly
+        # against the BLOCK_KN_SIZE*4 = 512 N-stride and per-thread 4 columns.
+        if c.partition_weight_shape[1] % 8 != 0:
+            return (
+                False,
+                "Output features must be a multiple of 8 for the RDNA3 "
+                "W4A16 kernel (qzeros packing)",
+            )
+
+        if c.has_g_idx and c.partition_weight_shape[0] != c.full_weight_shape[0]:
+            return (
+                False,
+                "Act-order with TP-partitioned input features is not "
+                "supported by the RDNA3 W4A16 kernel",
+            )
+
+        return True, None
+
+    # ----- Weight prep (identical layout/shuffle as ExllamaLinearKernel) -----
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        c = self.config
+        device = getattr(layer, self.w_q_name).device
+
+        # Synthesize zero points if the checkpoint doesn't carry them.
+        if not c.zero_points:
+            self.w_zp_name = "qzeros"
+            groups = c.partition_weight_shape[0] // c.group_size
+            out_features = c.partition_weight_shape[1]
+
+            if c.weight_type.has_bias():
+                # GPTQv1 quirk: the kernel adds 1 to the stored zero, so we
+                # encode (bias - 1) here. See exllama.py for the link to the
+                # documentation of this checkpoint-format wart.
+                zeros = torch.full(
+                    (groups, out_features),
+                    c.weight_type.bias - 1,
+                    dtype=torch.int32,
+                    device=device,
+                )
+            else:
+                raise NotImplementedError(
+                    "RDNA3 W4A16 kernel: zero-bias 4-bit quant requires "
+                    "explicit zero points (GPTQv1 +1 quirk)."
+                )
+            zeros = pack_quantized_values_into_int32(zeros, c.weight_type, packed_dim=1)
+            setattr(
+                layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False)
+            )
+
+        # Act-order: convert g_idx to the inverse permutation array exllama
+        # expects (kernel reads a[perm[k]] instead of using groups indirected
+        # by g_idx[k]).
+        if c.has_g_idx:
+
+            def transform_w_g_idx(x):
+                return torch.argsort(x).to(torch.int)
+
+            self._transform_param(layer, self.w_gidx_name, transform_w_g_idx)  # type: ignore
+        else:
+            self.w_gidx_name = "g_idx"
+            empty_g_idx = torch.nn.Parameter(
+                torch.empty((0,), dtype=torch.int, device=device),
+                requires_grad=False,
+            )
+            setattr(layer, self.w_gidx_name, empty_g_idx)
+
+        def transform_w_q(x):
+            assert isinstance(x, BasevLLMParameter)
+            assert self.w_gidx_name is not None
+            g_idx = getattr(layer, self.w_gidx_name)
+
+            permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
+            x_cont = x.data.contiguous()
+            # Same 4-bit shuffle as exllama. The RDNA3 kernel reads weights in
+            # the same shuffled int32 layout and uses the (qa & 0x000F000F)
+            # bit-trick on top.
+            ops.gptq_shuffle(x_cont, g_idx, c.weight_type.size_bits)
+            return x_cont
+
+        def transform_w_s(x):
+            assert isinstance(x, BasevLLMParameter)
+            permute_param_layout_(x, input_dim=0, output_dim=1)
+            x.data = x.data.contiguous()
+            # Keep scales in the activation dtype (fp16 OR bf16) — the kernel
+            # branches on dtype internally.
+            return x.to(dtype=c.act_type)
+
+        self._transform_param(layer, self.w_q_name, transform_w_q)
+        self._transform_param(layer, self.w_s_name, transform_w_s)
+
+    # ----- Forward --------------------------------------------------------
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        c = self.config
+
+        x_2d = x.reshape(-1, x.shape[-1])
+        out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
+
+        w_q, w_s, w_zp, w_g_idx = self._get_weight_params(layer)
+
+        assert w_zp is not None, "Zero points are required by RDNA3 W4A16"
+        assert w_g_idx is not None, "g_idx tensor (possibly empty) required"
+
+        output = ops.gptq_gemm_rdna3(x_2d, w_q, w_zp, w_s, w_g_idx, False)
+
+        if bias is not None:
+            output.add_(bias)
+        return output.reshape(out_shape)

vllm/model_executor/kernels/linear/mixed_precision/triton_w4a16.py

@@ -0,0 +1,438 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Triton-based W4A16 GEMM kernel for ROCm MI300.
+
+Implements fused int4-weight dequantization + fp16 GEMM in a single kernel,
+using GPTQ sequential packing (8 int4 values per int32, shifts [0,4,...,28]).
+Plugs into the MPLinearKernel selection system and is preferred over
+MarlinLinearKernel/ExllamaLinearKernel on ROCm.
+
+Weight layout expected by this kernel (post-process_weights_after_loading):
+  qweight: [K, N//8]  int32  — rows=K (input), cols=N//8 (N is packed)
+  scales:  [K//G, N]  fp16/bf16
+  qzeros:  [K//G, N//8]  int32  (optional; None for symmetric uint4b8)
+
+Checkpoint layout from compressed_tensors_wNa16 create_weights:
+  weight_packed:     [N, K//8]  int32  (output_dim=0, input_dim=1, packed_dim=1)
+  weight_scale:      [N, K//G]  fp16   (output_dim=0, input_dim=1)
+  weight_zero_point: [N//8, K//G]  int32 (output_dim=0, packed_dim=0)
+"""
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
+from vllm.platforms import current_platform
+from vllm.scalar_type import scalar_types
+from vllm.triton_utils import tl, triton
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+TRITON_W4A16_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128, 256]
+TRITON_W4A16_SUPPORTED_QUANT_TYPES = [
+    scalar_types.uint4b8,  # symmetric GPTQ (bias=8)
+    scalar_types.uint4,  # asymmetric with explicit zeros
+]
+
+
+@triton.jit
+def triton_w4a16_gemm_kernel(
+    # Pointers
+    a_ptr,  # [M, K]  fp16/bf16 activations
+    b_ptr,  # [K, N//8]  int32 packed 4-bit weights (N is the packed dim)
+    scales_ptr,  # [K//G, N]  fp16/bf16 scales
+    zeros_ptr,  # [K//G, N//8]  int32 packed zeros (unused when HAS_ZP=False)
+    c_ptr,  # [M, N]  fp16/bf16 output
+    # Dimensions
+    M,
+    N,
+    K,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,  # stride in b along the packed N//8 dim
+    stride_cm,
+    stride_cn,
+    # Quantization parameters
+    group_size,
+    # Whether explicit zero points are provided
+    HAS_ZP: tl.constexpr,
+    # Zero bias used when HAS_ZP is False (e.g. 8 for uint4b8)
+    ZP_BIAS: tl.constexpr,
+    # Block sizes (tuned for MI300 wavefront=64)
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    """
+    Fused W4A16 GEMM: C[M,N] = A[M,K] @ dequant(B)[K,N]
+
+    B is stored as [K, N//8] int32 using GPTQ sequential packing:
+      each int32 packs 8 consecutive N-values at bit offsets [0,4,8,12,16,20,24,28].
+
+    Dequant: w_fp = (w_int4 - zero) * scale
+      HAS_ZP=True:  zero is loaded from zeros_ptr and unpacked
+      HAS_ZP=False: zero = ZP_BIAS constant (e.g. 8 for uint4b8 symmetric)
+    """
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    # Row/col offsets for this tile
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    # b/zeros are stored with N packed: N//8 int32 columns per K row
+    offs_bn = pid_n * (BLOCK_N // 8) + tl.arange(0, BLOCK_N // 8)
+
+    # GPTQ sequential shifts tiled across BLOCK_N:
+    #   [0,4,8,...,28] repeating for every group of 8 N-values.
+    # Build 1D shifts_1d of length BLOCK_N: column j gets shift (j % 8) * 4.
+    shifts_row = tl.arange(0, 8) * 4  # [8]
+    shifts_1d_2d = tl.broadcast_to(shifts_row[None, :], (BLOCK_N // 8, 8))
+    shifts_1d = tl.reshape(shifts_1d_2d, (BLOCK_N,))  # [BLOCK_N]
+    # Broadcast to [BLOCK_K, BLOCK_N] for weight unpacking
+    shifts = tl.broadcast_to(shifts_1d[None, :], (BLOCK_K, BLOCK_N))
+
+    # Scales column offsets: full N-width (one scale per output neuron)
+    offs_sn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    for k_start in range(0, tl.cdiv(K, BLOCK_K)):
+        offs_k = k_start * BLOCK_K + tl.arange(0, BLOCK_K)
+        mask_k = offs_k < K
+
+        # ---- Load activations A: [BLOCK_M, BLOCK_K] ----
+        a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
+        mask_a = (offs_m[:, None] < M) & mask_k[None, :]
+        a = tl.load(a_ptrs, mask=mask_a, other=0.0)
+
+        # ---- Load packed weights B: [BLOCK_K, BLOCK_N//8] int32 ----
+        b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn
+        mask_b = mask_k[:, None] & (offs_bn[None, :] < N // 8)
+        b_packed = tl.load(b_ptrs, mask=mask_b, other=0)
+
+        # ---- Unpack int4 weights → [BLOCK_K, BLOCK_N] ----
+        # tl.interleave(x, x) doubles the last dim by interleaving.
+        # Starting from [BLOCK_K, BLOCK_N//8], three interleaves give
+        # [BLOCK_K, BLOCK_N], where each int32 is replicated 8 times.
+        b = tl.interleave(b_packed, b_packed)
+        b = tl.interleave(b, b)
+        b = tl.interleave(b, b)
+        # Extract the correct 4-bit nibble for each output column
+        b = (b >> shifts) & 0xF
+
+        # ---- Compute scale/zero group row index ----
+        g_idx = (k_start * BLOCK_K) // group_size
+
+        # ---- Load scales: [BLOCK_N] → broadcast to [BLOCK_K, BLOCK_N] ----
+        scale_offset = g_idx * N + offs_sn
+        scale_mask = offs_sn < N
+        scales = tl.load(scales_ptr + scale_offset, mask=scale_mask, other=1.0)
+        scales = tl.broadcast_to(scales[None, :], (BLOCK_K, BLOCK_N))
+
+        # ---- Load / compute zeros ----
+        if HAS_ZP:
+            # Load packed zeros row: [BLOCK_N//8] int32
+            zero_offset = g_idx * (N // 8) + offs_bn
+            zero_mask = offs_bn < N // 8
+            z_packed = tl.load(zeros_ptr + zero_offset, mask=zero_mask, other=0)
+            # Unpack to [BLOCK_N] using same interleave+shift pattern
+            z = tl.interleave(z_packed, z_packed)
+            z = tl.interleave(z, z)
+            z = tl.interleave(z, z)
+            z = (z >> shifts_1d) & 0xF
+            z = tl.broadcast_to(z[None, :], (BLOCK_K, BLOCK_N))
+        else:
+            z = tl.full((BLOCK_K, BLOCK_N), ZP_BIAS, dtype=tl.int32)
+
+        # ---- Dequantize: (w - zero) * scale ----
+        b_fp = (b - z).to(a.dtype) * scales
+
+        # ---- Accumulate ----
+        accumulator += tl.dot(a, b_fp, out_dtype=tl.float32)
+
+    # ---- Store output C: [BLOCK_M, BLOCK_N] ----
+    c = accumulator.to(c_ptr.type.element_ty)
+    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
+    mask_c = (offs_m[:, None] < M) & (offs_n[None, :] < N)
+    tl.store(c_ptrs, c, mask=mask_c)
+
+
+def triton_w4a16_gemm(
+    a: torch.Tensor,  # [M, K] fp16/bf16
+    b_q: torch.Tensor,  # [K, N//8] int32
+    scales: torch.Tensor,  # [K//G, N] fp16/bf16
+    qzeros: torch.Tensor | None,  # [K//G, N//8] int32, or None
+    group_size: int,
+    zp_bias: int = 8,  # bias for uint4b8 when qzeros is None
+) -> torch.Tensor:
+    """
+    Fused W4A16 GEMM using GPTQ-packed int4 weights.
+
+    Args:
+        a:          Activation matrix [M, K], float16 or bfloat16.
+        b_q:        Packed weight matrix [K, N//8], int32 (GPTQ sequential).
+        scales:     Per-group scales [K//G, N], same dtype as a.
+        qzeros:     Per-group packed zero points [K//G, N//8] int32, or None
+                    for symmetric quantization (uses zp_bias instead).
+        group_size: Quantization group size (resolved from -1 to K by caller).
+        zp_bias:    Constant zero used when qzeros is None (default 8 for uint4b8).
+
+    Returns:
+        Output matrix [M, N], same dtype as a.
+    """
+    assert a.is_contiguous(), "Activation matrix must be contiguous"
+    assert b_q.is_contiguous(), "Weight matrix must be contiguous"
+    assert scales.is_contiguous(), "Scales must be contiguous"
+
+    M, K = a.shape
+    N = b_q.shape[1] * 8
+
+    assert b_q.shape == (K, N // 8), (
+        f"b_q shape mismatch: {b_q.shape} vs ({K}, {N // 8})"
+    )
+    assert scales.shape == (K // group_size, N), (
+        f"scales shape mismatch: {scales.shape} vs ({K // group_size}, {N})"
+    )
+    if qzeros is not None:
+        assert qzeros.shape == (K // group_size, N // 8), (
+            f"qzeros shape mismatch: {qzeros.shape}"
+        )
+
+    c = torch.empty((M, N), dtype=a.dtype, device=a.device)
+
+    has_zp = qzeros is not None
+    # Provide a dummy pointer when HAS_ZP=False (Triton requires a valid ptr)
+    zeros_ptr = qzeros if has_zp else b_q
+
+    if current_platform.is_rocm():
+        from vllm.platforms.rocm import on_gfx1x
+
+        if on_gfx1x():
+            # Tuned for RDNA 3.5 (gfx1151, 40 CUs, 32-wide wavefronts).
+            if M <= 32:
+                BLOCK_M, BLOCK_N, BLOCK_K = 32, 32, 64
+            elif M <= 64:
+                BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
+            else:
+                BLOCK_M, BLOCK_N, BLOCK_K = 128, 32, 64
+        else:
+            # Tuned for MI300 (gfx942, 304 CUs, 64-wide wavefronts).
+            if M <= 32:
+                BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32
+            elif M <= 64:
+                BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
+            else:
+                BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32
+    else:
+        if M <= 32:
+            BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32
+        elif M <= 64:
+            BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32
+        else:
+            BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32
+
+    # The kernel loads scales/zeros for a single group per BLOCK_K tile
+    # (one g_idx per iteration). If BLOCK_K > group_size, rows at the tail
+    # of the tile dequantize with the wrong group's scales, silently
+    # corrupting the output. Clamp BLOCK_K to group_size to keep one
+    # scale group per tile.
+    if group_size < BLOCK_K:
+        BLOCK_K = group_size
+
+    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
+
+    triton_w4a16_gemm_kernel[grid](
+        a,
+        b_q,
+        scales,
+        zeros_ptr,
+        c,
+        M,
+        N,
+        K,
+        a.stride(0),
+        a.stride(1),
+        b_q.stride(0),
+        b_q.stride(1),
+        c.stride(0),
+        c.stride(1),
+        group_size=group_size,
+        HAS_ZP=has_zp,
+        ZP_BIAS=zp_bias,
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+    )
+    return c
+
+
+class TritonW4A16LinearKernel(MPLinearKernel):
+    """
+    Triton-based W4A16 GEMM kernel for ROCm (MI300 and newer).
+
+    Supports GPTQ-format int4 weights (uint4b8 symmetric, uint4 asymmetric)
+    with grouped quantization. Weight tensors are transposed from the
+    compressed-tensors checkpoint layout to the kernel's [K, N//8] layout.
+    """
+
+    SUPPORTED_QUANT_TYPES = TRITON_W4A16_SUPPORTED_QUANT_TYPES
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        # Triton handles capability checks itself
+        return 0
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        if not (current_platform.is_rocm() or current_platform.is_cuda()):
+            return False, "TritonW4A16LinearKernel requires CUDA or ROCm"
+
+        if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
+            return (
+                False,
+                f"Quant type {c.weight_type} not supported; "
+                f"supported: {cls.SUPPORTED_QUANT_TYPES}",
+            )
+
+        if c.act_type not in (torch.float16, torch.bfloat16):
+            return False, "Only float16/bfloat16 activations are supported"
+
+        N = c.partition_weight_shape[1]
+        if N % 8 != 0:
+            return (
+                False,
+                f"Output features ({N}) must be divisible by 8 "
+                "(8 int4 values packed per int32)",
+            )
+
+        if c.has_g_idx:
+            return (
+                False,
+                "Activation reordering (g_idx) is not supported by "
+                "TritonW4A16LinearKernel",
+            )
+
+        gs = c.group_size
+        if (
+            gs not in TRITON_W4A16_SUPPORTED_GROUP_SIZES
+            and gs != c.full_weight_shape[0]
+        ):
+            return (
+                False,
+                f"Group size {gs} not supported; "
+                f"supported: {TRITON_W4A16_SUPPORTED_GROUP_SIZES} "
+                f"or full K ({c.full_weight_shape[0]})",
+            )
+
+        K = c.partition_weight_shape[0]
+        eff_gs = gs if gs != -1 else K
+        if K % eff_gs != 0:
+            return (False, f"Input features {K} not divisible by group size {eff_gs}")
+
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """
+        Convert compressed-tensors checkpoint layout to kernel layout.
+
+        Checkpoint (from compressed_tensors_wNa16.create_weights):
+          weight_packed:     [N, K//8]  int32   input_dim=1, output_dim=0, packed_dim=1
+          weight_scale:      [N, K//G]  fp16    input_dim=1, output_dim=0
+          weight_zero_point: [N//8, K//G] int32  output_dim=0, packed_dim=0
+
+        Kernel needs:
+          qweight: [K, N//8]  int32   (transpose weight_packed)
+          scales:  [K//G, N]  fp16    (transpose weight_scale)
+          qzeros:  [K//G, N//8] int32 (transpose weight_zero_point)
+        """
+
+        # ---- Transform qweight: [N, K//8] → [K//8, N] → back to [K, N//8] ----
+        # permute_param_layout_(x, input_dim=0, output_dim=1) rearranges so that
+        # the input(K) dimension is at physical dim 0 and output(N) at dim 1.
+        # Checkpoint has input_dim=1, output_dim=0, packed_dim=1 (K is packed).
+        # After permute we get [K//8, N] (K packed at dim 0, N at dim 1).
+        # The kernel wants [K, N//8] (K at dim 0, N packed at dim 1), so we
+        # then transpose: [K//8, N].T = [N, K//8] — that's not right.
+        #
+        # Actually we need to change WHAT is packed:
+        #   Original packing: K packed into K//8 (8 K-values per int32)
+        #   Kernel packing:   N packed into N//8 (8 N-values per int32)
+        # These require a full repack, not just a transpose.
+        #
+        # Simple approach: unpack → transpose the full [N, K] → repack as [K, N//8].
+        # This is done CPU-side at load time (one-time cost).
+        def repack_w_q(x: BasevLLMParameter) -> BasevLLMParameter:
+            # x.data is [N, K//8] int32, K packed (GPTQ checkpoint format)
+            # Step 1: bring to [N, K//8] with output(N) at dim 0
+            permute_param_layout_(x, input_dim=1, output_dim=0, packed_dim=1)
+            w = x.data  # [N, K//8] int32
+
+            N_dim, K8 = w.shape
+            K_dim = K8 * 8
+            # Step 2: unpack to [N, K] int32 (vectorized)
+            shifts = torch.arange(8, device=w.device, dtype=torch.int32) * 4
+            w_unpacked = ((w.unsqueeze(-1) >> shifts) & 0xF).reshape(N_dim, K_dim)
+            # Step 3: transpose to [K, N] int32
+            w_KN = w_unpacked.t().contiguous()
+            # Step 4: repack N into N//8 int32 values → [K, N//8] (vectorized)
+            N8 = N_dim // 8
+            w_repacked = torch.sum(
+                (w_KN.view(K_dim, N8, 8) & 0xF) << shifts,
+                dim=2,
+                dtype=torch.int32,
+            )
+            x.data = w_repacked.contiguous()
+            return x
+
+        def repack_w_s(x: BasevLLMParameter) -> BasevLLMParameter:
+            # x.data is [N, K//G] fp16, bring to [K//G, N]
+            permute_param_layout_(x, input_dim=1, output_dim=0)
+            x.data = x.data.t().contiguous()
+            return x
+
+        self._transform_param(layer, self.w_q_name, repack_w_q)
+        self._transform_param(layer, self.w_s_name, repack_w_s)
+
+        if self.w_zp_name is not None:
+            zp = getattr(layer, self.w_zp_name, None)
+            if zp is not None:
+                # Checkpoint: [N//8, K//G] int32 (N packed at dim 0, K//G at dim 1)
+                # Kernel needs: [K//G, N//8] — just transpose
+                replace_parameter(
+                    layer,
+                    self.w_zp_name,
+                    torch.nn.Parameter(zp.data.t().contiguous(), requires_grad=False),
+                )
+
+    def apply_weights(
+        self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None
+    ) -> torch.Tensor:
+        c = self.config
+        w_q, w_s, w_zp, _ = self._get_weight_params(layer)
+
+        x_2d = x.reshape(-1, x.shape[-1]).contiguous()
+        out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
+
+        K = c.partition_weight_shape[0]
+        group_size = c.group_size if c.group_size != -1 else K
+
+        # For symmetric types (uint4b8), use the scalar bias; no zeros tensor
+        zp_bias = c.weight_type.bias if c.weight_type.has_bias() else 0
+
+        output = triton_w4a16_gemm(
+            a=x_2d,
+            b_q=w_q,
+            scales=w_s,
+            qzeros=w_zp,
+            group_size=group_size,
+            zp_bias=zp_bias,
+        )
+
+        if bias is not None:
+            output.add_(bias)
+
+        return output.reshape(out_shape)

vllm/model_executor/kernels/linear/mixed_precision/xpu.py

@@ -0,0 +1,231 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+from torch.nn.parameter import Parameter
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.platforms import current_platform
+from vllm.scalar_type import scalar_types
+
+from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
+
+_XPUWNA16_SUPPORTED_QUANT_TYPES = (scalar_types.uint4, scalar_types.uint4b8)
+
+logger = init_logger(__name__)
+
+
+class XPUwNa16LinearKernel(MPLinearKernel):
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return -1
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        if not current_platform.is_xpu():
+            return False, "XPUwNa16 only supported on XPU"
+
+        if c.act_type != torch.bfloat16 and c.act_type != torch.float16:
+            return False, "XPUwNa16 only supports BF16/FP16 activations"
+
+        if c.weight_type not in _XPUWNA16_SUPPORTED_QUANT_TYPES:
+            return (
+                False,
+                f"Quant type ({c.weight_type}) not supported by "
+                "XPUwNa16, supported types are: "
+                f"{_XPUWNA16_SUPPORTED_QUANT_TYPES}",
+            )
+        if c.group_size != -1 and c.group_size % 32 != 0:
+            return (
+                False,
+                f"Group size ({c.group_size}) not supported by "
+                "XPUwNa16, supported group sizes are multiples of 32",
+            )
+
+        if c.partition_weight_shape[0] % 32 != 0:
+            return (
+                False,
+                f"Input size ({c.partition_weight_shape[0]}) not supported by "
+                "XPUwNa16, supported sizes are multiples of 32",
+            )
+
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        # Default names since marlin requires empty parameters for these,
+        # TODO: remove this requirement from marlin (allow optional tensors)
+        if self.w_gidx_name is None:
+            self.w_gidx_name = "g_idx"
+        if self.w_zp_name is None:
+            self.w_zp_name = "w_zp"
+
+        need_transpose = False
+        qweight_shape = getattr(layer, self.w_q_name).shape
+        scale_shape = getattr(layer, self.w_s_name).shape
+        # gptq marlin and compressed tensors wna16 expect different default
+        # layouts for weight and scale, so we check the shapes to determine
+        # if we need to transpose
+        if qweight_shape[0] != scale_shape[0]:
+            need_transpose = True
+
+        if need_transpose:
+            getattr(layer, self.w_q_name).data = (
+                getattr(layer, self.w_q_name).data.t().contiguous()
+            )
+            getattr(layer, self.w_s_name).data = getattr(layer, self.w_s_name).data
+        else:
+            getattr(layer, self.w_s_name).data = (
+                getattr(layer, self.w_s_name).data.t().contiguous()
+            )
+
+        if self.config.zero_points:
+            # (FIXME): maybe zero points should also be transposed.
+            getattr(layer, self.w_zp_name).data = (
+                getattr(layer, self.w_zp_name).data.t().contiguous()
+            )
+        else:
+            weight_zero_point = torch.Tensor([8]).to(torch.int8).to("xpu")
+            setattr(
+                layer, self.w_zp_name, Parameter(weight_zero_point, requires_grad=False)
+            )
+        if self.config.has_g_idx:
+            setattr(
+                layer,
+                self.w_gidx_name,
+                Parameter(
+                    getattr(layer, self.w_gidx_name).data.t().contiguous(),
+                    requires_grad=False,
+                ),
+            )
+        else:
+            setattr(layer, self.w_gidx_name, None)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        reshaped_x = x.reshape(-1, x.shape[-1])
+        w_q, w_s, w_zp, w_gidx = self._get_weight_params(layer)
+        out = torch.ops._xpu_C.int4_gemm_w4a16(
+            reshaped_x,
+            w_q.t(),
+            bias if bias is not None else None,
+            w_s,
+            w_zp,
+            self.config.group_size,
+            w_gidx,
+        )
+        return out
+
+
+class XPUW4A8IntLinearKernel(MPLinearKernel):
+    """XPU kernel for W4A8 integer quantization using oneDNN int4_gemm_w4a8.
+
+    Weights are symmetric group-quantized int4 packed as uint4.
+    Activations are dynamically quantized per-token to symmetric int8.
+    """
+
+    @classmethod
+    def get_min_capability(cls) -> int:
+        return -1
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        if not current_platform.is_xpu():
+            return False, "XPUW4A8Int only supported on XPU"
+        if c.act_type not in (torch.bfloat16, torch.float16):
+            return False, "XPUW4A8Int requires BF16/FP16 activations"
+        if c.weight_type != scalar_types.int4:
+            return (
+                False,
+                f"XPUW4A8Int requires int4 weights, got {c.weight_type}",
+            )
+        if c.zero_points:
+            return False, "XPUW4A8Int only supports symmetric weight quantization"
+        if c.group_size != -1 and c.group_size % 32 != 0:
+            return (
+                False,
+                f"Group size ({c.group_size}) not supported by XPUW4A8Int, "
+                "must be a multiple of 32",
+            )
+        in_size, out_size = c.partition_weight_shape
+        if in_size % 8 != 0 or out_size % 8 != 0:
+            return (
+                False,
+                f"in/out sizes ({in_size}, {out_size}) must be multiples of 8",
+            )
+
+        if c.act_type != torch.float16:
+            logger.warning_once(
+                "XPUW4A8IntLinearKernel is running with model dtype %s, "
+                "but int4_gemm_w4a8 produces float16 output. Recommend "
+                "setting --dtype float16 for best performance.",
+                c.act_type,
+            )
+
+        return True, None
+
+    def _pack_int4_weight(self, w: torch.Tensor) -> torch.Tensor:
+        # w is [N, K] int8 with values in [-8, 7]
+        w_u4 = w.to(torch.int32) + 8  # shift to [0, 15]
+        w_u4 = w_u4.reshape(w.shape[0], w.shape[1] // 8, 8)  # [N, K/8, 8]
+        shifts = torch.arange(0, 32, 4, dtype=torch.int32, device=w.device)
+        packed = ((w_u4 & 0xF) << shifts[None, None, :]).sum(dim=2).to(torch.int32)
+        return packed
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.weight_scale.data = layer.weight_scale.data.t().contiguous()
+
+        device = layer.weight_packed.device
+        # TODO: support asymmetric quantization
+        weight_zero_point = torch.tensor([8], dtype=torch.int8, device=device)
+        layer.weight_zero_point = Parameter(weight_zero_point, requires_grad=False)
+
+        # weight_packed is [out, in] int8, signed int4 values in [-8, 7]
+        w = layer.weight_packed.data  # [out, in]
+
+        # TODO: implement asym case
+        packed = self._pack_int4_weight(w)  # [out, in/8] packed uint4
+
+        replace_parameter(
+            layer,
+            self.w_q_name,
+            torch.nn.Parameter(packed, requires_grad=False),
+        )
+
+        # Free the original unpacked int8 weight (still registered as "weight")
+        # to avoid double-storing both int8 [N, K] and int32 [N, K/8] in memory.
+        layer.register_parameter("weight", None)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        reshaped_x = x.reshape(-1, x.shape[-1])  # [M, K]
+        from vllm._xpu_ops import xpu_ops as ops
+
+        # TODO: static and asymmetric quantization case
+        # Common code for CompressedTensorsW4A8Int does not read act symmetry data
+        quant_x, x_scale, x_zero = ops.dynamic_per_token_int8_quant_ref(
+            reshaped_x, True, 8
+        )
+
+        out = torch.ops._xpu_C.int4_gemm_w4a8(
+            quant_x,
+            x_scale,
+            x_zero,
+            layer.weight_packed.t(),
+            layer.weight_scale,
+            layer.weight_zero_point,
+            self.config.group_size,
+            None,  # g_idx not currently supported
+            bias,
+        )
+
+        return out.to(x.dtype)

vllm/model_executor/kernels/linear/mixed_precision/zentorch.py

@@ -0,0 +1,211 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Zentorch W4A16 GPTQ weight-only-quantized linear kernel for AMD Zen CPUs.
+
+Selected by ``choose_mp_linear_kernel`` ahead of the generic oneDNN-backed
+``CPUWNA16LinearKernel``. When ``can_implement`` rejects a layer, the selector
+falls through to the next kernel in ``_POSSIBLE_KERNELS[PlatformEnum.CPU]``.
+"""
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
+from vllm.platforms import current_platform
+from vllm.scalar_type import scalar_types
+
+from .cpu import CPUWNA16LinearKernel
+from .MPLinearKernel import MPLinearLayerConfig
+
+logger = init_logger(__name__)
+
+
+def _import_unpack_from_int32():
+    """Import compressed-tensors' ``unpack_from_int32`` across versions."""
+    try:
+        from compressed_tensors.compressors.pack_quantized.helpers import (
+            unpack_from_int32,
+        )
+    except ImportError:
+        from compressed_tensors.compressors.quantized_compressors.pack_quantized import (  # type: ignore[import-not-found]  # noqa: E501
+            unpack_from_int32,
+        )
+    return unpack_from_int32
+
+
+class ZentorchWNA16LinearKernel(CPUWNA16LinearKernel):
+    """W4A16 GPTQ kernel backed by ``torch.ops.zentorch.zentorch_woq_linear``."""
+
+    @classmethod
+    def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
+        ok, reason = super().can_implement(c)
+        if not ok:
+            return ok, reason
+
+        if not current_platform.is_zen_cpu():
+            return False, "ZentorchWNA16 requires an AMD Zen CPU."
+
+        if not has_zentorch_op(["zentorch_woq_repack_weight", "zentorch_woq_linear"]):
+            return (
+                False,
+                "torch.ops.zentorch.{zentorch_woq_repack_weight, "
+                "zentorch_woq_linear} are not registered.",
+            )
+
+        if c.has_g_idx:
+            return False, "ZentorchWNA16 does not support activation re-ordering."
+        return True, None
+
+    def _zentorch_woq_eligible(self, layer: torch.nn.Module) -> bool:
+        """Eligibility predicate for the zentorch W4A16 GPTQ fast path.
+
+        Constraints (any failure -> ``cpu_gemm_wna16`` path via ``super()``
+        with ``layer`` untouched).
+        """
+        if (
+            self.w_gidx_name is not None
+            and getattr(layer, self.w_gidx_name, None) is not None
+        ) or (getattr(self.config, "has_g_idx", False)):
+            return False
+
+        weight_packed = getattr(layer, self.w_q_name, None)
+        weight_scale = getattr(layer, self.w_s_name, None)
+        if weight_packed is None or weight_scale is None:
+            return False
+
+        bits = self.config.weight_type.mantissa
+        pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
+        # 4-bit -> 8 values per int32;
+        if pack_factor != 8:
+            return False
+
+        # GPTQ-only. AWQ packs along the output dim instead.
+        in_dim = getattr(weight_packed, "input_dim", None)
+        pk_dim = getattr(weight_packed, "packed_dim", None)
+        if in_dim is None or pk_dim is None or in_dim != pk_dim:
+            return False
+
+        is_ct_format = in_dim == pk_dim == 1
+        if not is_ct_format:
+            return False
+
+        if weight_packed.dim() != 2 or weight_scale.dim() != 2:
+            return False
+
+        # 4-bit -> 8 values per int32; in_features must be divisible by num_groups.
+        in_features = weight_packed.shape[1] * 8
+        num_groups = weight_scale.shape[1]
+        return num_groups > 0 and in_features % num_groups == 0
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Repack CT GPTQ weights into the zentorch WOQ layout.
+
+        Falls back to ``CPUWNA16LinearKernel.process_weights_after_loading``
+        via ``super()`` when the layer doesn't satisfy
+        ``_zentorch_woq_eligible``.
+
+        On success, ``layer._zentorch_processed_weights`` is set to ``True``
+        """
+        if getattr(layer, "_zentorch_processed_weights", False):
+            return
+
+        if not self._zentorch_woq_eligible(layer):
+            logger.info_once(
+                "[zen_cpu] ZentorchWNA16 fast path not eligible for this "
+                "layer (AWQ pack layout, g_idx, or non-int32 storage); "
+                "falling back to CPUWNA16LinearKernel (cpu_gemm_wna16)."
+            )
+            super().process_weights_after_loading(layer)
+            return
+
+        if (not self.config.zero_points) and (self.w_zp_name is not None):
+            setattr(layer, self.w_zp_name, None)
+
+        if (not self.config.has_g_idx) and (self.w_gidx_name is not None):
+            setattr(layer, self.w_gidx_name, None)
+
+        weight_q = getattr(layer, self.w_q_name)
+        weight_s = getattr(layer, self.w_s_name)
+        weight_packed = weight_q.data if hasattr(weight_q, "data") else weight_q
+        weight_scale = weight_s.data if hasattr(weight_s, "data") else weight_s
+
+        bits = self.config.weight_type.mantissa
+        pack_factor = torch.iinfo(weight_packed.dtype).bits // bits
+        out_features, num_groups = weight_scale.shape[0], weight_scale.shape[1]
+        in_features = weight_packed.shape[1] * pack_factor
+        original_shape = torch.Size([out_features, in_features])
+        unpack_from_int32 = _import_unpack_from_int32()
+        repack_op = torch.ops.zentorch.zentorch_woq_repack_weight.default
+
+        weight_unpacked = unpack_from_int32(
+            weight_packed,
+            bits,
+            original_shape,
+            packed_dim=weight_q.packed_dim,
+        )
+
+        zp_param = (
+            getattr(layer, self.w_zp_name, None) if self.w_zp_name is not None else None
+        )
+        needs_unsigned_offset = self.config.weight_type == scalar_types.uint4
+
+        if needs_unsigned_offset:
+            weight_unpacked = (weight_unpacked.to(torch.int32) + 8).clamp(0, 15)
+        repacked = repack_op(weight_unpacked.to(torch.int8).contiguous())
+
+        if zp_param is None:
+            zp_tc = None
+        else:
+            zp_tensor = zp_param.data if hasattr(zp_param, "data") else zp_param
+            zp = unpack_from_int32(
+                zp_tensor,
+                bits,
+                (out_features, num_groups),
+                packed_dim=zp_param.packed_dim,
+            )
+            if needs_unsigned_offset:
+                zp = (zp.to(torch.int32) + 8).clamp(0, 15)
+            zp_tc = zp.to(torch.int8).t().contiguous()
+
+        layer._zentorch_woq_packed = repacked.t()
+        layer._zentorch_woq_scale = weight_scale.t().contiguous()
+        layer._zentorch_woq_zero_point = zp_tc
+
+        for param_name in (self.w_q_name, self.w_s_name, self.w_zp_name):
+            if param_name is None:
+                continue
+            param = getattr(layer, param_name, None)
+            if param is None:
+                continue
+            if hasattr(param, "data"):
+                param.data = torch.empty(0)
+            else:
+                setattr(layer, param_name, torch.empty(0))
+
+        layer._zentorch_kind = "compressed_tensors_w4a16_gptq"
+        layer._zentorch_processed_weights = True
+        logger.info_once(
+            "[zen_cpu] Using zentorch_woq_linear for W4A16 GPTQ "
+            "(weight_type=%s, has_zp=%s)",
+            self.config.weight_type,
+            zp_tc is not None,
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if getattr(layer, "_zentorch_processed_weights", False):
+            return torch.ops.zentorch.zentorch_woq_linear.default(
+                x,
+                layer._zentorch_woq_packed,
+                layer._zentorch_woq_scale,
+                layer._zentorch_woq_zero_point,
+                bias,
+            )
+        return super().apply_weights(layer, x, bias)
+
+
+__all__ = ["ZentorchWNA16LinearKernel"]

vllm/model_executor/kernels/linear/mxfp4/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.model_executor.kernels.linear.mxfp4.base import (
+    MxFp4LinearKernel,
+    MxFp4LinearLayerConfig,
+)
+
+__all__ = [
+    "MxFp4LinearKernel",
+    "MxFp4LinearLayerConfig",
+]

vllm/model_executor/kernels/linear/mxfp4/base.py

@@ -0,0 +1,67 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class MxFp4LinearLayerConfig:
+    """Configuration for an MXFP4 linear layer.
+
+    All MXFP4 layers share the same structure: packed uint8 weights (2 FP4 values per
+    byte) and per-block weight scales (group size 32).
+    """
+
+    pass
+
+
+class MxFp4LinearKernel(ABC):
+    """Base class for MXFP4 quantized linear kernels.
+
+    Each subclass implements a specific GEMM backend (CUTLASS, Marlin, etc).
+    The kernel selection mechanism iterates over registered subclasses in
+    priority order,calling ``is_supported`` and ``can_implement`` to find the best
+    match for the current hardware.
+    """
+
+    def __init__(self, config: MxFp4LinearLayerConfig) -> None:
+        assert self.can_implement(config)[0]
+        assert self.is_supported()[0]
+        self.config = config
+
+    @classmethod
+    @abstractmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        """Return whether this kernel can run on the current platform."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        """Return whether this kernel can handle *config*."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Transform weights into the format required by this kernel.
+
+        Called once after checkpoint weights have been loaded onto the
+        device.  Implementations should repack / swizzle / pad weights
+        and scales in-place on *layer*.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Run the quantized GEMM."""
+        raise NotImplementedError

vllm/model_executor/kernels/linear/mxfp4/flashinfer.py

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from torch.nn.parameter import Parameter
+
+from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
+    swizzle_mxfp4_scales,
+)
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import has_flashinfer_cutedsl
+
+from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
+
+_MXFP4_GROUP_SIZE = 32
+
+
+class FlashInferMxFp4LinearKernel(MxFp4LinearKernel):
+    """MXFP4 W4A4 GEMM via FlashInfer CUTLASS (SM100+)."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if current_platform.has_device_capability(100) and has_flashinfer_cutedsl():
+            return True, None
+        return False, "FlashInfer + >=sm_100 (Blackwell) required"
+
+    @classmethod
+    def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        N, scale_K = layer.weight_scale.shape
+        K = scale_K * _MXFP4_GROUP_SIZE
+
+        # swizzle pads N to the next multiple of 128 for CUTLASS tiling
+        padded_N = ((N + 127) // 128) * 128
+        layer.weight_scale = Parameter(
+            swizzle_mxfp4_scales(layer.weight_scale.data, N, K).reshape(padded_N, -1),
+            requires_grad=False,
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        from vllm.utils.flashinfer import (
+            flashinfer_mxfp4_quantize,
+            flashinfer_scaled_fp4_mm,
+        )
+
+        weight = layer.weight
+        out_shape = x.shape[:-1] + (layer.output_size_per_partition,)
+        x_2d = x.reshape(-1, x.shape[-1])
+
+        x_fp4, x_scale = flashinfer_mxfp4_quantize(x_2d.contiguous())
+        out = flashinfer_scaled_fp4_mm(
+            x_fp4,
+            weight,
+            x_scale,
+            layer.weight_scale,
+            alpha=None,
+            out_dtype=x.dtype,
+            backend="cute-dsl",
+            block_size=_MXFP4_GROUP_SIZE,
+            use_nvfp4=False,
+        )
+
+        if bias is not None:
+            out = out + bias
+        return out.view(out_shape)

vllm/model_executor/kernels/linear/mxfp4/humming.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.humming_utils import (
+    convert_linear_layer_to_humming_standard,
+    prepare_humming_layer,
+)
+from vllm.platforms import current_platform
+
+from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
+
+
+class HummingMxFp4LinearKernel(MxFp4LinearKernel):
+    """Humming GEMM Kernel for MXFP4."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "Humming only supported on CUDA"
+
+        if not current_platform.has_device_capability(75):
+            return False, "Humming only supported on SM75+"
+
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu)
+        name_map = {"weight": "weight", "weight_scale": "weight_scale"}
+
+        quant_config = {
+            "quant_method": "humming",
+            "dtype": "float4e2m1",
+            "scale_dtype": "float8e8m0",
+            "group_size": 32,
+            "weight_scale_type": "group",
+        }
+
+        convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
+        prepare_humming_layer(layer, quant_config)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        from vllm.utils.humming import HummingMethod
+
+        flatten_inputs = x.view(-1, x.size(-1))
+        output = HummingMethod.forward_layer(
+            layer=layer,
+            inputs=flatten_inputs,
+            compute_config=layer.compute_config,
+        )
+        return output.view(*x.shape[:-1], output.size(-1))

vllm/model_executor/kernels/linear/mxfp4/marlin.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
+
+
+class MarlinMxFp4LinearKernel(MxFp4LinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
+            is_fp4_marlin_supported,
+        )
+
+        if is_fp4_marlin_supported():
+            return True, None
+        return False, "Marlin FP4 not available"
+
+    @classmethod
+    def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
+            prepare_fp4_layer_for_marlin,
+        )
+
+        prepare_fp4_layer_for_marlin(layer)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
+            apply_fp4_marlin_linear,
+        )
+
+        return apply_fp4_marlin_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale,
+            weight_global_scale=None,
+            workspace=layer.workspace,
+            size_n=layer.output_size_per_partition,
+            size_k=layer.input_size_per_partition,
+            bias=bias,
+        )

vllm/model_executor/kernels/linear/mxfp4/xpu.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.mxfp4_utils import (
+    xpu_mxfp4_quantize as quant_mxfp4,
+)
+from vllm.model_executor.utils import replace_parameter
+from vllm.platforms import current_platform
+
+from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
+
+
+class XPUMxFp4LinearKernel(MxFp4LinearKernel):
+    """MXFP4 W4A4 GEMM on XPU."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_xpu():
+            return False, "XPUMxFp4 only support on XPU"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight.view(torch.float4_e2m1fn_x2)
+        replace_parameter(layer, "weight", weight.data.t())
+
+        weight_scale = layer.weight_scale.view(torch.float8_e8m0fnu)
+        weight_scale = weight_scale.t().contiguous()
+        replace_parameter(layer, "weight_scale", weight_scale.data)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        out_dtype = x.dtype
+        x_fp4, x_blockscale = quant_mxfp4(x)
+        return torch.ops._xpu_C.fp4_gemm(
+            x_fp4,
+            layer.weight,
+            x_blockscale,
+            layer.weight_scale,
+            out_dtype,
+            bias,
+        )

vllm/model_executor/kernels/linear/mxfp8/Mxfp8LinearKernel.py

@@ -0,0 +1,56 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class Mxfp8LinearLayerConfig:
+    """Configuration for an MXFP8 linear layer.
+
+    All MXFP8 layers share the same structure: FP8-E4M3 weights with
+    uint8 (E8M0) per-block scales at block size 32.
+    """
+
+    pass
+
+
+class Mxfp8LinearKernel(ABC):
+    """Base class for MXFP8 quantized linear kernels.
+
+    Each subclass implements a specific GEMM backend (FlashInfer CUTLASS,
+    Marlin, emulation).
+    """
+
+    def __init__(self, c: Mxfp8LinearLayerConfig) -> None:
+        assert self.can_implement(c)[0]
+        assert self.is_supported()[0]
+        self.config = c
+
+    @classmethod
+    @abstractmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        raise NotImplementedError

vllm/model_executor/kernels/linear/mxfp8/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
+    Mxfp8LinearKernel,
+    Mxfp8LinearLayerConfig,
+)
+
+__all__ = [
+    "Mxfp8LinearKernel",
+    "Mxfp8LinearLayerConfig",
+]

vllm/model_executor/kernels/linear/mxfp8/emulation.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from torch.nn.parameter import Parameter
+
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_BLOCK_SIZE,
+    MXFP8_SCALE_DTYPE,
+    dequant_mxfp8_to_bf16,
+)
+
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
+
+
+class EmulationMxfp8LinearKernel(Mxfp8LinearKernel):
+    """Software emulation fallback for MXFP8 (dequant to BF16)."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight.data  # [N, K]
+        N, K = weight.shape
+        scale_k = K // MXFP8_BLOCK_SIZE
+
+        weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
+
+        # Dequantize MXFP8 -> BF16 ONCE here, at load time, so apply_weights runs
+        # a plain BF16 linear with no per-step dequant -- i.e. run as if from a
+        # BF16 checkpoint. The 1-byte MXFP8 weight is replaced by BF16 (2x its
+        # size, but linear weights are small vs the MoE experts); the tiny E8M0
+        # scale is kept for the dtype/ndim asserts but is otherwise unused.
+        # Opt out (VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD=0) to keep the MXFP8
+        # weight and dequant per-step in apply_weights instead.
+        import vllm.envs as envs
+
+        if envs.VLLM_MXFP8_EMULATION_DEQUANT_AT_LOAD:
+            weight = dequant_mxfp8_to_bf16(weight.contiguous(), weight_scale)
+        layer.weight = Parameter(weight.contiguous(), requires_grad=False)
+        layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        weight = layer.weight
+        # Load-time dequant path: weights are already BF16/FP16 (>= 2-byte), so
+        # run a plain linear -- no per-step dequant. (MXFP8 weights are 1-byte.)
+        if weight.element_size() >= 2:
+            # F.linear requires x and weight share a dtype; .to() is a no-op when
+            # they already match (e.g. both BF16).
+            output = torch.nn.functional.linear(x, weight.to(x.dtype), bias)
+            return output.to(x.dtype)
+
+        # Fallback: weights still in MXFP8 -- dequant on the fly (other archs /
+        # if a future caller skips the load-time conversion above).
+        weight_scale = layer.weight_scale
+        if weight_scale.dtype != MXFP8_SCALE_DTYPE:
+            raise ValueError(
+                f"Emulation backend requires {MXFP8_SCALE_DTYPE} "
+                f"weight_scale dtype, got {weight_scale.dtype}."
+            )
+        if weight_scale.ndim != 2:
+            raise ValueError(
+                f"Emulation backend requires 2D weight_scale, "
+                f"got {weight_scale.ndim}D. "
+                f"Ensure process_weights_after_loading was called."
+            )
+
+        # Cast to x's dtype: dequant yields BF16, but F.linear needs both operands
+        # to match (e.g. an FP16 model). No-op when x is already BF16.
+        weight_bf16 = dequant_mxfp8_to_bf16(weight, weight_scale).to(x.dtype)
+        output = torch.nn.functional.linear(x, weight_bf16, bias)
+        return output.to(x.dtype)

vllm/model_executor/kernels/linear/mxfp8/flashinfer.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+from torch.nn.parameter import Parameter
+
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_BLOCK_SIZE,
+    mxfp8_e4m3_quantize,
+    swizzle_mxfp8_scale,
+)
+from vllm.platforms import current_platform
+from vllm.utils import flashinfer as vllm_flashinfer
+from vllm.utils.flashinfer import has_flashinfer_cutedsl
+
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
+
+
+class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel):
+    """MXFP8 W8A8 GEMM via FlashInfer CUTLASS (SM100+)."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if current_platform.has_device_capability(100):
+            return True, None
+        return False, "requires >=sm_100 (Blackwell)"
+
+    @classmethod
+    def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight.data  # [N, K]
+        N, K = weight.shape
+
+        scale_k = K // MXFP8_BLOCK_SIZE
+        weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous()
+        weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K)
+
+        layer.weight = Parameter(weight.contiguous(), requires_grad=False)
+        layer.weight_scale = Parameter(
+            weight_scale_swizzled.contiguous(), requires_grad=False
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        weight = layer.weight
+        weight_scale = layer.weight_scale
+        out_dtype = x.dtype
+        N, K = weight.shape
+
+        input_shape = x.shape
+        input_2d = x.view(-1, K)
+        min_dim = 128
+
+        assert min_dim <= K, (
+            f"mm_mxfp8 requires K >= {min_dim}, got K={K}. "
+            f"in_features is too small for mm_mxfp8."
+        )
+        assert K % MXFP8_BLOCK_SIZE == 0, (
+            f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}."
+        )
+        assert min_dim <= N, (
+            f"mm_mxfp8 requires N >= {min_dim}, got N={N}. "
+            f"out_features is too small for mm_mxfp8."
+        )
+
+        input_mxfp8, input_scale = mxfp8_e4m3_quantize(
+            input_2d, is_sf_swizzled_layout=True
+        )
+
+        if not weight.is_contiguous():
+            weight = weight.contiguous()
+
+        output = vllm_flashinfer.mm_mxfp8(
+            input_mxfp8,
+            weight.t(),
+            input_scale,
+            weight_scale,
+            out_dtype=out_dtype,
+            backend="cutlass",
+        )
+
+        if bias is not None:
+            output = output + bias
+
+        output_shape = (*input_shape[:-1], N)
+        return output.view(output_shape)
+
+
+class FlashInferCutedslMxfp8LinearKernel(Mxfp8LinearKernel):
+    """MXFP8 W8A8 GEMM via FlashInfer CuTe-DSL (SM100/SM103)."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not (
+            current_platform.is_cuda()
+            and current_platform.is_device_capability_family(100)
+        ):
+            return False, "requires sm_100/sm_103 (Blackwell)"
+        if not has_flashinfer_cutedsl():
+            return False, "requires FlashInfer CuTe-DSL module"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight.data  # [N, K]
+        N, K = weight.shape
+
+        scale_k = K // MXFP8_BLOCK_SIZE
+        weight_scale_2d = layer.weight_scale.data[:N, :scale_k].contiguous()
+        weight_scale_swizzled = swizzle_mxfp8_scale(weight_scale_2d, M=N, K=K)
+
+        # Store weight column-major [K, N] as mm_mxfp8 expects for operand B.
+        layer.weight = Parameter(weight.contiguous().t(), requires_grad=False)
+        layer.weight_scale = Parameter(
+            weight_scale_swizzled.contiguous(), requires_grad=False
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        weight = layer.weight  # [K, N], column-major
+        weight_scale = layer.weight_scale
+        out_dtype = x.dtype
+        K, N = weight.shape
+
+        input_shape = x.shape
+        input_2d = x.view(-1, K)
+        min_dim = 128
+
+        assert min_dim <= K, (
+            f"mm_mxfp8 requires K >= {min_dim}, got K={K}. "
+            f"in_features is too small for mm_mxfp8."
+        )
+        assert K % MXFP8_BLOCK_SIZE == 0, (
+            f"mm_mxfp8 requires K to be divisible by {MXFP8_BLOCK_SIZE}, got K={K}."
+        )
+        assert min_dim <= N, (
+            f"mm_mxfp8 requires N >= {min_dim}, got N={N}. "
+            f"out_features is too small for mm_mxfp8."
+        )
+
+        input_mxfp8, input_scale = mxfp8_e4m3_quantize(
+            input_2d, is_sf_swizzled_layout=True
+        )
+
+        output = vllm_flashinfer.mm_mxfp8(
+            input_mxfp8,
+            weight,
+            input_scale,
+            weight_scale,
+            out_dtype=out_dtype,
+            backend="cute-dsl",
+        )
+
+        if bias is not None:
+            output = output + bias
+
+        output_shape = (*input_shape[:-1], N)
+        return output.view(output_shape)

vllm/model_executor/kernels/linear/mxfp8/humming.py

@@ -0,0 +1,63 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.humming_utils import (
+    convert_linear_layer_to_humming_standard,
+    prepare_humming_layer,
+)
+from vllm.platforms import current_platform
+
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
+
+
+class HummingMxfp8LinearKernel(Mxfp8LinearKernel):
+    """Humming GEMM Kernel for MXFP8."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "Humming only supported on CUDA"
+
+        if not current_platform.has_device_capability(75):
+            return False, "Humming only supported on SM75+"
+
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.weight_scale.data = layer.weight_scale.data.view(torch.float8_e8m0fnu)
+        name_map = {"weight": "weight", "weight_scale": "weight_scale"}
+
+        quant_config = {
+            "quant_method": "humming",
+            "dtype": "float8e4m3",
+            "scale_dtype": "float8e8m0",
+            "group_size": 32,
+            "weight_scale_type": "group",
+        }
+
+        convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
+        prepare_humming_layer(layer, quant_config)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        from vllm.utils.humming import HummingMethod
+
+        flatten_inputs = x.view(-1, x.size(-1))
+        output = HummingMethod.forward_layer(
+            layer=layer,
+            inputs=flatten_inputs,
+            compute_config=layer.compute_config,
+        )
+        return output.view(*x.shape[:-1], output.size(-1))

vllm/model_executor/kernels/linear/mxfp8/marlin.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
+
+
+class MarlinMxfp8LinearKernel(Mxfp8LinearKernel):
+    """MXFP8 W8A16 GEMM via Marlin (SM80+)."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
+            is_fp8_marlin_supported,
+        )
+
+        if is_fp8_marlin_supported():
+            return True, None
+        return False, "Marlin FP8 not available"
+
+    @classmethod
+    def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
+            prepare_mxfp8_layer_for_marlin,
+        )
+
+        prepare_mxfp8_layer_for_marlin(layer)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
+            apply_mxfp8_marlin_linear,
+        )
+
+        return apply_mxfp8_marlin_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale,
+            workspace=layer.workspace,
+            size_n=layer.output_size_per_partition,
+            size_k=layer.input_size_per_partition,
+            bias=bias,
+        )

vllm/model_executor/kernels/linear/mxfp8/rocm_native.py

@@ -0,0 +1,247 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Native MXFP8 linear GEMM for AMD CDNA4 (gfx950) via Triton ``tl.dot_scaled``.
+
+Consumes the FP8 E4M3 weights + E8M0 block scales directly (no dequant-to-BF16);
+activations are MXFP8-quantized per token. Uses the CDNA4 hardware microscaling
+matrix cores. Falls back (via the kernel selector) to the BF16
+``EmulationMxfp8LinearKernel`` on archs without native MX or for shapes with
+``K % 128 != 0``.
+"""
+
+import torch
+from torch.nn.parameter import Parameter
+
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_BLOCK_SIZE,
+    MXFP8_SCALE_DTYPE,
+    dequant_mxfp8_to_bf16,
+    mxfp8_e4m3_quantize,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
+
+
+@triton.jit
+def _mxfp8_linear_kernel(
+    x_ptr,
+    xs_ptr,
+    w_ptr,
+    ws_ptr,
+    out_ptr,
+    M,
+    N,
+    K,
+    stride_xm,
+    stride_xk,
+    stride_xsm,
+    stride_xsk,
+    stride_wn,
+    stride_wk,
+    stride_wsn,
+    stride_wsk,
+    stride_om,
+    stride_on,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_K)
+    offs_sk = tl.arange(0, BLOCK_K // 32)
+    m_mask = offs_m < M
+    n_mask = offs_n < N
+
+    x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk
+    xs_ptrs = xs_ptr + offs_m[:, None] * stride_xsm + offs_sk[None, :] * stride_xsk
+    w_ptrs = w_ptr + offs_n[:, None] * stride_wn + offs_k[None, :] * stride_wk
+    ws_ptrs = ws_ptr + offs_n[:, None] * stride_wsn + offs_sk[None, :] * stride_wsk
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+    for _ in range(0, tl.cdiv(K, BLOCK_K)):
+        x = tl.load(x_ptrs, mask=m_mask[:, None], other=0.0)
+        w = tl.load(w_ptrs, mask=n_mask[:, None], other=0.0)
+        xs = tl.load(xs_ptrs, mask=m_mask[:, None], other=0)
+        ws = tl.load(ws_ptrs, mask=n_mask[:, None], other=0)
+        acc += tl.dot_scaled(x, xs, "e4m3", w.T, ws, "e4m3")
+        x_ptrs += BLOCK_K * stride_xk
+        w_ptrs += BLOCK_K * stride_wk
+        xs_ptrs += (BLOCK_K // 32) * stride_xsk
+        ws_ptrs += (BLOCK_K // 32) * stride_wsk
+
+    o_ptrs = out_ptr + offs_m[:, None] * stride_om + offs_n[None, :] * stride_on
+    tl.store(
+        o_ptrs, acc.to(out_ptr.dtype.element_ty), mask=m_mask[:, None] & n_mask[None, :]
+    )
+
+
+def _mxfp8_dot_scaled_linear(
+    x: torch.Tensor,  # [M, K] bf16/fp16
+    w: torch.Tensor,  # [N, K] fp8 e4m3
+    w_scale: torch.Tensor,  # [N, K//32] uint8 (E8M0)
+) -> torch.Tensor:
+    M, K = x.shape
+    N = w.shape[0]
+    x_q, x_scale = mxfp8_e4m3_quantize(x)
+    out = torch.empty((M, N), dtype=x.dtype, device=x.device)
+    BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages = _select_cfg(M, N, K)
+    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
+    _mxfp8_linear_kernel[grid](
+        x_q,
+        x_scale,
+        w,
+        w_scale,
+        out,
+        M,
+        N,
+        K,
+        x_q.stride(0),
+        x_q.stride(1),
+        x_scale.stride(0),
+        x_scale.stride(1),
+        w.stride(0),
+        w.stride(1),
+        w_scale.stride(0),
+        w_scale.stride(1),
+        out.stride(0),
+        out.stride(1),
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+        num_warps=num_warps,
+        num_stages=num_stages,
+    )
+    return out
+
+
+def _select_cfg(M, N, K):
+    """(BLOCK_M, BLOCK_N, BLOCK_K, num_warps, num_stages) — graph-tuned on gfx950.
+
+    The M-bucketed, shape-adaptive tile selection here is the speedup over the
+    upstream 2-bucket launcher. Tiles are pipelined (num_stages>=2, larger BLOCK_K)
+    and occupancy- and shape-aware: keyed on the LOCAL (M, N, K), so it adapts to the
+    TP-sharded shapes (e.g. MiniMax-M3 TP=4 vs TP=8, where local N and K differ) —
+    large-K prefill uses BLOCK_K=256; short-K (K=768) widens N. BLOCK_K must divide K
+    (the K-loop is unmasked), so every BLOCK_K below is guarded to be K-divisible
+    (served K: 384/768/1024/2048/6144).
+    """
+    if M <= 64:
+        # decode (M in {1,32,64}): tiny-M GEMV is weight-BW + GPU-OCCUPANCY bound. The
+        # lever is NARROW BLOCK_N=16 (maximize N-tiles so more CUs stream the weight in
+        # parallel) + LARGE BLOCK_K (fewer K-iters, bigger coalesced weight loads).
+        # Tuned by CUDA-graph replay latency. Optimal at both TP=4 and TP=8.
+        if K % 1024 == 0:  # K=2048, 6144 -> graph-best 16x16x1024 (all M)
+            return 16, 16, 1024, 2, 2
+        if K % 512 == 0:
+            return 16, 16, 512, 2, 3
+        if K % 256 == 0:  # K=768 (shared_down) -> graph-best 16x32x256
+            return 16, 32, 256, 4, 3
+        return 16, 32, 128, 4, 3
+    # mid-M (65..256) on SMALL local-N: still occupancy-bound (a 64x64 tile makes too
+    # few N-tiles), so the narrow-BLOCK_N decode-style tile fills the CUs better.
+    # N<=1536 covers the real fused-qkv local N at TP=8: q heads shard but the GQA KV
+    # (4) + sparse-indexer (4) heads are < TP=8, so vLLM replicates them to 1/rank ->
+    # N = 1024 + 4*128 = 1536 (not 2560/2=1280). For the wider 1280<N<=1536 band the
+    # narrow tile only wins up to M=128 (at M=256 the 64x64 tile is better), so cap it
+    # there; N<=1280 keeps the narrow tile through M=256. TP=4 qkv N=2560 is unchanged.
+    if (M <= 256 and N <= 1280) or (M <= 128 and N <= 1536):
+        if K % 1024 == 0:
+            return 16, 16, 1024, 2, 2
+        if K % 512 == 0:
+            return 16, 16, 512, 2, 3
+        if K % 256 == 0:
+            return 16, 16, 256, 2, 3
+        return 16, 16, 128, 2, 3
+    # right-sized launch grid (host-side), used to gate the tall 256-BLOCK_M tile.
+    occ = triton.cdiv(M, 256) * triton.cdiv(N, 128)
+    if K <= 1024:  # short-K (shared_down K=384/768; TP=8 o_proj K=1024)
+        if M <= 256:
+            return (64, 64, 256, 8, 2) if K % 256 == 0 else (64, 64, 128, 8, 2)
+        # large prefill. (The former 256x128x256 tall tile was faster only on triton
+        # 3.6; on triton 3.7 its large BLOCK_M register/LDS footprint spills or hits
+        # "out of resources", so use 128x128x256 -- within the known-good footprint.)
+        if M >= 4096 and K >= 1024 and K % 256 == 0 and occ >= 256:
+            return 128, 128, 256, 8, 3
+        return 128, 128, 128, 8, 3
+    # large-K (K >= 2048). BLOCK_K is K-divisibility-guarded (the K-loop is unmasked):
+    # served large-K is 2048/6144 (%256==0), but fall back to 128 (always divides, since
+    # the entry requires K%128==0) for any other K to stay correct.
+    if M <= 256:  # conc~128 decode + small prefill chunk: occupancy tile
+        if K % 512 == 0:
+            return 64, 64, 512, 8, 2
+        return (64, 64, 256, 8, 2) if K % 256 == 0 else (64, 64, 128, 8, 2)
+    if M <= 1024:  # medium prefill chunk: BN=64 keeps small-N occupied
+        return (128, 64, 256, 8, 3) if K % 256 == 0 else (128, 64, 128, 8, 3)
+    # large prefill (M > 1024). The previously graph-tuned tall 256x128x256 and deep
+    # 128x128x512 tiles won only on triton 3.6; on triton 3.7 their larger BLOCK_M /
+    # BLOCK_K register+LDS footprint spills (or hits "out of resources" on stricter
+    # ROCm/triton builds). The 128x128x256 tile is equal-or-faster on triton 3.7 (the
+    # M=4096,N=2048,K=6144 shape: ~104us vs the 256x128x256 tile's ~184us), within ~5%
+    # on 3.6, and inside the footprint of the tiles used elsewhere in this selector.
+    # Covers the qkv-class local N=1536 (TP=8 qkv / TP=4 shared_gate_up) and the deep-K
+    # / very-large-M shapes.
+    if K % 256 == 0 and (1280 < N <= 1536 or (occ >= 128 and (K >= 4096 or M >= 4096))):
+        return 128, 128, 256, 8, 3
+    # small local-N (e.g. TP=8 shared_gate_up N=768): a 64-wide BLOCK_N doubles the
+    # N-tile count -> better CU fill than 128x128 at this mid-large M (~1.4x there).
+    if N <= 1024 and K % 256 == 0:
+        return 128, 64, 256, 8, 3
+    return (128, 128, 256, 8, 2) if K % 256 == 0 else (128, 256, 128, 8, 3)
+
+
+class RocmDotScaledMxfp8LinearKernel(Mxfp8LinearKernel):
+    """Native CDNA4 (gfx950) MXFP8 linear via Triton ``tl.dot_scaled``."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_rocm():
+            return False, "not ROCm"
+        # supports_mx() == gfx95x (CDNA4 native microscaling hardware). On other
+        # archs dot_scaled would upcast to BF16, so the kernel selector falls
+        # through to the BF16 emulation (hipBLASLt) path instead.
+        if not current_platform.supports_mx():
+            return False, "native MX requires CDNA4 (gfx95x)"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight = layer.weight.data  # [N, K] fp8
+        N, K = weight.shape
+        scale_k = K // MXFP8_BLOCK_SIZE
+        weight_scale = layer.weight_scale.data[:N, :scale_k].contiguous()
+        layer.weight = Parameter(weight.contiguous(), requires_grad=False)
+        layer.weight_scale = Parameter(weight_scale, requires_grad=False)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if layer.weight_scale.dtype != MXFP8_SCALE_DTYPE:
+            raise ValueError(
+                f"Expected {MXFP8_SCALE_DTYPE} weight_scale, got "
+                f"{layer.weight_scale.dtype}."
+            )
+        out_shape = (*x.shape[:-1], layer.weight.shape[0])
+        x2d = x.reshape(-1, x.shape[-1])
+        if x2d.shape[-1] % 128 == 0:
+            out = _mxfp8_dot_scaled_linear(x2d, layer.weight, layer.weight_scale)
+        else:
+            # dot_scaled tiling needs K % 128 == 0; dequantize fallback otherwise.
+            w_bf16 = dequant_mxfp8_to_bf16(layer.weight, layer.weight_scale)
+            out = torch.nn.functional.linear(x2d, w_bf16).to(x.dtype)
+        out = out.reshape(out_shape)
+        if bias is not None:
+            out = out + bias
+        return out

vllm/model_executor/kernels/linear/mxfp8/xpu.py

@@ -0,0 +1,51 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    xpu_mxfp8_quantize as quant_mxfp8,
+)
+from vllm.model_executor.utils import replace_parameter
+from vllm.platforms import current_platform
+
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
+
+
+class XPUMxFp8LinearKernel(Mxfp8LinearKernel):
+    """MXFP8 W8A8 GEMM on XPU."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_xpu():
+            return False, "XPUMxFp8 only support on XPU"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight_scale = layer.weight_scale.view(torch.float8_e8m0fnu)
+        weight_scale = weight_scale.t().contiguous()
+        replace_parameter(layer, "weight", layer.weight.t())
+        replace_parameter(layer, "weight_scale", weight_scale.data)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        out_dtype = x.dtype
+        x_fp8, x_scale = quant_mxfp8(x)
+        return torch.ops._xpu_C.fp8_gemm(
+            x_fp8,
+            layer.weight,
+            out_dtype,
+            x_scale,
+            layer.weight_scale,
+            bias,
+        )

vllm/model_executor/kernels/linear/nvfp4/__init__.py

@@ -0,0 +1,12 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.model_executor.kernels.linear.nvfp4.base import (
+    NvFp4LinearKernel,
+    NvFp4LinearLayerConfig,
+)
+
+__all__ = [
+    "NvFp4LinearKernel",
+    "NvFp4LinearLayerConfig",
+]

vllm/model_executor/kernels/linear/nvfp4/base.py

@@ -0,0 +1,76 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
+
+
+@dataclass
+class NvFp4LinearLayerConfig:
+    """Configuration for an NVFP4 linear layer.
+
+    All NVFP4 layers share the same structure: packed uint8 weights (2 FP4 values per
+    byte), FP8-E4M3 per-block weight scales (group size 16), and scalar global
+    scales for both weights and activations.
+    """
+
+    pass
+
+
+class NvFp4LinearKernel(ABC):
+    """Base class for NVFP4 quantized linear kernels.
+
+    Each subclass implements a specific GEMM backend (CUTLASS, Marlin, etc).
+    The kernel selection mechanism iterates over registered subclasses in
+    priority order,calling ``is_supported`` and ``can_implement`` to find the best
+    match for the current hardware.
+    """
+
+    def __init__(self, config: NvFp4LinearLayerConfig) -> None:
+        assert self.can_implement(config)[0]
+        assert self.is_supported()[0]
+        self.config = config
+
+    def input_quant_key(self) -> QuantKey | None:
+        """Return the input quantization key supported by this kernel. If the kernel
+        does not support input quantization outside of the kernel, return None.
+        """
+        return None
+
+    @classmethod
+    @abstractmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        """Return whether this kernel can run on the current platform."""
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        """Return whether this kernel can handle *config*."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Transform weights into the format required by this kernel.
+
+        Called once after checkpoint weights have been loaded onto the
+        device.  Implementations should repack / swizzle / pad weights
+        and scales in-place on *layer*.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Run the quantized GEMM."""
+        raise NotImplementedError

vllm/model_executor/kernels/linear/nvfp4/cutlass.py

@@ -0,0 +1,77 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm._custom_ops import (
+    cutlass_scaled_fp4_mm,
+    scaled_fp4_quant,
+)
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
+    cutlass_fp4_supported,
+    pad_nvfp4_weight_for_cutlass,
+    slice_nvfp4_output,
+    swizzle_blockscale,
+)
+
+from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
+
+
+class CutlassNvFp4LinearKernel(NvFp4LinearKernel):
+    """NVFP4 GEMM via the vLLM CUTLASS kernel."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if cutlass_fp4_supported():
+            return True, None
+        return False, "CUTLASS FP4 kernels not available"
+
+    @classmethod
+    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.weight_scale = torch.nn.Parameter(
+            swizzle_blockscale(layer.weight_scale.data), requires_grad=False
+        )
+        padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
+            layer.weight.data
+        )
+        layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
+        layer.weights_padding_cols = weights_padding_cols
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        output_size = layer.output_size_per_partition
+        output_dtype = x.dtype
+        output_shape = [*x.shape[:-1], output_size]
+        weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)
+
+        x_fp4, x_blockscale = scaled_fp4_quant(
+            x,
+            layer.input_global_scale_inv,
+            is_sf_swizzled_layout=True,
+            backend="cutlass",
+            padded_n=x.shape[-1] + weights_padding_bytes * 2,
+        )
+
+        out = cutlass_scaled_fp4_mm(
+            x_fp4,
+            layer.weight,
+            x_blockscale,
+            layer.weight_scale,
+            layer.alpha,
+            output_dtype,
+        )
+
+        out = slice_nvfp4_output(out, output_size)
+
+        if bias is not None:
+            out = out + bias
+        return out.view(*output_shape)

vllm/model_executor/kernels/linear/nvfp4/emulation.py

@@ -0,0 +1,49 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
+    kE2M1ToFloat_handle,
+    run_nvfp4_emulations,
+)
+
+from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
+
+
+class EmulationNvFp4LinearKernel(NvFp4LinearKernel):
+    """Software emulation fallback for NVFP4 (dequant → BF16 matmul)."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        # Always available as a last-resort fallback.
+        return True, None
+
+    @classmethod
+    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Move the E2M1 lookup table to the device now, because
+        # `.to(device)` is not allowed during CUDA graph capture.
+        kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(layer.weight.device)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        out = run_nvfp4_emulations(
+            x=x,
+            input_global_scale=layer.input_global_scale_inv,
+            weight=layer.weight,
+            weight_scale_swizzled=layer.weight_scale,
+            weight_global_scale=layer.weight_global_scale,
+            swizzle=False,
+        )
+        if bias is not None:
+            out = out + bias
+        return out

vllm/model_executor/kernels/linear/nvfp4/fbgemm.py

@@ -0,0 +1,69 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm._custom_ops import scaled_fp4_quant
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
+    slice_nvfp4_output,
+    swizzle_blockscale,
+)
+from vllm.utils.import_utils import has_fbgemm_gpu
+
+from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
+
+
+class FbgemmNvFp4LinearKernel(NvFp4LinearKernel):
+    """NVFP4 GEMM via FBGEMM."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if has_fbgemm_gpu():
+            return True, None
+        return False, "fbgemm_gpu required"
+
+    @classmethod
+    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        swizzled = swizzle_blockscale(layer.weight_scale.data)
+        layer.weight_scale = torch.nn.Parameter(
+            swizzled.view(-1).view(torch.uint8), requires_grad=False
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        import fbgemm_gpu  # noqa: F401 - registers torch.ops.fbgemm.*
+
+        output_size = layer.output_size_per_partition
+        output_dtype = x.dtype
+        output_shape = [*x.shape[:-1], output_size]
+
+        x_fp4, x_blockscale = scaled_fp4_quant(
+            x,
+            layer.input_global_scale_inv,
+            is_sf_swizzled_layout=True,
+            backend="fbgemm",
+        )
+
+        out = torch.ops.fbgemm.f4f4bf16(
+            x_fp4,
+            layer.weight,
+            x_blockscale.view(-1).view(torch.uint8),
+            layer.weight_scale,
+            layer.alpha,
+            use_mx=False,
+        ).to(output_dtype)
+
+        out = slice_nvfp4_output(out, output_size)
+
+        if bias is not None:
+            out = out + bias
+        return out.view(*output_shape)

vllm/model_executor/kernels/linear/nvfp4/flashinfer.py

@@ -0,0 +1,372 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm._custom_ops import scaled_fp4_quant
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    as_quantized_activation,
+)
+from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
+    pad_nvfp4_activation_for_cutlass,
+    pad_nvfp4_weight_for_cutlass,
+    slice_nvfp4_output,
+    swizzle_blockscale,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    kNvfp4Dynamic,
+)
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import (
+    flashinfer_scaled_fp4_mm,
+    has_flashinfer,
+    has_flashinfer_b12x_gemm,
+)
+
+from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
+
+
+class FlashInferCuteDslNvFp4LinearKernel(NvFp4LinearKernel):
+    """NVFP4 GEMM via FlashInfer's cutedsl backend."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_device_capability_family(100):
+            return False, "FlashInfer cutedsl requires sm_10x"
+        if not has_flashinfer():
+            return False, "FlashInfer required"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # cutedsl uses the same swizzled + padded layout as cutlass.
+        layer.weight_scale = torch.nn.Parameter(
+            swizzle_blockscale(layer.weight_scale.data), requires_grad=False
+        )
+        padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
+            layer.weight.data
+        )
+        layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
+        layer.weights_padding_cols = weights_padding_cols
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        output_size = layer.output_size_per_partition
+        output_dtype = x.dtype
+        output_shape = [*x.shape[:-1], output_size]
+
+        x_fp4, x_blockscale = scaled_fp4_quant(
+            x,
+            layer.input_global_scale_inv,
+            is_sf_swizzled_layout=True,
+            backend="flashinfer-cutedsl",
+        )
+
+        x_fp4 = pad_nvfp4_activation_for_cutlass(
+            x_fp4, getattr(layer, "weights_padding_cols", 0)
+        )
+
+        out = flashinfer_scaled_fp4_mm(
+            x_fp4,
+            layer.weight,
+            x_blockscale,
+            layer.weight_scale,
+            layer.alpha,
+            output_dtype,
+            backend="cute-dsl",
+        )
+
+        out = slice_nvfp4_output(out, output_size)
+
+        if bias is not None:
+            out = out + bias
+        return out.view(*output_shape)
+
+
+class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel):
+    """NVFP4 GEMM via FlashInfer's CUTLASS wrapper."""
+
+    def input_quant_key(self) -> QuantKey | None:
+        """This kernel supports dynamic quantization of the input. By
+        convention, pre-quantized blockscales must use the swizzled layout."""
+        return kNvfp4Dynamic
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
+            cutlass_fp4_supported,
+        )
+
+        if (
+            cutlass_fp4_supported()
+            and current_platform.has_device_capability(100)
+            and has_flashinfer()
+        ):
+            return True, None
+        return False, "FlashInfer + >=sm_100 required"
+
+    @classmethod
+    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.weight_scale = torch.nn.Parameter(
+            swizzle_blockscale(layer.weight_scale.data), requires_grad=False
+        )
+        padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
+            layer.weight.data
+        )
+        layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
+        layer.weights_padding_cols = weights_padding_cols
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor | QuantizedActivation,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        output_size = layer.output_size_per_partition
+        weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)
+
+        qa = as_quantized_activation(x, self.input_quant_key())
+        if qa is not None:
+            x_fp4, x_blockscale = qa.data, qa.scale
+            x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_bytes)
+            output_dtype = qa.orig_dtype
+            output_shape = [*qa.orig_shape[:-1], output_size]
+        else:
+            assert isinstance(x, torch.Tensor)
+            output_dtype = x.dtype
+            output_shape = [*x.shape[:-1], output_size]
+            x_fp4, x_blockscale = scaled_fp4_quant(
+                x,
+                layer.input_global_scale_inv,
+                is_sf_swizzled_layout=True,
+                backend="flashinfer-cutlass",
+                padded_n=x.shape[-1] + weights_padding_bytes * 2,
+            )
+
+        out = flashinfer_scaled_fp4_mm(
+            x_fp4,
+            layer.weight,
+            x_blockscale,
+            layer.weight_scale,
+            layer.alpha,
+            output_dtype,
+            backend="cutlass",
+        )
+
+        out = slice_nvfp4_output(out, output_size)
+
+        if bias is not None:
+            out = out + bias
+        return out.view(*output_shape)
+
+
+class FlashInferTrtllmNvFp4LinearKernel(NvFp4LinearKernel):
+    """NVFP4 GEMM via FlashInfer's TensorRT-LLM wrapper."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if has_flashinfer():
+            return True, None
+        return False, "FlashInfer required"
+
+    @classmethod
+    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a
+
+        weight = layer.weight.data
+        weight_scale = layer.weight_scale.data
+        epilogue_tile_m = 128
+
+        layer.weight = torch.nn.Parameter(
+            shuffle_matrix_a(weight.view(torch.uint8), epilogue_tile_m),
+            requires_grad=False,
+        )
+        layer.weight_scale = torch.nn.Parameter(
+            shuffle_matrix_sf_a(weight_scale.view(torch.uint8), epilogue_tile_m)
+            .reshape(weight_scale.shape)
+            .view(torch.float8_e4m3fn),
+            requires_grad=False,
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        output_size = layer.output_size_per_partition
+        output_dtype = x.dtype
+        output_shape = [*x.shape[:-1], output_size]
+
+        x_fp4, x_blockscale = scaled_fp4_quant(
+            x,
+            layer.input_global_scale_inv,
+            is_sf_swizzled_layout=True,
+            backend="flashinfer-trtllm",
+        )
+
+        out = flashinfer_scaled_fp4_mm(
+            x_fp4,
+            layer.weight,
+            x_blockscale,
+            layer.weight_scale,
+            layer.alpha,
+            output_dtype,
+            backend="trtllm",
+        )
+
+        out = slice_nvfp4_output(out, output_size)
+
+        if bias is not None:
+            out = out + bias
+        return out.view(*output_shape)
+
+
+class FlashInferCudnnNvFp4LinearKernel(NvFp4LinearKernel):
+    """NVFP4 GEMM via FlashInfer's cuDNN wrapper."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if has_flashinfer():
+            return True, None
+        return False, "FlashInfer required"
+
+    @classmethod
+    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # cuDNN uses the same swizzled + padded layout as CUTLASS
+        layer.weight_scale = torch.nn.Parameter(
+            swizzle_blockscale(layer.weight_scale.data), requires_grad=False
+        )
+        padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
+            layer.weight.data
+        )
+        layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
+        layer.weights_padding_cols = weights_padding_cols
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        output_size = layer.output_size_per_partition
+        output_dtype = x.dtype
+        output_shape = [*x.shape[:-1], output_size]
+        weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)
+
+        x_fp4, x_blockscale = scaled_fp4_quant(
+            x,
+            layer.input_global_scale_inv,
+            is_sf_swizzled_layout=True,
+            backend="flashinfer-cudnn",
+            padded_n=x.shape[-1] + weights_padding_bytes * 2,
+        )
+
+        out = flashinfer_scaled_fp4_mm(
+            x_fp4,
+            layer.weight,
+            x_blockscale,
+            layer.weight_scale,
+            layer.alpha,
+            output_dtype,
+            backend="cudnn",
+        )
+
+        out = slice_nvfp4_output(out, output_size)
+
+        if bias is not None:
+            out = out + bias
+        return out.view(*output_shape)
+
+
+class FlashInferB12xNvFp4LinearKernel(NvFp4LinearKernel):
+    """NVFP4 GEMM via FlashInfer's b12x CuTe DSL warp-level MMA kernel (SM120+)."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if current_platform.has_device_capability(120) and has_flashinfer_b12x_gemm():
+            return True, None
+        return (
+            False,
+            "FlashInfer b12x requires SM120+ and FlashInfer "
+            "with Sm120BlockScaledDenseGemmKernel",
+        )
+
+    @classmethod
+    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        layer.weight_scale = torch.nn.Parameter(
+            swizzle_blockscale(layer.weight_scale.data), requires_grad=False
+        )
+        padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
+            layer.weight.data
+        )
+        layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
+        layer.weights_padding_cols = weights_padding_cols
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        output_size = layer.output_size_per_partition
+        output_dtype = x.dtype
+        output_shape = [*x.shape[:-1], output_size]
+
+        x_fp4, x_blockscale = scaled_fp4_quant(
+            x,
+            layer.input_global_scale_inv,
+            is_sf_swizzled_layout=True,
+            backend="b12x",
+        )
+
+        x_fp4 = pad_nvfp4_activation_for_cutlass(
+            x_fp4, getattr(layer, "weights_padding_cols", 0)
+        )
+
+        out = flashinfer_scaled_fp4_mm(
+            x_fp4,
+            layer.weight,
+            x_blockscale,
+            layer.weight_scale,
+            layer.alpha,
+            output_dtype,
+            backend="b12x",
+        )
+
+        out = slice_nvfp4_output(out, output_size)
+
+        if bias is not None:
+            out = out + bias
+        return out.view(*output_shape)

vllm/model_executor/kernels/linear/nvfp4/humming.py

@@ -0,0 +1,73 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.utils.humming_utils import (
+    prepare_humming_layer,
+)
+from vllm.platforms import current_platform
+
+from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
+
+logger = init_logger(__name__)
+
+
+class HummingNvFp4LinearKernel(NvFp4LinearKernel):
+    """Humming GEMM Kernel for NVFP4."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "Humming only supported on CUDA"
+
+        if not current_platform.has_device_capability(75):
+            return False, "Humming only supported on SM75+"
+
+        return True, None
+
+    @classmethod
+    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Route through humming's compressed-tensors nvfp4 loader (same path as
+        # the MoE oracle); the native group_tensor schema mishandles a scalar
+        # global scale.
+        quant_config = {
+            "quant_method": "compressed-tensors",
+            "format": "nvfp4-pack-quantized",
+            "type": "float",
+            "num_bits": 4,
+            "strategy": "group",
+            "group_size": 16,
+        }
+        # CT pack-quantized reads `weight_packed`; the scheme renamed it to `weight`.
+        if not hasattr(layer, "weight_packed"):
+            layer.weight_packed = layer.weight
+            del layer.weight
+        # The CT linear scheme inverts the global scale (1/scale) for
+        # marlin/cutlass; humming wants the original.
+        layer.weight_global_scale = torch.nn.Parameter(
+            1.0 / layer.weight_global_scale, requires_grad=False
+        )
+        prepare_humming_layer(layer, quant_config)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        from vllm.utils.humming import HummingMethod
+
+        flatten_inputs = x.view(-1, x.size(-1))
+        output = HummingMethod.forward_layer(
+            layer=layer,
+            inputs=flatten_inputs,
+            compute_config=layer.compute_config,
+        )
+        return output.view(*x.shape[:-1], output.size(-1))

vllm/model_executor/kernels/linear/nvfp4/marlin.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
+    apply_fp4_marlin_linear,
+    is_fp4_marlin_supported,
+    prepare_fp4_layer_for_marlin,
+)
+
+from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
+
+logger = init_logger(__name__)
+
+
+class MarlinNvFp4LinearKernel(NvFp4LinearKernel):
+    """NVFP4 weight-only GEMM via Marlin (W4A16)."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if is_fp4_marlin_supported():
+            return True, None
+        return False, "Marlin FP4 not available"
+
+    @classmethod
+    def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        logger.warning_once(
+            "Your GPU does not have native support for FP4 computation but "
+            "FP4 quantization is being used. Weight-only FP4 compression "
+            "will be used leveraging the Marlin kernel. This may degrade "
+            "performance for compute-heavy workloads."
+        )
+        prepare_fp4_layer_for_marlin(layer)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return apply_fp4_marlin_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=layer.weight_scale,
+            weight_global_scale=layer.weight_global_scale,
+            workspace=layer.workspace,
+            size_n=layer.output_size_per_partition,
+            size_k=layer.input_size_per_partition,
+            bias=bias,
+        )

vllm/model_executor/kernels/linear/scaled_mm/BlockScaledMMLinearKernel.py

@@ -0,0 +1,209 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import ClassVar
+
+import torch
+from typing_extensions import Self
+
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+    process_fp8_weight_block_strategy,
+)
+from vllm.model_executor.utils import replace_parameter
+
+from ..base import (
+    FP8Params,
+    MMLinearKernel,
+)
+from .ScaledMMLinearKernel import FP8ScaledMMLinearLayerConfig
+
+
+@dataclass
+class FP8BlockParams(FP8Params):
+    weight_scale_inv: torch.Tensor | None
+    weight_scale: torch.Tensor | None
+
+    WEIGHT_SCALE_INV: ClassVar[str] = "weight_scale_inv"
+
+    @classmethod
+    def from_layer(cls, layer: torch.nn.Module) -> Self:
+        return cls(
+            weight=getattr(layer, cls.WEIGHT),
+            weight_scale_inv=getattr(layer, cls.WEIGHT_SCALE_INV, None),
+            weight_scale=getattr(layer, cls.WEIGHT_SCALE, None),
+            input_scale=getattr(layer, cls.INPUT_SCALE, None),
+            input_scale_ub=getattr(layer, cls.INPUT_SCALE_UB, None),
+        )
+
+
+class Fp8BlockScaledMMLinearKernel(
+    MMLinearKernel[FP8ScaledMMLinearLayerConfig, FP8BlockParams], ABC
+):
+    # Set to False in subclasses that accept BF16 input directly (e.g. FlashInfer)
+    # and therefore do not need the input quantization step in apply_weights.
+    apply_input_quant: ClassVar[bool] = True
+
+    def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
+        super().__init__(config)
+        act_scale_descriptor = config.activation_quant_key.scale
+        self.weight_group_shape = config.weight_quant_key.scale.group_shape
+        self.quant_fp8 = QuantFP8(
+            static=act_scale_descriptor.static,
+            group_shape=act_scale_descriptor.group_shape,
+            num_token_padding=self.get_output_padding(),
+            use_ue8m0=False,
+        )
+        self.use_triton = False
+
+    @classmethod
+    def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
+        act_quant_key = config.activation_quant_key
+        if act_quant_key.scale.static:
+            return (
+                False,
+                "Only dynamic per token group activation quantization is supported.",
+            )
+
+        return True, None
+
+    def _get_layer_params(self, layer: torch.nn.Module, **kwargs) -> FP8BlockParams:
+        return FP8BlockParams.from_layer(layer)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        params = self._get_layer_params(layer)
+        # Fp8LinearMethod registered weight scale
+        # buffer as weight_scale_inv unlike compressed tensors.
+        weight_scale = (
+            params.weight_scale
+            if params.weight_scale_inv is None
+            else params.weight_scale_inv
+        )
+        scale_attr_name = (
+            params.WEIGHT_SCALE
+            if params.weight_scale_inv is None
+            else params.WEIGHT_SCALE_INV
+        )
+        new_weight, new_weight_scale = process_fp8_weight_block_strategy(
+            params.weight,
+            weight_scale,
+        )
+
+        replace_parameter(layer, params.WEIGHT, new_weight.data)
+        replace_parameter(layer, scale_attr_name, new_weight_scale.data)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        out_dtype = self.config.out_dtype
+        params = self._get_layer_params(layer)
+        weight = params.weight
+        weight_scale = (
+            params.weight_scale
+            if params.weight_scale_inv is None
+            else params.weight_scale_inv
+        )
+        input_scale = params.input_scale
+        scale_up = params.input_scale_ub
+
+        # View input as 2D matrix for fp8 methods
+        input_2d = x.view(-1, x.shape[-1])
+        output_shape = [*x.shape[:-1], weight.shape[0]]
+
+        if self.apply_input_quant:
+            q_input, input_scale = self.quant_fp8(
+                input_2d, input_scale, scale_up, use_triton=self.use_triton
+            )
+        else:
+            q_input = input_2d
+            # Provide a concrete placeholder so apply_block_scaled_mm args are
+            # always Tensors. Subclasses with apply_input_quant=False must not
+            # use As in apply_block_scaled_mm.
+            input_scale = (
+                input_scale if input_scale is not None else input_2d.new_empty(1)
+            )
+
+        output = self.apply_block_scaled_mm(
+            A=q_input,
+            B=weight,
+            As=input_scale,
+            Bs=weight_scale,
+        )
+
+        if bias is not None:
+            output = output + bias
+        return output.to(dtype=out_dtype).view(*output_shape)
+
+    @abstractmethod
+    def apply_block_scaled_mm(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class Fp8BlockScaledDynamicMMLinearKernel(Fp8BlockScaledMMLinearKernel, ABC):
+    """Dynamic FP8 block-scaled kernel that dispatches at runtime.
+
+    Extends Fp8BlockScaledMMLinearKernel to inherit apply_weights and overrides
+    apply_block_scaled_mm to dispatch between two sub-kernels using torch.cond.
+
+    Subclasses must define:
+        base_type:     The primary kernel class.
+        fallback_type: The fallback kernel class.
+    """
+
+    base_type: ClassVar[type[Fp8BlockScaledMMLinearKernel]]
+    fallback_type: ClassVar[type[Fp8BlockScaledMMLinearKernel]]
+
+    def __init__(self, config: "FP8ScaledMMLinearLayerConfig") -> None:
+        super().__init__(config)
+        self.base = self.base_type(config)
+        self.fallback = self.fallback_type(config)
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        is_base_supported, reason_1 = cls.base_type.is_supported(compute_capability)
+        is_fallback_supported, reason_2 = cls.fallback_type.is_supported(
+            compute_capability
+        )
+        if is_base_supported and is_fallback_supported:
+            return True, None
+        if not is_base_supported and not is_fallback_supported:
+            return (
+                False,
+                f"base is not supported due to {reason_1}; "
+                f"fallback is not supported due to {reason_2}",
+            )
+        if not is_base_supported:
+            return False, f"base is not supported due to {reason_1}"
+        return False, f"fallback is not supported due to {reason_2}"
+
+    @classmethod
+    def can_implement(
+        cls, config: "FP8ScaledMMLinearLayerConfig"
+    ) -> tuple[bool, str | None]:
+        can_implement_base, reason_1 = cls.base_type.can_implement(config)
+        can_implement_fallback, reason_2 = cls.fallback_type.can_implement(config)
+        if can_implement_base and can_implement_fallback:
+            return True, None
+        if not can_implement_base and not can_implement_fallback:
+            return (
+                False,
+                f"base cannot implement due to {reason_1}; "
+                f"fallback cannot implement due to {reason_2}",
+            )
+        if not can_implement_base:
+            return False, f"base cannot implement due to {reason_1}"
+        return False, f"fallback cannot implement due to {reason_2}"

vllm/model_executor/kernels/linear/scaled_mm/ScaledMMLinearKernel.py

@@ -0,0 +1,201 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from abc import ABC, abstractmethod
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+import torch
+
+from vllm.model_executor.layers.fusion.quant_activation import (
+    QuantizedActivation,
+    as_quantized_activation,
+)
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+)
+from vllm.platforms import current_platform
+
+from ..base import MMLinearLayerConfig
+
+
+@dataclass
+class Int8ScaledMMLinearLayerConfig(MMLinearLayerConfig):
+    # TODO: Change to QuantKey like FP8ScaledMMLinearLayerConfig
+    is_static_input_scheme: bool
+    is_channelwise: bool
+    input_symmetric: bool
+
+
+@dataclass
+class FP8ScaledMMLinearLayerConfig(MMLinearLayerConfig):
+    weight_quant_key: QuantKey
+    activation_quant_key: QuantKey
+    weight_shape: tuple[int, int]
+    input_dtype: torch.dtype
+    out_dtype: torch.dtype
+
+
+_FP8ParamsT = tuple[
+    torch.Tensor,  # weight
+    torch.Tensor,  # weight_scale
+    torch.Tensor | None,  # input_scale,
+    torch.Tensor | None,  # input_scale_ub,
+]
+_Int8ParamsT = tuple[
+    torch.Tensor,  # weight
+    torch.Tensor,  # weight_scale
+    torch.Tensor | None,  # input_scale,
+    torch.Tensor | None,  # input_zp
+    torch.Tensor | None,  # azp_adj
+]
+
+_ParamsT = TypeVar("_ParamsT", _Int8ParamsT, _FP8ParamsT)
+_ConfigT = TypeVar("_ConfigT", bound=MMLinearLayerConfig)
+
+
+class ScaledMMLinearKernel(Generic[_ConfigT, _ParamsT], ABC):
+    @classmethod
+    @abstractmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def can_implement(cls, c: _ConfigT) -> tuple[bool, str | None]:
+        raise NotImplementedError
+
+    def __init__(self, c: _ConfigT, layer_param_names: Sequence[str]) -> None:
+        assert self.can_implement(c)[0]
+        assert self.is_supported()[0]
+        self.config = c
+        self.layer_param_names = layer_param_names
+
+    def input_quant_key(self) -> QuantKey | None:
+        """The activation quant key this kernel can consume pre-quantized.
+
+        Manual fusion uses this to decide whether to hoist activation
+        quantization out of apply_weights into an upstream fused kernel.
+        Return None when the kernel needs in-kernel quantization (custom
+        padding or swizzling, dynamic scales, etc.). Kernels that return a
+        key must consume the activation via as_quantized_activation.
+        """
+        return None
+
+    @abstractmethod
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        raise NotImplementedError
+
+    @abstractmethod
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    # return a covariant type in the subclass
+    @abstractmethod
+    def _get_layer_params(self, layer) -> _ParamsT:
+        raise NotImplementedError
+
+
+class FP8ScaledMMLinearKernel(
+    ScaledMMLinearKernel[FP8ScaledMMLinearLayerConfig, _FP8ParamsT], ABC
+):
+    def __init__(
+        self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
+    ) -> None:
+        act_scale_descriptor = c.activation_quant_key.scale
+        self.quant_fp8 = QuantFP8(
+            static=act_scale_descriptor.static,
+            group_shape=act_scale_descriptor.group_shape,
+            num_token_padding=self.get_output_padding(),
+        )
+        self.fp8_dtype = current_platform.fp8_dtype()
+        super().__init__(c, layer_param_names)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        pass
+
+    def _get_layer_params(self, layer) -> _FP8ParamsT:
+        w, w_s, x_s, x_s_ub = self.layer_param_names
+        return (
+            getattr(layer, w),
+            getattr(layer, w_s),
+            getattr(layer, x_s, None),
+            getattr(layer, x_s_ub, None),
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor | QuantizedActivation,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        fp8_dtype = self.fp8_dtype
+        maybe_out_dtype = self.config.out_dtype
+        w, w_s, x_s, x_s_ub = self._get_layer_params(layer)
+
+        qa = as_quantized_activation(x, self.input_quant_key())
+        if qa is not None:
+            x_data, x_s = qa.data, qa.scale
+            orig_shape, orig_dtype = qa.orig_shape, qa.orig_dtype
+            assert x_data.dtype == fp8_dtype
+        else:
+            assert isinstance(x, torch.Tensor)
+            x_data = x
+            orig_shape, orig_dtype = x.shape, x.dtype
+
+        x_2d = x_data.view(-1, x_data.shape[-1])
+        output_shape = [*orig_shape[:-1], w.shape[1]]
+        out_dtype = orig_dtype if maybe_out_dtype is None else maybe_out_dtype
+
+        x_2d_q = x_2d
+        if qa is None:
+            x_2d_q, x_s = self.quant_fp8(x_2d, x_s, x_s_ub)
+        return self.apply_scaled_mm(
+            A=x_2d_q,
+            B=w,
+            out_dtype=out_dtype,
+            As=x_s,
+            Bs=w_s,
+            bias=bias,
+            output_shape=output_shape,
+        )
+
+    @abstractmethod
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+    def get_output_padding(self) -> int | None:
+        return None
+
+
+class Int8ScaledMMLinearKernel(
+    ScaledMMLinearKernel[Int8ScaledMMLinearLayerConfig, _Int8ParamsT], ABC
+):
+    def _get_layer_params(self, layer) -> _Int8ParamsT:
+        w_q, w_s, i_s, i_zp, azp_adj = self.layer_param_names
+        return (
+            getattr(layer, w_q),
+            getattr(layer, w_s),
+            getattr(layer, i_s, None),
+            getattr(layer, i_zp, None),
+            getattr(layer, azp_adj, None),
+        )

vllm/model_executor/kernels/linear/scaled_mm/__init__.py

@@ -0,0 +1,71 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
+    AiterInt8ScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.BlockScaledMMLinearKernel import (
+    Fp8BlockScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
+    CPUFp8BlockScaledMMKernel,
+    CPUInt8ScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
+    CutlassFP8ScaledMMLinearKernel,
+    CutlassInt8ScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.flashinfer import (
+    FlashInferFP8ScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.marlin import (
+    MarlinFP8ScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.pytorch import (
+    ChannelWiseTorchFP8ScaledMMLinearKernel,
+    PerTensorTorchFP8ScaledMMLinearKernel,
+    RowWiseTorchFP8ScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.rocm import (
+    ROCmFP8ScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
+    FP8ScaledMMLinearKernel,
+    FP8ScaledMMLinearLayerConfig,
+    Int8ScaledMMLinearKernel,
+    Int8ScaledMMLinearLayerConfig,
+    ScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.triton import (
+    TritonInt8ScaledMMLinearKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.xpu import (
+    XPUFp8BlockScaledMMKernel,
+)
+from vllm.model_executor.kernels.linear.scaled_mm.zentorch import (
+    ZentorchInt8ScaledMMLinearKernel,
+)
+
+__all__ = [
+    "FP8ScaledMMLinearKernel",
+    "FP8ScaledMMLinearLayerConfig",
+    "Int8ScaledMMLinearKernel",
+    "Int8ScaledMMLinearLayerConfig",
+    "ScaledMMLinearKernel",
+    "ScaledMMLinearLayerConfig",
+    "AiterInt8ScaledMMLinearKernel",
+    "CPUInt8ScaledMMLinearKernel",
+    "CutlassFP8ScaledMMLinearKernel",
+    "CutlassInt8ScaledMMLinearKernel",
+    "FlashInferFP8ScaledMMLinearKernel",
+    "MarlinFP8ScaledMMLinearKernel",
+    "ChannelWiseTorchFP8ScaledMMLinearKernel",
+    "PerTensorTorchFP8ScaledMMLinearKernel",
+    "RowWiseTorchFP8ScaledMMLinearKernel",
+    "ROCmFP8ScaledMMLinearKernel",
+    "TritonInt8ScaledMMLinearKernel",
+    "ZentorchInt8ScaledMMLinearKernel",
+    "Fp8BlockScaledMMLinearKernel",
+    "CPUFp8BlockScaledMMKernel",
+    "XPUFp8BlockScaledMMKernel",
+]

vllm/model_executor/kernels/linear/scaled_mm/aiter.py

@@ -0,0 +1,431 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm._aiter_ops import (
+    rocm_aiter_ops,
+)
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+)
+from vllm.model_executor.utils import replace_parameter
+from vllm.platforms import current_platform
+
+from .BlockScaledMMLinearKernel import (
+    Fp8BlockScaledMMLinearKernel,
+)
+from .cutlass import CutlassInt8ScaledMMLinearKernel
+from .ScaledMMLinearKernel import (
+    FP8ScaledMMLinearKernel,
+    FP8ScaledMMLinearLayerConfig,
+    Int8ScaledMMLinearLayerConfig,
+)
+
+logger = init_logger(__name__)
+
+
+class AiterInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_rocm():
+            return False, "Requires ROCm."
+
+        if compute_capability is not None and compute_capability < 90:
+            return False, "requires compute capability 90 and above."
+
+        try:
+            import aiter  # noqa: F401 # deliberately attempt to import aiter
+        except Exception:
+            return False, "requires `aiter` to be installed."
+
+        if not rocm_aiter_ops.is_linear_enabled():
+            return (
+                False,
+                "requires setting `VLLM_ROCM_USE_AITER=1` "
+                "and `VLLM_ROCM_USE_AITER_LINEAR=1`. "
+                "`VLLM_ROCM_USE_AITER_LINEAR` default is True.",
+            )
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        if not c.input_symmetric:
+            return False, "supports symmetric quantization only."
+        return True, None
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """
+        `AiterInt8ScaledMMLinearKernel` implements a fused version of
+            `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
+        where scale_a * a and scale_b * b are implemented using numpy-style
+        broadcasting.
+        Currently only support per-tensor-per-tensor GEMM
+        and per-token-per-channel GEMM through AITER
+        w8a8 scaled gemm. `AiterInt8ScaledMMLinearKernel` also does not support
+        ATIER block scaled GEMM and mix-precision GEMM.
+        """
+        w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
+
+        # ops.scaled_int8_quant supports both dynamic and static quant:
+        # * dynamic, i_s is None and x_s computed from x.
+        # * static, i_s is scalar and x_s is i_s.
+        symmetric = azp_adj is None
+        assert symmetric, (
+            "AiterInt8ScaledMMLinearKernel only supports symmetric quantization."
+        )
+        x_q, x_s, x_zp = ops.scaled_int8_quant(x, i_s, i_zp, symmetric=symmetric)
+
+        assert x_zp is None, (
+            "AiterInt8ScaledMMLinearKernel only supports symmetric quantization."
+        )
+        out_dtype = x.dtype
+
+        assert w_q.shape[0] % 16 == 0 and w_q.shape[1] % 16 == 0
+        assert out_dtype is torch.bfloat16 or out_dtype is torch.float16
+        assert bias is None or bias.shape[0] == w_q.shape[1] and bias.dtype == out_dtype
+
+        m = x_q.shape[0]  # a
+        n = w_q.shape[1]  # b
+
+        per_tensor_scale_a = x_s.numel() == 1
+        per_tensor_scale_b = w_s.numel() == 1
+        per_token_scale_a = x_s.numel() == m
+        per_channel_scale_b = w_s.numel() == n
+
+        # @TODO:
+        # Maybe broadcast the per-tensor-scale into per-channel-scale
+        # if one of the scale is a per-channel-scale.
+        # For now, it only supports:
+        # - per-tensor-per-tensor a8w8 scaled GEMM, and
+        # - per-token-per-channel a8w8 scaled GEMM
+        assert (per_tensor_scale_a and per_tensor_scale_b) or (
+            per_token_scale_a and per_channel_scale_b
+        ), (
+            "Currently only support per-tensor-per-tensor GEMM "
+            " and per-token-per-channel GEMM through AITER"
+            " w8a8 scaled gemm. `AiterInt8ScaledMMLinearKernel` "
+            "does not support AITER block scaled GEMM."
+        )
+
+        # gemm_a8w8_CK(a, b, scale_a, scale_b, bias) expects
+        # a to be [M, K]
+        # b to be [N, K]
+        # CutlassInt8ScaledMMLinearKernel prepare weight `w_q` in [K, N] format
+        return rocm_aiter_ops.w8a8_gemm(x_q, w_q.t(), x_s, w_s, bias, out_dtype)
+
+
+class AiterPreshuffledPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_rocm():
+            return False, "requires ROCm."
+        if not rocm_aiter_ops.is_linear_fp8_enabled():
+            return (
+                False,
+                "requires setting `VLLM_ROCM_USE_AITER=1` "
+                "and `VLLM_ROCM_USE_AITER_LINEAR=1`. "
+                "`VLLM_ROCM_USE_AITER_LINEAR` default is True.",
+            )
+        try:
+            import aiter  # noqa: F401
+        except Exception:
+            return False, "requires aiter library to be installed."
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        is_ptpc = (
+            c.activation_quant_key.scale.group_shape.is_per_token()
+            and c.weight_quant_key.scale.group_shape.is_per_channel()
+        )
+        if c.weight_shape is None:
+            return False, "weight_shape is required for Aiter kernels"
+        N, K = c.weight_shape
+        fp8_dtype = current_platform.fp8_dtype()
+
+        if c.out_dtype is not torch.bfloat16:
+            return False, "requires bfloat16 output dtype."
+
+        if not is_ptpc:
+            return (
+                False,
+                "requires per token activation scales and per channel weight scales.",
+            )
+
+        if not (N % 16 == 0 and K % 16 == 0):
+            return (
+                False,
+                f"requires N and K dimensions divisible by 16, received "
+                f"N={N} and K={K}.",
+            )
+
+        # Aiter's shuffled per-token Gemm performs better than torch only when its
+        # tuned.
+        if not rocm_aiter_ops.is_shuffled_per_token_w8a8_gemm_tuned(N, K, fp8_dtype):
+            return (
+                False,
+                f"requires a tuned configuration for N: {N} and K: {K} "
+                f"and fp8 dtype {fp8_dtype}.",
+            )
+
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        w_name, *_ = self.layer_param_names
+        w, *_ = self._get_layer_params(layer)
+
+        replace_parameter(
+            layer,
+            w_name,
+            torch.nn.Parameter(
+                rocm_aiter_ops.shuffle_weight(w.t().contiguous()).data,
+                requires_grad=False,
+            ),
+        )
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        return rocm_aiter_ops.preshuffled_per_token_w8a8_gemm(
+            A, B, As, Bs, bias, out_dtype
+        )
+
+
+class AiterHipbMMPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_rocm():
+            return False, "requires ROCm."
+
+        if not rocm_aiter_ops.is_linear_hipbmm_enabled():
+            return (
+                False,
+                "requires setting `VLLM_ROCM_USE_AITER=1`, "
+                "`VLLM_ROCM_USE_AITER_LINEAR=1`, "
+                "and `VLLM_ROCM_USE_AITER_LINEAR_HIPBMM=1`.",
+            )
+        try:
+            import aiter  # noqa: F401
+        except Exception:
+            return False, "requires aiter library to be installed."
+
+        if not hasattr(aiter, "hipb_mm"):
+            return False, "requires aiter hipb_mm support."
+
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        is_ptpc = (
+            c.activation_quant_key.scale.group_shape.is_per_token()
+            and c.weight_quant_key.scale.group_shape.is_per_channel()
+        )
+        if c.weight_shape is None:
+            return False, "weight_shape is required for Aiter kernels"
+        N, K = c.weight_shape
+
+        if c.out_dtype is not torch.bfloat16:
+            return False, "requires bfloat16 output dtype."
+
+        if not is_ptpc:
+            return (
+                False,
+                "requires per token activation scales and per channel weight scales.",
+            )
+
+        if not (N >= 16 and N % 16 == 0 and K % 16 == 0):
+            return (
+                False,
+                "requires N >= 16 and both N and K divisible by 16, "
+                f"received N={N} and K={K}.",
+            )
+
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        w_name, w_s_name, *_ = self.layer_param_names
+        w, w_s, *_ = self._get_layer_params(layer)
+
+        # Pre-apply the transposes that used to live in
+        # _rocm_aiter_hipb_mm_fp8_impl so the kernel can consume B/Bs directly.
+        # The `.t()` on the shuffled weight is kept as a non-contiguous view —
+        # materializing it with `.contiguous()` would re-arrange the bytes and
+        # break the `bpreshuffle` layout.
+        shuffled_w = rocm_aiter_ops.shuffle_weight(w.t().contiguous())
+        replace_parameter(
+            layer,
+            w_name,
+            torch.nn.Parameter(shuffled_w.t(), requires_grad=False),
+        )
+
+        if w_s.ndim > 1:
+            replace_parameter(
+                layer,
+                w_s_name,
+                torch.nn.Parameter(w_s.t().contiguous(), requires_grad=False),
+            )
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        output_shape[-1] = B.shape[1]
+        return rocm_aiter_ops.hipb_mm_fp8(A, B, As, Bs, bias, out_dtype).view(
+            *output_shape
+        )
+
+
+class AiterPerTokenFp8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        return AiterPreshuffledPerTokenFp8ScaledMMLinearKernel.is_supported(
+            compute_capability
+        )
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        is_ptpc = (
+            c.activation_quant_key.scale.group_shape.is_per_token()
+            and c.weight_quant_key.scale.group_shape.is_per_channel()
+        )
+        if c.weight_shape is None:
+            return False, "weight_shape is required for Aiter kernels"
+        N, K = c.weight_shape
+        fp8_dtype = current_platform.fp8_dtype()
+
+        if not is_ptpc:
+            return (
+                False,
+                "requires per token activation scales and per channel weight scales.",
+            )
+
+        # Aiter's per-token Gemm performs better than torch oonly when its
+        # tuned.
+        if not rocm_aiter_ops.is_per_token_w8a8_gemm_tuned(N, K, fp8_dtype):
+            return (
+                False,
+                f"requires a tuned configuration for N: {N} and K: {K} "
+                f"and fp8 dtype {fp8_dtype}.",
+            )
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        w_name, *_ = self.layer_param_names
+        w, *_ = self._get_layer_params(layer)
+
+        replace_parameter(
+            layer,
+            w_name,
+            torch.nn.Parameter(w.t(), requires_grad=False),
+        )
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        return rocm_aiter_ops.w8a8_gemm(A, B, As, Bs, bias, out_dtype)
+
+
+class AiterFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
+    def __init__(self, config: FP8ScaledMMLinearLayerConfig):
+        super().__init__(config)
+        n, k = config.weight_shape
+
+        self.use_triton = (
+            not current_platform.is_fp8_fnuz()
+            and rocm_aiter_ops.is_triton_gemm_w8a8_tuned(n, k)
+        )
+
+    @classmethod
+    def is_supported(cls, compute_capability=None):
+        return (
+            rocm_aiter_ops.is_linear_enabled(),
+            "Only supported on ROCm platform \
+                with aiter package installed.",
+        )
+
+    @classmethod
+    def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
+        can_implement_base, reason = super().can_implement(config)
+        if not can_implement_base:
+            return can_implement_base, reason
+
+        act_quant_desc = config.activation_quant_key.scale
+        if act_quant_desc.group_shape != GroupShape(1, 128):
+            return (
+                False,
+                "Supports only dynamic per token group activation "
+                "quantization with group_shape=(1,128).",
+            )
+        return True, None
+
+    def apply_block_scaled_mm(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+    ) -> torch.Tensor:
+        if As.dtype != Bs.dtype:
+            from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+                _upcast_e8m0_to_fp32,
+            )
+
+            if As.dtype == torch.float8_e8m0fnu:
+                As = _upcast_e8m0_to_fp32(As).contiguous()
+            else:
+                As = As.to(torch.float32)
+
+            if Bs.dtype == torch.float8_e8m0fnu:
+                Bs = _upcast_e8m0_to_fp32(Bs).contiguous()
+            else:
+                Bs = Bs.to(torch.float32)
+
+        out_dtype = self.config.out_dtype
+        if self.use_triton:
+            gemm_a8w8_blockscale_op = rocm_aiter_ops.triton_gemm_a8w8_blockscale
+        else:
+            gemm_a8w8_blockscale_op = rocm_aiter_ops.gemm_a8w8_blockscale
+
+        return gemm_a8w8_blockscale_op(
+            A, B, As, Bs, list(self.weight_group_shape), output_dtype=out_dtype
+        )

vllm/model_executor/kernels/linear/scaled_mm/cpu.py

@@ -0,0 +1,327 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm import envs
+from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    convert_to_channelwise,
+)
+from vllm.model_executor.layers.utils import check_cpu_sgl_kernel
+from vllm.platforms import current_platform
+from vllm.platforms.interface import CpuArchEnum
+
+from .BlockScaledMMLinearKernel import (
+    Fp8BlockScaledMMLinearKernel,
+    FP8ScaledMMLinearLayerConfig,
+)
+from .ScaledMMLinearKernel import (
+    Int8ScaledMMLinearKernel,
+    Int8ScaledMMLinearLayerConfig,
+)
+
+
+class CPUInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cpu():
+            return False, "requires CPU."
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        w_q_name, _, _, _, _ = self.layer_param_names
+        weight = getattr(layer, w_q_name)
+        dtype = weight.dtype
+        N, K = weight.size()
+        if (
+            current_platform.get_cpu_architecture() == CpuArchEnum.X86
+            and envs.VLLM_CPU_SGL_KERNEL
+            and self.config.input_symmetric
+            and check_cpu_sgl_kernel(N, K, dtype)
+        ):
+            self.linear_method = self._apply_weights_sgl
+            self.process_weights_for_sgl(layer)
+        else:
+            self.linear_method = self._apply_weights_onednn
+            self.process_weights_for_onednn(layer)
+
+    def process_weights_for_onednn(self, layer: torch.nn.Module) -> None:
+        # WEIGHT
+        # Transpose to [K, N] for convenience
+        w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names
+        weight = getattr(layer, w_q_name)
+        replace_parameter(
+            layer,
+            w_q_name,
+            torch.nn.Parameter(weight.t().data, requires_grad=False),
+        )
+
+        # WEIGHT SCALE
+        # oneDNN kernels support only per-tensor and per-channel.
+        # If we have a fused module (QKV, MLP) with per tensor scales (thus N
+        # scales being passed to the kernel), convert to the per-channel case.
+        is_fused_module = len(layer.logical_widths) > 1
+        weight_scale = getattr(layer, w_s_name)
+        if is_fused_module and not self.config.is_channelwise:
+            weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
+        replace_parameter(
+            layer,
+            w_s_name,
+            torch.nn.Parameter(weight_scale.data, requires_grad=False),
+        )
+
+        # INPUT SCALE
+        if self.config.is_static_input_scheme:
+            input_scale = getattr(layer, i_s_name)
+
+            if self.config.input_symmetric:
+                replace_parameter(
+                    layer,
+                    i_s_name,
+                    torch.nn.Parameter(input_scale.max(), requires_grad=False),
+                )
+            else:
+                input_zero_point = getattr(layer, i_zp_name)
+
+                # reconstruct the ranges
+                int8_traits = torch.iinfo(torch.int8)
+                azps = input_zero_point.to(dtype=torch.int32)
+                range_max = (input_scale * (int8_traits.max - azps)).max()
+                range_min = (input_scale * (int8_traits.min - azps)).min()
+
+                scale = (range_max - range_min) / (int8_traits.max - int8_traits.min)
+                replace_parameter(
+                    layer, i_s_name, torch.nn.Parameter(scale, requires_grad=False)
+                )
+
+                azp = (
+                    (int8_traits.min - range_min / scale).round().to(dtype=torch.int32)
+                )
+                replace_parameter(
+                    layer, i_zp_name, torch.nn.Parameter(azp, requires_grad=False)
+                )
+
+        # Different from cutlass, oneDNN kernels only need the AZP adjustment
+        # term for dynamic quantization. And s_b should be folded into the
+        # term. Such as:
+        # s_a * s_b * [(A - zp_a)B] + bias =
+        # s_a * (s_b * AB) - s_a * s_b * zp_a * B + bias =
+        # s_a * GEMM_output - s_a * zp_a * adj + bias
+        if not (self.config.input_symmetric and self.config.is_static_input_scheme):
+            weight = getattr(layer, w_q_name)
+            weight_scale = getattr(layer, w_s_name)
+            azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.float32)
+            azp_adj = azp_adj * weight_scale.squeeze()
+            setattr(
+                layer,
+                azp_adj_name,
+                torch.nn.Parameter(azp_adj, requires_grad=False),
+            )
+
+        weight = getattr(layer, w_q_name)
+        self.dnnl_handler = ops.create_onednn_scaled_mm(
+            weight,
+            getattr(layer, w_s_name),
+            torch.get_default_dtype(),
+            getattr(layer, i_s_name) is None,
+            not self.config.input_symmetric,
+            32,
+        )
+        # weight is prepacked and maintained by the dnnl_handler,
+        # release the original weight
+        setattr(layer, w_q_name, None)
+        del weight
+
+    def process_weights_for_sgl(self, layer: torch.nn.Module) -> None:
+        w_q_name, w_s_name, _, _, _ = self.layer_param_names
+        # WEIGHT
+        weight = getattr(layer, w_q_name)
+        packed_weight = torch.ops._C.convert_weight_packed(weight)
+        replace_parameter(
+            layer, w_q_name, torch.nn.Parameter(packed_weight, requires_grad=False)
+        )
+
+        if layer.bias is not None:
+            bias = layer.bias
+            layer.register_parameter(
+                "bias_fp32", torch.nn.Parameter(bias.float().data, requires_grad=False)
+            )
+
+        # WEIGHT SCALE
+        # CPU SGL kernels only support per-channel.
+        # For per-tensor quant, convert to the per-channel case.
+        weight_scale = getattr(layer, w_s_name)
+        if not self.config.is_channelwise:
+            weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
+        replace_parameter(
+            layer,
+            w_s_name,
+            torch.nn.Parameter(weight_scale.data, requires_grad=False),
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.linear_method(
+            layer,
+            x,
+            bias,
+        )
+
+    def _apply_weights_onednn(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        x_shape = x.shape
+        x = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x
+        w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
+
+        # ops.scaled_int8_quant supports both dynamic and static quant:
+        # * dynamic, i_s is None and x_s computed from x.
+        # * static, i_s is scalar and x_s is i_s.
+        x_q, x_s, x_zp = ops.onednn_scaled_int8_quant(
+            x, i_s, i_zp, self.config.input_symmetric
+        )
+
+        m = x.size(0)
+        n = self.dnnl_handler.n
+        out = torch.empty((m, n), dtype=x.dtype)
+        ops.onednn_scaled_mm(self.dnnl_handler, x_q, out, x_s, x_zp, azp_adj, bias)
+        out = out.reshape(x_shape[:-1] + (n,)) if len(x_shape) > 2 else out
+        return out
+
+    def _apply_weights_sgl(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        w_q, w_s, _, _, _ = self._get_layer_params(layer)
+        return torch.ops._C.int8_scaled_mm_with_quant(
+            x,
+            w_q,
+            w_s,
+            layer.bias_fp32 if bias is not None else None,
+            x.dtype,
+            True,
+        )
+
+
+class CPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
+    """FP8 W8A16 block-quantized GEMM via AMX BRGEMM on CPU."""
+
+    # Input stays BF16 — no FP8 activation quantization.
+    apply_input_quant = False
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cpu():
+            return False, "requires CPU platform."
+        if not torch.cpu._is_amx_tile_supported():
+            return False, "requires AMX tile support (Sapphire Rapids or newer)."
+        if not ops._supports_cpu_fp8_w8a16:
+            return False, "fp8_scaled_mm_cpu op not available."
+        return True, None
+
+    @classmethod
+    def can_implement(
+        cls, config: FP8ScaledMMLinearLayerConfig
+    ) -> tuple[bool, str | None]:
+        # Validate weight block shape
+        weight_gs = config.weight_quant_key.scale.group_shape
+        if weight_gs.col <= 0 or weight_gs.col != 128:
+            return False, (
+                "CPU FP8 kernel requires K-dimension block size of 128, "
+                f"got {weight_gs.col}."
+            )
+        if weight_gs.row <= 0 or weight_gs.row % 32 != 0:
+            return False, (
+                "CPU FP8 kernel requires N-dimension block size to be "
+                f"a positive multiple of 32, got {weight_gs.row}."
+            )
+        if config.out_dtype not in (torch.bfloat16, torch.float32):
+            return False, "Only bfloat16/float32 output dtype supported."
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # Skip the base class process (FP8 padding / fnuz normalization)
+        # which is GPU-oriented.  Instead, VNNI-prepack weights for AMX.
+        params = self._get_layer_params(layer)
+        packed_weight = torch.ops._C.convert_weight_packed(params.weight)
+        replace_parameter(
+            layer,
+            params.WEIGHT,
+            torch.nn.Parameter(packed_weight, requires_grad=False),
+        )
+
+        # Re-wrap scale as a plain Parameter so the kernel can read it
+        # without weight-loader metadata interfering.
+        scale_attr = (
+            params.WEIGHT_SCALE_INV
+            if params.weight_scale_inv is not None
+            else params.WEIGHT_SCALE
+        )
+        weight_scale = (
+            params.weight_scale_inv
+            if params.weight_scale_inv is not None
+            else params.weight_scale
+        )
+        assert weight_scale is not None
+        replace_parameter(
+            layer,
+            scale_attr,
+            torch.nn.Parameter(weight_scale.data, requires_grad=False),
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        params = self._get_layer_params(layer)
+        weight_scale = (
+            params.weight_scale_inv
+            if params.weight_scale_inv is not None
+            else params.weight_scale
+        )
+
+        x_2d = x.reshape(-1, x.shape[-1]) if x.dim() > 2 else x
+        out = torch.ops._C.fp8_scaled_mm_cpu(
+            x_2d,
+            params.weight,
+            weight_scale,
+            list(self.weight_group_shape),
+            bias,
+            x.dtype,
+            True,  # is_vnni (weight already prepacked)
+        )
+        return out.reshape(x.shape[:-1] + (out.size(-1),)) if x.dim() > 2 else out
+
+    def apply_block_scaled_mm(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+    ) -> torch.Tensor:
+        raise NotImplementedError(
+            "CPUFp8BlockScaledMMKernel overrides apply_weights directly."
+        )

vllm/model_executor/kernels/linear/scaled_mm/cutlass.py

@@ -0,0 +1,343 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+from collections.abc import Sequence
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
+from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    QuantKey,
+    kFp8StaticTensorSym,
+)
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    CUTLASS_BLOCK_FP8_SUPPORTED,
+    convert_to_channelwise,
+)
+from vllm.model_executor.utils import set_weight_attrs
+from vllm.platforms import current_platform
+
+from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
+from .ScaledMMLinearKernel import (
+    FP8ScaledMMLinearKernel,
+    FP8ScaledMMLinearLayerConfig,
+    Int8ScaledMMLinearKernel,
+    Int8ScaledMMLinearLayerConfig,
+)
+
+
+class CutlassInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "requires CUDA."
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names
+        config = self.config
+        # WEIGHT
+        # Cutlass kernels need transposed weight.
+        weight = getattr(layer, w_q_name)
+        replace_parameter(
+            layer,
+            w_q_name,
+            torch.nn.Parameter(weight.t().data, requires_grad=False),
+        )
+
+        # WEIGHT SCALE
+        # Cutlass kernels support only per-tensor and per-channel.
+        # If we have a fused module (QKV, MLP) with per tensor scales (thus N
+        # scales being passed to the kernel), convert to the per-channel case.
+        is_fused_module = len(layer.logical_widths) > 1
+        weight_scale = getattr(layer, w_s_name)
+        if is_fused_module and not config.is_channelwise:
+            weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
+        replace_parameter(
+            layer,
+            w_s_name,
+            torch.nn.Parameter(weight_scale.data, requires_grad=False),
+        )
+
+        # INPUT SCALE
+        if config.is_static_input_scheme:
+            input_scale = getattr(layer, i_s_name)
+
+            if config.input_symmetric:
+                replace_parameter(
+                    layer,
+                    i_s_name,
+                    torch.nn.Parameter(input_scale.max(), requires_grad=False),
+                )
+                setattr(layer, i_zp_name, None)
+            else:
+                input_zero_point = getattr(layer, i_zp_name)
+
+                # reconstruct the ranges
+                int8_traits = torch.iinfo(torch.int8)
+                azps = input_zero_point.to(dtype=torch.int32)
+                range_max = (input_scale * (int8_traits.max - azps)).max()
+                range_min = (input_scale * (int8_traits.min - azps)).min()
+
+                scale = (range_max - range_min) / (int8_traits.max - int8_traits.min)
+                replace_parameter(
+                    layer, i_s_name, torch.nn.Parameter(scale, requires_grad=False)
+                )
+
+                # AZP loaded as int8 but used as int32
+                azp = (int8_traits.min - range_min / scale).to(dtype=torch.int32)
+                replace_parameter(
+                    layer, i_zp_name, torch.nn.Parameter(azp, requires_grad=False)
+                )
+
+        # azp_adj is the AZP adjustment term, used to account for weights.
+        # It does not depend on scales or azp, so it is the same for
+        # static and dynamic quantization.
+        # For more details, see csrc/quantization/w8a8/cutlass/Epilogues.md
+        # https://github.com/vllm-project/vllm/blob/main/csrc/quantization/w8a8/cutlass/Epilogues.md
+        if not config.input_symmetric:
+            weight = getattr(layer, w_q_name)
+            azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.int32)
+            if config.is_static_input_scheme:
+                # cutlass_w8a8 requires azp to be folded into azp_adj
+                # in the per-tensor case
+                azp_adj = getattr(layer, i_zp_name) * azp_adj
+            setattr(
+                layer,
+                azp_adj_name,
+                torch.nn.Parameter(azp_adj, requires_grad=False),
+            )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
+
+        # ops.scaled_int8_quant supports both dynamic and static quant:
+        # * dynamic, i_s is None and x_s computed from x.
+        # * static, i_s is scalar and x_s is i_s.
+        symmetric = azp_adj is None
+        x_q, x_s, x_zp = ops.scaled_int8_quant(
+            x.contiguous(), i_s, i_zp, symmetric=symmetric
+        )
+
+        if x_zp is not None:
+            # Currently, static is always per-tensor and dynamic is per-token
+            static = i_zp is not None
+            azp = None if static else x_zp
+            return ops.cutlass_scaled_mm_azp(
+                x_q,
+                w_q,
+                scale_a=x_s,
+                scale_b=w_s,
+                out_dtype=x.dtype,
+                azp_adj=azp_adj,
+                azp=azp,
+                bias=bias,
+            )
+        return ops.cutlass_scaled_mm(
+            x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias
+        )
+
+
+class CutlassFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    def __init__(
+        self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
+    ) -> None:
+        self.logical_output_size: int | None = None
+        super().__init__(c, layer_param_names)
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "requires CUDA."
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def input_quant_key(self) -> QuantKey | None:
+        """Only static per-tensor activation quantization is supported for external
+        quantization."""
+        if self.config.activation_quant_key == kFp8StaticTensorSym:
+            return kFp8StaticTensorSym
+        return None
+
+    @staticmethod
+    def _pad_to_alignment(
+        x: torch.Tensor, dim: int, alignment: int, value: float = 0.0
+    ) -> torch.Tensor:
+        """Pad tensor ``x`` along ``dim`` to the next multiple of
+        ``alignment``."""
+        remainder = x.shape[dim] % alignment
+        if remainder == 0:
+            return x
+        pad_size = alignment - remainder
+        pad_spec = [0] * (2 * x.dim())
+        pad_spec[-(2 * dim + 1)] = pad_size
+        return torch.nn.functional.pad(x, pad_spec, value=value)
+
+    @staticmethod
+    def padded_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None:
+        if loaded_weight.shape != param.shape:
+            slices = tuple(slice(0, s) for s in loaded_weight.shape)
+            param.data[slices].copy_(loaded_weight)
+        else:
+            param.data.copy_(loaded_weight.view(param.shape))
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight_name, weight_scale_name, _, _ = self.layer_param_names
+        weight = getattr(layer, weight_name)
+
+        # keep the logical output width so runtime can slice away static padding.
+        self.logical_output_size = weight.shape[1]
+
+        pad_k = (16 - weight.shape[0] % 16) % 16
+        pad_n = (16 - weight.shape[1] % 16) % 16
+        if pad_k == 0 and pad_n == 0:
+            return
+
+        # B is column-major [K, N]
+        padded_weight = torch.nn.functional.pad(
+            weight.t().contiguous(),
+            (0, pad_k, 0, pad_n),
+        ).t()
+        replace_parameter(layer, weight_name, padded_weight.data)
+        set_weight_attrs(
+            getattr(layer, weight_name),
+            {
+                "weight_loader": self.padded_weight_loader,
+            },
+        )
+
+        weight_scale = getattr(layer, weight_scale_name, None)
+        if weight_scale is not None and pad_n > 0 and weight_scale.numel() > 1:
+            flat_scale = weight_scale.reshape(-1)
+            padded_scale = self._pad_to_alignment(
+                flat_scale, dim=0, alignment=16, value=1.0
+            ).view(-1, *weight_scale.shape[1:])
+            replace_parameter(layer, weight_scale_name, padded_scale.data)
+            set_weight_attrs(
+                getattr(layer, weight_name),
+                {
+                    "weight_loader": self.padded_weight_loader,
+                },
+            )
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        padded_k, padded_n = B.shape
+        output_size = self.logical_output_size
+        assert output_size is not None
+        pad_k = padded_k - A.shape[1]
+        pad_n = padded_n - output_size
+
+        if pad_k > 0:
+            A = self._pad_to_alignment(A, dim=1, alignment=16)
+        if pad_n > 0 and bias is not None:
+            bias = self._pad_to_alignment(bias, dim=0, alignment=16)
+
+        output = ops.cutlass_scaled_mm(
+            A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
+        )
+
+        if pad_n > 0:
+            output = output[..., :output_size].contiguous()
+
+        return output.view(*output_shape[:-1], output_size)
+
+
+class CutlassFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
+    def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
+        super().__init__(config)
+        act_scale_descriptor = config.activation_quant_key.scale
+        self.quant_fp8 = QuantFP8(
+            static=act_scale_descriptor.static,
+            group_shape=act_scale_descriptor.group_shape,
+            num_token_padding=self.get_output_padding(),
+            use_ue8m0=False,
+            column_major_scales=True,
+        )
+
+    @classmethod
+    def is_supported(cls, compute_capability=None):
+        if not CUTLASS_BLOCK_FP8_SUPPORTED:
+            return (
+                False,
+                "The device compute capability of"
+                f"{compute_capability} is not supported.",
+            )
+        return True, None
+
+    @classmethod
+    def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
+        can_implement_base, reason = super().can_implement(config)
+        if not can_implement_base:
+            return can_implement_base, reason
+
+        act_quant_desc = config.activation_quant_key.scale
+        if act_quant_desc.group_shape != GroupShape(1, 128):
+            return (
+                False,
+                "Supports only dynamic per token group activation "
+                "quantization with group_shape=(1,128).",
+            )
+        return True, None
+
+    def apply_block_scaled_mm(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+    ) -> torch.Tensor:
+        out_dtype = self.config.out_dtype
+        return ops.cutlass_scaled_mm(
+            A,
+            B.T,
+            out_dtype=out_dtype,
+            scale_a=As,
+            scale_b=Bs.T,
+        )
+
+
+def cutlass_scaled_mm(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor,
+    Bs: torch.Tensor,
+    block_size: list[int],
+    output_dtype: torch.dtype = torch.float16,
+) -> torch.Tensor:
+    return ops.cutlass_scaled_mm(
+        A,
+        B.T,
+        out_dtype=output_dtype,
+        scale_a=As,
+        scale_b=Bs.T,
+    )

vllm/model_executor/kernels/linear/scaled_mm/deep_gemm.py

@@ -0,0 +1,158 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+import vllm.envs as envs
+from vllm.config import get_current_vllm_config
+from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+    deepgemm_post_process_fp8_weight_block,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+)
+from vllm.model_executor.utils import replace_parameter
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import (
+    fp8_gemm_nt,
+    is_deep_gemm_e8m0_used,
+    is_deep_gemm_supported,
+    should_auto_disable_deep_gemm,
+    should_use_deepgemm_for_fp8_linear,
+)
+from vllm.utils.torch_utils import direct_register_custom_op
+
+from .BlockScaledMMLinearKernel import (
+    Fp8BlockScaledMMLinearKernel,
+    FP8ScaledMMLinearLayerConfig,
+)
+
+
+class DeepGemmFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
+    def __init__(self, config: FP8ScaledMMLinearLayerConfig):
+        super().__init__(config)
+        self.use_deep_gemm_e8m0 = is_deep_gemm_e8m0_used()
+        act_scale_descriptor = config.activation_quant_key.scale
+        self.is_deep_gemm_supported = is_deep_gemm_supported()
+        self.quant_fp8 = QuantFP8(
+            static=False,
+            group_shape=act_scale_descriptor.group_shape,
+            use_ue8m0=self.use_deep_gemm_e8m0,
+            tma_aligned_scales=envs.VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES,
+            column_major_scales=True,
+        )
+
+    @classmethod
+    def is_supported(cls, compute_capability=None):
+        if not current_platform.is_cuda():
+            return False, "DeepGEMM is only supported on cuda platform"
+        if not is_deep_gemm_supported():
+            return False, "Currently, only Hopper and Blackwell GPUs are supported."
+        return True, None
+
+    @classmethod
+    def can_implement(cls, config):
+        can_implement_base, reason = super().can_implement(config)
+        if not can_implement_base:
+            return can_implement_base, reason
+        if config.out_dtype != torch.bfloat16:
+            return (False, "Supports only output dtype of bfloat16")
+
+        act_quant_desc = config.activation_quant_key.scale
+        if act_quant_desc.group_shape != GroupShape(1, 128):
+            return (
+                False,
+                "Supports only dynamic per token group activation "
+                "quantization with group_shape=(1,128).",
+            )
+        model_config = get_current_vllm_config().model_config
+
+        if model_config is None:
+            return False, "Model configuration is required."
+
+        model_type = getattr(model_config.hf_text_config, "model_type", None)
+        if should_auto_disable_deep_gemm(model_type):
+            return False, f"Should not use deepgemm for model {model_type}"
+
+        if not should_use_deepgemm_for_fp8_linear(
+            config.out_dtype, config.weight_shape
+        ):
+            return False, "The provided metadata is not supported."
+        return True, None
+
+    def process_weights_after_loading(self, layer):
+        super().process_weights_after_loading(layer)
+        params = self._get_layer_params(layer)
+        assert layer.weight_block_size is not None
+
+        if self.is_deep_gemm_supported:
+            weight_scale_invs = params.weight_scale_inv
+            scale_attr = (
+                params.WEIGHT_SCALE_INV
+                if weight_scale_invs is not None
+                else params.WEIGHT_SCALE
+            )
+            dg_weight, dg_weight_scale = deepgemm_post_process_fp8_weight_block(
+                wq=params.weight,
+                ws=weight_scale_invs
+                if weight_scale_invs is not None
+                else params.weight_scale,
+                quant_block_shape=tuple(layer.weight_block_size),
+                use_e8m0=self.use_deep_gemm_e8m0,
+                is_bmm=getattr(layer, "is_bmm", False),
+                bmm_batch_size=getattr(layer, "bmm_batch_size", 0),
+            )
+            replace_parameter(layer, params.WEIGHT, dg_weight)
+            replace_parameter(layer, scale_attr, dg_weight_scale)
+
+    def apply_block_scaled_mm(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+    ) -> torch.Tensor:
+        out_dtype = self.config.out_dtype
+        output = torch.empty(
+            (A.shape[0], B.shape[0]),
+            dtype=out_dtype,
+            device=A.device,
+        )
+        torch.ops.vllm.fp8_gemm_nt_op(A, As, B, Bs, output, self.use_deep_gemm_e8m0)
+        return output
+
+
+def _fp8_gemm_nt_op(
+    q_input: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output: torch.Tensor,
+    use_deep_gemm_e8m0: bool,
+) -> None:
+    fp8_gemm_nt(
+        (q_input, input_scale),
+        (weight, weight_scale),
+        output,
+        is_deep_gemm_e8m0_used=use_deep_gemm_e8m0,
+    )
+
+
+def _fp8_gemm_nt_op_fake(
+    q_input: torch.Tensor,
+    input_scale: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    output: torch.Tensor,
+    use_deep_gemm_e8m0: bool,
+) -> None:
+    return None
+
+
+direct_register_custom_op(
+    "fp8_gemm_nt_op",
+    _fp8_gemm_nt_op,
+    mutates_args=["output"],
+    fake_impl=_fp8_gemm_nt_op_fake,
+)

vllm/model_executor/kernels/linear/scaled_mm/flashinfer.py

@@ -0,0 +1,338 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+from typing import ClassVar
+
+import torch
+
+import vllm.envs as envs
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+    per_token_group_quant_fp8,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    QuantKey,
+    kFp8StaticTensorSym,
+)
+from vllm.platforms import current_platform
+from vllm.utils.flashinfer import (
+    flashinfer_fp8_blockscale_gemm,
+    flashinfer_scaled_fp8_mm,
+    has_flashinfer,
+    is_flashinfer_fp8_blockscale_gemm_supported,
+    should_use_flashinfer_for_blockscale_fp8_gemm,
+)
+from vllm.utils.torch_utils import direct_register_custom_op
+
+from .BlockScaledMMLinearKernel import (
+    Fp8BlockScaledDynamicMMLinearKernel,
+    Fp8BlockScaledMMLinearKernel,
+)
+from .deep_gemm import DeepGemmFp8BlockScaledMMKernel, fp8_gemm_nt
+from .ScaledMMLinearKernel import (
+    FP8ScaledMMLinearKernel,
+    FP8ScaledMMLinearLayerConfig,
+)
+
+
+class FlashInferFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "requires CUDA."
+
+        if not has_flashinfer():
+            return False, "requires FlashInfer to be installed."
+
+        if compute_capability is not None and compute_capability < 100:
+            return False, "requires compute capability 100 and above."
+
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        per_tensor_activation_scales = (
+            c.activation_quant_key.scale.group_shape.is_per_tensor()
+        )
+        per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
+
+        if not (per_tensor_activation_scales and per_tensor_weight_scales):
+            return False, "requires per tensor activation and weight scales."
+
+        return True, None
+
+    def input_quant_key(self) -> QuantKey | None:
+        if self.config.activation_quant_key == kFp8StaticTensorSym:
+            return kFp8StaticTensorSym
+        return None
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        return flashinfer_scaled_fp8_mm(
+            A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
+        )
+
+
+class FlashInferFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
+    # FlashInfer accepts BF16 input and handles FP8 conversion internally.
+    apply_input_quant: ClassVar[bool] = False
+
+    def __init__(self, config: FP8ScaledMMLinearLayerConfig) -> None:
+        super().__init__(config)
+
+    @classmethod
+    def can_implement(cls, config: FP8ScaledMMLinearLayerConfig):
+        can_implement_base, reason = super().can_implement(config)
+        if not can_implement_base:
+            return can_implement_base, reason
+
+        act_quant_desc = config.activation_quant_key.scale
+        if act_quant_desc.group_shape != GroupShape(1, 128):
+            return (
+                False,
+                "Supports only dynamic per token group activation "
+                "quantization with group_shape=(1,128).",
+            )
+
+        if not should_use_flashinfer_for_blockscale_fp8_gemm(
+            is_flashinfer_fp8_blockscale_gemm_supported(),
+            config.out_dtype,
+            config.input_dtype,
+            config.weight_quant_key.dtype,
+            config.weight_shape,
+        ):
+            return (
+                False,
+                "The provided metadata is not supported.",
+            )
+
+        return True, None
+
+    @classmethod
+    def is_supported(cls, compute_capability=None):
+        if not current_platform.is_cuda():
+            return False, "only cuda devices are supported."
+
+        if not is_flashinfer_fp8_blockscale_gemm_supported():
+            return False, "FlashInfer block-scale FP8 GEMM is not available."
+
+        return True, None
+
+    def apply_block_scaled_mm(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+    ) -> torch.Tensor:
+        # A is BF16 — FlashInfer handles FP8 conversion internally.
+        # As is a placeholder (apply_input_quant=False) and is not used here.
+        return torch.ops.vllm.flashinfer_fp8_blockscale_gemm(
+            A,  # BF16 input
+            B,  # FP8 weight
+            Bs,  # Weight scales
+        )
+
+
+class FlashInferFp8DeepGEMMDynamicBlockScaledKernel(
+    Fp8BlockScaledDynamicMMLinearKernel
+):
+    """
+    Conditional FlashInfer / DeepGEMM FP8 block-scaled GEMM.
+
+    Dispatches between two kernels based on input batch size:
+    - Small batches (M < 32): FlashInfer's swapAB trick for better utilisation.
+    - Large batches (M >= 32): DeepGEMM for peak throughput.
+
+    apply_input_quant is False because FlashInfer accepts BF16 input and
+    handles FP8 conversion internally.  The DeepGEMM branch therefore
+    quantises BF16→FP8 inside apply_mm via a closure before dispatching to
+    the DeepGEMM kernel — keeping both branches compatible with the single
+    BF16 tensor operand list passed by torch.cond.
+    """
+
+    base_type: ClassVar[type[FlashInferFp8BlockScaledMMKernel]] = (
+        FlashInferFp8BlockScaledMMKernel
+    )
+    fallback_type: ClassVar[type[DeepGemmFp8BlockScaledMMKernel]] = (
+        DeepGemmFp8BlockScaledMMKernel
+    )
+    apply_input_quant: ClassVar[bool] = False
+
+    def __init__(self, config: FP8ScaledMMLinearLayerConfig):
+        super().__init__(config)
+        self.base: FlashInferFp8BlockScaledMMKernel
+        self.fallback: DeepGemmFp8BlockScaledMMKernel
+
+    def process_weights_after_loading(self, layer: torch.nn.Module):
+        # DeepGEMM need post-processing; both kernels share the same
+        # parameter tensor layout so processing once is sufficient.
+        self.fallback.process_weights_after_loading(layer)
+
+    def apply_block_scaled_mm(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+    ) -> torch.Tensor:
+        group_size = self.weight_group_shape.col
+        use_deep_gemm_e8m0 = self.fallback.use_deep_gemm_e8m0
+
+        return torch.ops.vllm.dynamic_flashinfer_deepgemm_blockscale_gemm(
+            A, B, Bs, group_size, use_deep_gemm_e8m0
+        )
+
+
+def _flashinfer_fp8_blockscale_gemm_impl(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+) -> torch.Tensor:
+    return flashinfer_fp8_blockscale_gemm(
+        input=input,
+        weight=weight,
+        weight_scale=weight_scale,
+        out_dtype=torch.bfloat16,
+    )
+
+
+def _flashinfer_fp8_blockscale_gemm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Required fake/meta implementation for torch.compile graph tracing.
+    """
+    return torch.empty(
+        input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device
+    )
+
+
+direct_register_custom_op(
+    "flashinfer_fp8_blockscale_gemm",
+    _flashinfer_fp8_blockscale_gemm_impl,
+    fake_impl=_flashinfer_fp8_blockscale_gemm_fake,
+)
+
+
+def _dynamic_flashinfer_deepgemm_blockscale_gemm_impl(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    group_size: int,
+    use_deep_gemm_e8m0: bool,
+) -> torch.Tensor:
+    """
+    Conditional FlashInfer FP8 blockscale GEMM with batch-size-dependent selection.
+
+    This function switches between two optimized kernels based on the input batch size:
+    - For small batches (M < 32): Uses FlashInfer's DeepGEMM swapAB optimization.
+    - For larger batches (M >= 32): Uses the official DeepGEMM kernel.
+
+    The conditional logic must use torch.cond() instead of a simple if-else statement
+    to maintain compatibility with torch.compile graph compilation.
+
+    This batch-size-dependent selection is essential for maintaining model accuracy.
+    Benchmarks on GSM8K show a significant accuracy gap (88% vs 95%) for DeepSeek-V3.1
+    when using FlashInfer's DeepGEMM on M>=32. The M < 32 strategy fixes the accuracy
+    drop.
+
+    Args:
+        input: Input tensor of shape (batch_size, input_dim) in FP8 format
+        weight: Weight tensor of shape (output_dim, input_dim) in FP8 format
+        weight_scale: Scale factors for weight quantization (per-group)
+        group_size: Quantization group size for the weight tensor
+        use_deep_gemm_e8m0: Whether to use the E8M0 format in DeepGEMM quantization
+
+    Returns:
+        Output tensor of shape (batch_size, output_dim) in bfloat16 format
+    """
+
+    def run_flashinfer_deepgemm_swapAB(
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        weight_scale: torch.Tensor,
+    ) -> torch.Tensor:
+        return flashinfer_fp8_blockscale_gemm(
+            input=input,
+            weight=weight,
+            weight_scale=weight_scale,
+            out_dtype=torch.bfloat16,
+        )
+
+    def run_deepgemm(
+        input: torch.Tensor,
+        weight: torch.Tensor,
+        weight_scale: torch.Tensor,
+    ) -> torch.Tensor:
+        q_input, input_scale = per_token_group_quant_fp8(
+            input,
+            group_size=group_size,
+            column_major_scales=True,
+            use_ue8m0=use_deep_gemm_e8m0,
+        )
+        output = torch.empty(
+            (q_input.shape[0], weight.shape[0]),
+            dtype=torch.bfloat16,
+            device=q_input.device,
+        )
+        fp8_gemm_nt(
+            (q_input, input_scale),
+            (weight, weight_scale),
+            output,
+            is_deep_gemm_e8m0_used=use_deep_gemm_e8m0,
+        )
+        return output
+
+    if envs.VLLM_BATCH_INVARIANT:
+        return run_deepgemm(input, weight, weight_scale)
+
+    condition = input.shape[0] < 32
+
+    # PyTorch's torch.compile cannot handle input-dependent control flow in standard
+    # Python conditionals. torch.cond() explicitly registers both code paths in the
+    # computation graph, allowing torch.compile to capture both branches.
+    # without torch.cond, the M < 32 condition won't be able to be captured by torch
+    # compile
+    return torch.cond(
+        condition,
+        run_flashinfer_deepgemm_swapAB,
+        run_deepgemm,
+        (input, weight, weight_scale),
+    )
+
+
+def _dynamic_flashinfer_deepgemm_blockscale_gemm_fake(
+    input: torch.Tensor,
+    weight: torch.Tensor,
+    weight_scale: torch.Tensor,
+    group_size: int,
+    use_deep_gemm_e8m0: bool,
+) -> torch.Tensor:
+    """
+    Required fake/meta implementation for torch.compile graph tracing.
+    """
+    return torch.empty(
+        input.shape[0], weight.shape[0], dtype=torch.bfloat16, device=input.device
+    )
+
+
+direct_register_custom_op(
+    "dynamic_flashinfer_deepgemm_blockscale_gemm",
+    _dynamic_flashinfer_deepgemm_blockscale_gemm_impl,
+    fake_impl=_dynamic_flashinfer_deepgemm_blockscale_gemm_fake,
+)

vllm/model_executor/kernels/linear/scaled_mm/humming.py

@@ -0,0 +1,156 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.utils.humming_utils import (
+    convert_linear_layer_to_humming_standard,
+    prepare_humming_layer,
+)
+from vllm.platforms import current_platform
+
+from .ScaledMMLinearKernel import (
+    FP8ScaledMMLinearKernel,
+    FP8ScaledMMLinearLayerConfig,
+    Int8ScaledMMLinearKernel,
+    Int8ScaledMMLinearLayerConfig,
+)
+
+logger = init_logger(__name__)
+
+
+class HummingFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    """Humming GEMM Kernel for FP8."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "Humming only supported on CUDA"
+
+        if not current_platform.has_device_capability(75):
+            return False, "Humming only supported on SM75+"
+
+        return True, None
+
+    @classmethod
+    def can_implement(
+        cls, config: FP8ScaledMMLinearLayerConfig
+    ) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        from vllm.utils.humming import dtypes
+
+        name_map = {"weight": "weight", "weight_scale": "weight_scale"}
+        scale_torch_dtype = self.config.weight_quant_key.scale.dtype
+        scale_dtype = dtypes.DataType.from_torch_dtype(scale_torch_dtype)
+
+        quant_config = {
+            "quant_method": "humming",
+            "dtype": "float8e4m3",
+            "scale_dtype": scale_dtype,
+        }
+
+        assert self.config.weight_quant_key.scale2 is None
+        scale_group_shape = self.config.weight_quant_key.scale.group_shape
+        if scale_group_shape.is_per_tensor():
+            quant_config["weight_scale_type"] = "tensor"
+            if not hasattr(layer, "global_scale") and hasattr(layer, "weight_scale"):
+                del name_map["weight_scale"]
+                name_map["global_scale"] = "weight_scale"
+        elif scale_group_shape.is_per_channel():
+            quant_config["weight_scale_type"] = "channel"
+        elif scale_group_shape.is_per_group():
+            quant_config["weight_scale_type"] = "group"
+            quant_config["group_size"] = scale_group_shape.col
+        else:
+            assert scale_group_shape.row > 0 and scale_group_shape.col > 0
+            quant_config["weight_scale_type"] = "block"
+            quant_config["weight_scale_group_size_n"] = scale_group_shape.row
+            quant_config["weight_scale_group_size"] = scale_group_shape.col
+
+            if hasattr(layer, "weight_scale_inv"):
+                name_map["weight_scale"] = "weight_scale_inv"
+
+        convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
+        prepare_humming_layer(layer, quant_config)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        from vllm.utils.humming import HummingMethod
+
+        flatten_inputs = x.view(-1, x.size(-1))
+        output = HummingMethod.forward_layer(
+            layer=layer,
+            inputs=flatten_inputs,
+            compute_config=layer.compute_config,
+        )
+        return output.view(*x.shape[:-1], output.size(-1))
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        pass
+
+
+class HummingInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
+    """Humming GEMM Kernel for INT8."""
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "Humming only supported on CUDA"
+
+        if not current_platform.has_device_capability(75):
+            return False, "Humming only supported on SM75+"
+
+        return True, None
+
+    @classmethod
+    def can_implement(
+        cls, config: Int8ScaledMMLinearLayerConfig
+    ) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        weight_name, weight_scale_name, *_ = self.layer_param_names
+        name_map = {"weight": weight_name, "weight_scale": weight_scale_name}
+        quant_config = {"quant_method": "humming", "dtype": "int8"}
+        weight = getattr(layer, weight_name)
+        weight.data = weight.data + 128
+
+        convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
+        prepare_humming_layer(layer, quant_config)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        from vllm.utils.humming import HummingMethod
+
+        flatten_inputs = x.view(-1, x.size(-1))
+        output = HummingMethod.forward_layer(
+            layer=layer,
+            inputs=flatten_inputs,
+            compute_config=layer.compute_config,
+        )
+        return output.view(*x.shape[:-1], output.size(-1))

vllm/model_executor/kernels/linear/scaled_mm/marlin.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Sequence
+
+import torch
+
+import vllm.envs as envs
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+    process_fp8_weight_block_strategy,
+)
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
+    apply_fp8_marlin_linear,
+    is_fp8_marlin_supported,
+    prepare_fp8_layer_for_marlin,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kFp8Static128BlockSym,
+)
+from vllm.model_executor.utils import replace_parameter
+from vllm.platforms import current_platform
+
+from .ScaledMMLinearKernel import (
+    FP8ScaledMMLinearKernel,
+    FP8ScaledMMLinearLayerConfig,
+)
+
+
+class MarlinFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    """
+    FP8 Marlin kernel for GPUs that lack FP8 hardware support.
+    Leverages the Marlin kernel for fast weight-only FP8 quantization.
+    """
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cuda():
+            return False, "requires CUDA."
+        # Check if platform supports FP8 Marlin
+        if not is_fp8_marlin_supported():
+            return False, "FP8 Marlin requires compute capability 7.5 or higher"
+        if envs.VLLM_BATCH_INVARIANT:
+            return False, "FP8 Marlin not supported for batch invariant execution."
+        if (
+            compute_capability is not None
+            and compute_capability >= 89
+            and not envs.VLLM_TEST_FORCE_FP8_MARLIN
+        ):
+            return (
+                False,
+                "To apply FP8 Marlin on high-capability GPUs, please set "
+                "VLLM_TEST_FORCE_FP8_MARLIN=1",
+            )
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def __init__(
+        self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
+    ) -> None:
+        super().__init__(c, layer_param_names)
+        self.marlin_input_dtype = None
+        self.block_quant = self.config.weight_quant_key in {kFp8Static128BlockSym}
+        self.size_k_first = not self.block_quant
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        if self.block_quant:
+            weight, weight_scale_inv = process_fp8_weight_block_strategy(
+                layer.weight, layer.weight_scale_inv
+            )
+            # Update layer with new values
+            replace_parameter(layer, "weight", weight.data)
+            replace_parameter(layer, "weight_scale_inv", weight_scale_inv.data)
+        # Non-block: callers must pass weight in (K, N) layout.
+
+        layer.input_scale = None
+        prepare_fp8_layer_for_marlin(
+            layer, self.size_k_first, input_dtype=self.marlin_input_dtype
+        )
+        del layer.input_scale
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.block_quant:
+            weight_scale = layer.weight_scale_inv
+        else:
+            weight_scale = layer.weight_scale
+        return apply_fp8_marlin_linear(
+            input=x,
+            weight=layer.weight,
+            weight_scale=weight_scale,
+            workspace=layer.workspace,
+            size_n=layer.output_size_per_partition,
+            size_k=layer.input_size_per_partition,
+            input_dtype=self.marlin_input_dtype,
+            bias=bias,
+        )
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        pass

vllm/model_executor/kernels/linear/scaled_mm/pytorch.py

@@ -0,0 +1,242 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import math
+
+import torch
+
+from vllm.config import CompilationMode, get_current_vllm_config
+from vllm.platforms import current_platform
+
+from .ScaledMMLinearKernel import (
+    FP8ScaledMMLinearKernel,
+    FP8ScaledMMLinearLayerConfig,
+)
+
+
+def _get_num_tokens(output_shape: list) -> int:
+    # torch._scaled_mm works with 2D tensors, so input tensors are
+    # flattened if they are 3D. If output_shape is 3D, num_tokens is
+    # the product of all dims except the last (hidden dim).
+    return math.prod(output_shape[:-1])
+
+
+class TorchFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    """
+    Base class for FP8 linear kernels using Torch.
+    Each subclass represents a kernel variant for
+    specific device capabilities and torch versions.
+    """
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not (current_platform.is_cuda_alike() or current_platform.is_cpu()):
+            return False, "requires ROCm, CUDA or CPU."
+
+        if compute_capability is not None and compute_capability < 89:
+            return False, "requires compute capability 89 and above."
+
+        return True, None
+
+    def get_output_padding(self) -> int | None:
+        # Note: we pad the input because torch._scaled_mm is more performant
+        # for matrices with batch dimension > 16.
+        # This could change in the future.
+        # We also don't pad when using torch.compile,
+        # as it breaks with dynamic shapes.
+        #
+        # The perf gain is still relevant as of 16/1/2026
+        # torch version == 2.9.0. More details in the link below:
+        # https://github.com/vllm-project/vllm/issues/32269
+        vllm_config = get_current_vllm_config().compilation_config
+        pad_output = vllm_config.mode < CompilationMode.VLLM_COMPILE
+        return 17 if pad_output else None
+
+
+class PerTensorTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        per_tensor_activation_scales = (
+            c.activation_quant_key.scale.group_shape.is_per_tensor()
+        )
+        per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
+
+        if not (per_tensor_activation_scales and per_tensor_weight_scales):
+            return False, "requires per tensor activation and weight scales."
+        return True, None
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        # torch._scaled_mm under torch.compile does not support 0-D scales
+        if As.dim() == 0:
+            As = As.view(1)
+        if Bs.dim() == 0:
+            Bs = Bs.view(1)
+
+        output = torch._scaled_mm(
+            A, B, out_dtype=out_dtype, scale_a=As, scale_b=Bs, bias=bias
+        )
+        # A fix for discrepancy in scaled_mm which returns tuple
+        # for torch < 2.5 and a single value in torch >= 2.5
+        if type(output) is tuple and len(output) == 2:
+            output = output[0]
+
+        num_tokens = _get_num_tokens(output_shape)
+        return torch.narrow(output, 0, 0, num_tokens).view(*output_shape)
+
+
+class RowWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_rocm():
+            return False, "requires ROCm."
+
+        from vllm.platforms.rocm import on_mi3xx
+
+        if not on_mi3xx():
+            return False, "requires MI3xx."
+
+        if compute_capability is not None and compute_capability < 94:
+            return False, "requires compute capability 94 and above."
+
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        per_tensor_activation_scales = (
+            c.activation_quant_key.scale.group_shape.is_per_tensor()
+        )
+        per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
+
+        if c.out_dtype == torch.float16:
+            # hipblaslt rowwise _scaled_mm only supports BFloat16
+            return False, "supports BFloat16 output data type only."
+
+        if per_tensor_activation_scales or per_tensor_weight_scales:
+            return False, "cannot be used with per tensor activation and weight scales."
+
+        return True, None
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        #  Note:
+        #  For now it has only been validated on ROCm platform.
+        #  fp8 rowwise scaling in torch._scaled_mm is introduced in
+        #  https://github.com/pytorch/pytorch/pull/144432 using
+        #  hipBLASLt and ROCm 6.3, which only exists in torch 2.7 and above.
+        #
+        #  For CUDA platform please validate if the torch._scaled_mm supports
+        #  rowwise scaled GEMM before using it
+
+        # torch._scaled_mm rowwise requires scale_a = (m, 1), scale_b = (1, n).
+        # CompressedTensors stores weight_scale as (n, 1), so `.t()` yields (1, n).
+        # ModelOpt FP8_PER_CHANNEL_PER_TOKEN stores it as 1-D (n,); reshape to
+        # (1, n) so both paths satisfy the rowwise contract.
+        scale_b = Bs.view(1, -1) if Bs.dim() == 1 else Bs.t()
+        if As.dim() == 1:
+            As = As.view(-1, 1)
+
+        # Fused GEMM_DQ Rowwise GEMM
+        output = torch._scaled_mm(
+            A,
+            B,
+            out_dtype=out_dtype,
+            scale_a=As,
+            scale_b=scale_b,
+            bias=bias,
+        )
+
+        num_tokens = _get_num_tokens(output_shape)
+        return torch.narrow(output, 0, 0, num_tokens).view(*output_shape)
+
+
+class ChannelWiseTorchFP8ScaledMMLinearKernel(TorchFP8ScaledMMLinearKernel):
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        per_tensor_activation_scales = (
+            c.activation_quant_key.scale.group_shape.is_per_tensor()
+        )
+        per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
+
+        if per_tensor_activation_scales and per_tensor_weight_scales:
+            return False, "cannot be used with per tensor activation and weight scales."
+
+        return True, None
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        # Use unfused DQ due to limitations with scaled_mm
+
+        # Symmetric quantized GEMM by definition computes the following:
+        #   C = (s_x * X) (s_w * W) + bias
+        # This is equivalent to dequantizing the weights and activations
+        # before applying a GEMM.
+        #
+        # In order to compute quantized operands, a quantized kernel
+        # will rewrite the above like so:
+        #   C = s_w * s_x * (X * W) + bias
+        #
+        # For the scaled_mm fallback case, we break this down, since it
+        # does not support s_w being a vector.
+
+        # Input scaling factors are no longer optional in _scaled_mm starting
+        # from pytorch 2.5. Allocating a dummy tensor to pass as scales
+        dummy_tensor = torch.ones(1, dtype=torch.float32, device=A.device)
+
+        # GEMM
+        # This computes C = (X * W).
+        # Output in fp32 to allow subsequent ops to happen in-place
+        output = torch._scaled_mm(
+            A,
+            B,
+            scale_a=dummy_tensor,
+            scale_b=dummy_tensor,
+            out_dtype=torch.float32,
+        )
+        # A fix for discrepancy in scaled_mm which returns tuple
+        # for torch < 2.5 and a single value in torch >= 2.5
+        if type(output) is tuple and len(output) == 2:
+            output = output[0]
+
+        # Unpad (undo num_token_padding)
+        num_tokens = _get_num_tokens(output_shape)
+        output = torch.narrow(output, 0, 0, num_tokens)
+        x_scale = torch.narrow(As, 0, 0, num_tokens)
+
+        # DQ
+        # C = sw * sx * (X * W) + bias
+        output = output * x_scale * Bs.t()
+        if bias is not None:
+            output = output + bias
+        return output.to(out_dtype).view(*output_shape)

vllm/model_executor/kernels/linear/scaled_mm/rocm.py

@@ -0,0 +1,118 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.platforms import current_platform
+from vllm.utils.platform_utils import num_compute_units
+from vllm.utils.torch_utils import direct_register_custom_op
+
+from .ScaledMMLinearKernel import (
+    FP8ScaledMMLinearKernel,
+    FP8ScaledMMLinearLayerConfig,
+)
+
+
+def rocm_per_tensor_float_w8a8_scaled_mm_impl(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    out_dtype: torch.dtype,
+    As: torch.Tensor,
+    Bs: torch.Tensor,
+    bias: torch.Tensor,
+) -> torch.Tensor:
+    if (
+        A.shape[0] <= 4
+        and B.shape[0] % 16 == 0  # M TODO: needed?
+        and B.shape[1] % 16 == 0  # K
+        and ((bias is None) or (bias.dtype == out_dtype))
+    ):
+        output = ops.wvSplitKQ(
+            B.t(),
+            A,
+            out_dtype,
+            As,
+            Bs,
+            num_compute_units(),
+            bias,
+        )
+    # Fallback
+    else:
+        output = torch._scaled_mm(
+            A,
+            B,
+            out_dtype=out_dtype,
+            scale_a=As,
+            scale_b=Bs,
+            bias=bias,
+        )
+    return output
+
+
+def rocm_per_tensor_float_w8a8_scaled_mm_fake(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    out_dtype: torch.dtype,
+    As: torch.Tensor,
+    Bs: torch.Tensor,
+    bias: torch.Tensor,
+) -> torch.Tensor:
+    return A.new_empty((*A.shape[:-1], B.shape[1]), dtype=out_dtype)
+
+
+if current_platform.is_rocm():
+    direct_register_custom_op(
+        op_name="rocm_per_tensor_float_w8a8_scaled_mm_impl",
+        op_func=rocm_per_tensor_float_w8a8_scaled_mm_impl,
+        fake_impl=rocm_per_tensor_float_w8a8_scaled_mm_fake,
+    )
+
+
+class ROCmFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_rocm():
+            return False, "requires ROCm."
+
+        from vllm.platforms.rocm import on_gfx12x, on_mi3xx
+
+        if not (on_mi3xx() or on_gfx12x()):
+            return False, "requires MI3xx or gfx12x"
+
+        if not envs.VLLM_ROCM_USE_SKINNY_GEMM:
+            return False, "requires VLLM_ROCM_USE_SKINNY_GEMM to be enabled."
+
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        per_tensor_activation_scales = (
+            c.activation_quant_key.scale.group_shape.is_per_tensor()
+        )
+        per_tensor_weight_scales = c.weight_quant_key.scale.group_shape.is_per_tensor()
+
+        if not (per_tensor_activation_scales and per_tensor_weight_scales):
+            return False, "requires per tensor activation and weight scales."
+
+        return True, None
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        output = torch.ops.vllm.rocm_per_tensor_float_w8a8_scaled_mm_impl(
+            A, B, out_dtype, As, Bs, bias
+        )
+        return torch.narrow(output, 0, 0, A.shape[0]).view(*output_shape)

vllm/model_executor/kernels/linear/scaled_mm/triton.py

@@ -0,0 +1,220 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm import (  # noqa: E501
+    triton_scaled_mm,
+)
+from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    convert_to_channelwise,
+)
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import direct_register_custom_op
+
+from .BlockScaledMMLinearKernel import (
+    Fp8BlockScaledMMLinearKernel,
+)
+from .cutlass import CutlassInt8ScaledMMLinearKernel
+from .ScaledMMLinearKernel import (
+    Int8ScaledMMLinearLayerConfig,
+)
+
+
+class TritonInt8ScaledMMLinearKernel(CutlassInt8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if current_platform.is_cuda_alike():
+            return True, None
+        return False, "requires ROCm or CUDA."
+
+    @classmethod
+    def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        w_q, _, i_s, _, _ = self._get_layer_params(layer)
+        w_q_name, w_s_name, i_s_name, i_zp_name, azp_adj_name = self.layer_param_names
+
+        replace_parameter(
+            layer,
+            w_q_name,
+            torch.nn.Parameter(w_q.t().data, requires_grad=False),
+        )
+
+        # WEIGHT SCALE
+        # Triton kernel supports only per-tensor and per-channel.
+        # If we have a fused module (QKV, MLP) with per tensor scales (thus N
+        # scales being passed to the kernel), convert to the per-channel case.
+        is_fused_module = len(layer.logical_widths) > 1
+        weight_scale = getattr(layer, w_s_name)
+        if is_fused_module and not self.config.is_channelwise:
+            weight_scale = convert_to_channelwise(weight_scale, layer.logical_widths)
+        replace_parameter(
+            layer,
+            w_s_name,
+            torch.nn.Parameter(weight_scale.data, requires_grad=False),
+        )
+
+        # INPUT SCALE
+        if self.config.is_static_input_scheme:
+            assert i_s is not None
+
+            if self.config.input_symmetric:
+                replace_parameter(
+                    layer,
+                    i_s_name,
+                    torch.nn.Parameter(i_s.max(), requires_grad=False),
+                )
+                setattr(layer, i_zp_name, None)
+            else:
+                input_zero_point = getattr(layer, i_zp_name)
+
+                # Reconstruct the ranges to find a single scale and azp
+                int8_traits = torch.iinfo(torch.int8)
+                azps = input_zero_point.to(dtype=torch.int32)
+                range_max = (i_s * (int8_traits.max - azps)).max()
+                range_min = (i_s * (int8_traits.min - azps)).min()
+
+                scale = (range_max - range_min) / (int8_traits.max - int8_traits.min)
+                replace_parameter(
+                    layer,
+                    i_s_name,
+                    torch.nn.Parameter(scale, requires_grad=False),
+                )
+
+                # AZP loaded as int8 but used as int32
+                azp = (int8_traits.min - range_min / scale).to(dtype=torch.int32)
+                replace_parameter(
+                    layer,
+                    i_zp_name,
+                    torch.nn.Parameter(azp, requires_grad=False),
+                )
+        else:
+            setattr(layer, i_s_name, None)
+            setattr(layer, i_zp_name, None)
+
+        # azp_adj is the AZP adjustment term, used to account for weights.
+        # It does not depend on scales or azp, so it is the same for
+        # static and dynamic quantization.
+        # See csrc/quantization/w8a8/cutlass/Epilogues.md for the math.
+        if not self.config.input_symmetric:
+            weight = getattr(layer, w_q_name)
+            # weight is already transposed to [K, N], sum over K (dim=0)
+            azp_adj = weight.sum(dim=0, keepdim=True, dtype=torch.int32)
+            if self.config.is_static_input_scheme:
+                # Fold azp into azp_adj for the per-tensor case
+                azp_adj = getattr(layer, i_zp_name) * azp_adj
+            setattr(
+                layer,
+                azp_adj_name,
+                torch.nn.Parameter(azp_adj, requires_grad=False),
+            )
+        else:
+            setattr(layer, azp_adj_name, None)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        w_q, w_s, i_s, i_zp, azp_adj = self._get_layer_params(layer)
+
+        symmetric = azp_adj is None
+        x_q, x_s, x_zp = ops.scaled_int8_quant(
+            x.contiguous(), i_s, i_zp, symmetric=symmetric
+        )
+
+        out = triton_scaled_mm(
+            x_q, w_q, scale_a=x_s, scale_b=w_s, out_dtype=x.dtype, bias=bias
+        )
+
+        if azp_adj is not None:
+            # Asymmetric quantization: subtract the zero-point correction.
+            # D = scale_a * scale_b * (A_q @ B_q - azp * azp_adj) + bias
+            # triton_scaled_mm already computed scale_a * scale_b * (A_q @ B_q) + bias
+            # so we subtract scale_a * scale_b * azp * azp_adj
+            #
+            # x_s: [M, 1] or scalar, w_s: [N, 1] or scalar, azp_adj: [1, N]
+            # Reshape w_s from [N, 1] to [1, N] for proper broadcasting.
+            w_s_row = w_s.view(1, -1) if w_s.dim() > 0 else w_s
+            static = i_zp is not None
+            if not static and x_zp is not None:
+                # Dynamic per-token: azp is per-token, azp_adj is per-channel
+                # x_zp: [M, 1], azp_adj: [1, N]
+                out -= x_s * w_s_row * (x_zp * azp_adj).to(x.dtype)
+            else:
+                # Static per-tensor: azp already folded into azp_adj
+                out -= (x_s * w_s_row * azp_adj).to(x.dtype)
+
+        return out
+
+
+class TritonFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
+    @classmethod
+    def is_supported(cls, compute_capability=None):
+        if not (current_platform.is_cuda_alike() or current_platform.is_xpu()):
+            return False, "only CUDA-alike and XPU devices are supported."
+        return True, None
+
+    def apply_block_scaled_mm(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+    ) -> torch.Tensor:
+        return torch.ops.vllm.w8a8_triton_block_scaled_mm_func(
+            A,
+            B,
+            As,
+            Bs,
+            list(self.weight_group_shape),
+            self.config.out_dtype,
+        )
+
+
+# TODO we should be able to change the type of block_size to GroupShape
+# after we resolve GroupShape compilation issue
+# https://github.com/vllm-project/vllm/issues/25270
+def _w8a8_triton_block_scaled_mm_func(
+    qx: torch.Tensor,
+    weight: torch.Tensor,
+    x_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    block_size: list[int],
+    output_dtype: torch.dtype,
+) -> torch.Tensor:
+    from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+        w8a8_triton_block_scaled_mm,
+    )
+
+    return w8a8_triton_block_scaled_mm(
+        qx, weight, x_scale, weight_scale, block_size, output_dtype
+    )
+
+
+def _w8a8_triton_block_scaled_mm_fake(
+    qx: torch.Tensor,
+    weight: torch.Tensor,
+    x_scale: torch.Tensor,
+    weight_scale: torch.Tensor,
+    block_size: list[int],
+    output_dtype: torch.dtype,
+) -> torch.Tensor:
+    return torch.empty(
+        (qx.size(0), weight.size(0)), dtype=output_dtype, device=qx.device
+    )
+
+
+direct_register_custom_op(
+    "w8a8_triton_block_scaled_mm_func",
+    _w8a8_triton_block_scaled_mm_func,
+    fake_impl=_w8a8_triton_block_scaled_mm_fake,
+)

vllm/model_executor/kernels/linear/scaled_mm/xpu.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Sequence
+
+import torch
+
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    kFp8DynamicTensorSym,
+    kFp8DynamicTokenSym,
+    kFp8StaticChannelSym,
+    kFp8StaticTensorSym,
+    kFp8StaticTokenSym,
+)
+from vllm.model_executor.utils import replace_parameter
+from vllm.platforms import current_platform
+
+from .BlockScaledMMLinearKernel import Fp8BlockScaledMMLinearKernel
+from .ScaledMMLinearKernel import FP8ScaledMMLinearKernel, FP8ScaledMMLinearLayerConfig
+
+
+class XPUW8A8FP8LinearKernel(FP8ScaledMMLinearKernel):
+    _SUPPORTED_ACT_QUANT_KEYS = {
+        kFp8DynamicTensorSym,
+        kFp8DynamicTokenSym,
+        kFp8StaticTensorSym,
+        kFp8StaticTokenSym,
+    }
+    _SUPPORTED_WEIGHT_QUANT_KEYS = {
+        kFp8StaticChannelSym,
+        kFp8StaticTensorSym,
+    }
+
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_xpu():
+            return False, "XPUW8A8FP8Linear only support on XPU"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        if c.weight_quant_key not in cls._SUPPORTED_WEIGHT_QUANT_KEYS:
+            return (
+                False,
+                "XPUW8A8FP8Linear only support per-channel and per-tensor quantization",
+            )
+        if c.activation_quant_key not in cls._SUPPORTED_ACT_QUANT_KEYS:
+            return (
+                False,
+                "XPUW8A8FP8Linear only support per-tensor and per-token activation "
+                "quantization",
+            )
+        if c.weight_quant_key.dtype not in {torch.float8_e5m2, torch.float8_e4m3fn}:
+            return False, "XPUW8A8FP8Linear only support FP8 weight dtype"
+        if c.activation_quant_key.dtype not in {
+            torch.float8_e5m2,
+            torch.float8_e4m3fn,
+        }:
+            return False, "XPUW8A8FP8Linear only support FP8 activation dtype"
+        return True, None
+
+    def __init__(
+        self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
+    ) -> None:
+        super().__init__(c, layer_param_names)
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Ensure weight is stored as C-contiguous [K, N] (KN layout).
+
+        Checkpoints store weight as [N, K]; fp8_gemm requires [K, N],
+        C-contiguous.  Three incoming layouts are possible:
+          • [N, K] C-contiguous   ← direct checkpoint   → .t().contiguous()
+          • [K, N] Fortran-order  ← fp8.py's weight.t() → .contiguous()
+          • [K, N] C-contiguous   ← already correct     → no-op
+
+        For square weights (K == N) the shape is ambiguous; contiguity is used
+        as a proxy: C-contiguous ≡ checkpoint [N, K] (needs transpose);
+        Fortran-order ≡ fp8.py already transposed (needs only contiguous).
+        """
+        K = getattr(layer, "input_size_per_partition", self.config.weight_shape[1])
+        N = getattr(layer, "output_size_per_partition", self.config.weight_shape[0])
+        w = layer.weight
+
+        if w.shape not in {(K, N), (N, K)}:
+            raise ValueError(
+                f"XPUFP8ScaledMM expects weight shape ({K},{N}) or ({N},{K}), "
+                f"but got {tuple(w.shape)}"
+            )
+
+        needs_transpose = w.shape == (N, K) if K != N else w.is_contiguous()
+        layer_weight = w.t() if needs_transpose else w
+        replace_parameter(layer, "weight", layer_weight)
+        ws = layer.weight_scale
+        if ws.numel() == 1:
+            replace_parameter(layer, "weight_scale", ws.reshape(1))
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        # B is C-contiguous [K, N] from process_weights_after_loading.
+        # fp8_gemm routes on scale dtype (float32) and numel:
+        #   As [1]   → per-tensor  (numel==1 branch)
+        #   As [M,1] → per-token   (group={1,K} branch, broadcast across K)
+        #   Bs [1]   → per-tensor
+        #   Bs [N]   → per-channel (mask=bit1 branch)
+        # No shape manipulation needed here.
+        output = torch.ops._xpu_C.fp8_gemm(A, B, out_dtype, As, Bs, bias)
+        return output.view(*output_shape)
+
+
+class XPUW8A16FP8LinearKernel(FP8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_xpu():
+            return False, "XPUW8A16FP8Linear only support on XPU"
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: FP8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        if c.weight_quant_key not in {kFp8StaticChannelSym, kFp8StaticTensorSym}:
+            return (
+                False,
+                "XPUW8A16FP8Linear only support per-channel and per-tensor "
+                "quantization",
+            )
+        if c.weight_quant_key.dtype not in {torch.float8_e5m2, torch.float8_e4m3fn}:
+            return False, "XPUW8A16FP8Linear only support FP8 weight dtype"
+        return True, None
+
+    def __init__(
+        self, c: FP8ScaledMMLinearLayerConfig, layer_param_names: Sequence[str]
+    ) -> None:
+        assert self.can_implement(c)[0]
+        assert self.is_supported()[0]
+        self.config = c
+        self.layer_param_names = layer_param_names
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        # fp8_gemm_w8a16 expects weight in [in, out] layout.
+        # Transpose if weight is still in [out, in] layout.
+        # For square matrices, use contiguity as tie-breaker:
+        # checkpoint weights are contiguous, .t() views are not.
+        weight = layer.weight
+        out_features, in_features = self.config.weight_shape
+
+        if weight.shape == (out_features, in_features) and (
+            in_features != out_features or weight.is_contiguous()
+        ):
+            replace_parameter(layer, "weight", weight.data.t())
+        # else: already in [in, out] layout — no-op
+
+        weight_scale = layer.weight_scale.t().contiguous()
+        replace_parameter(layer, "weight_scale", weight_scale.data)
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        weight = layer.weight
+        weight_scale = layer.weight_scale
+        return torch.ops._xpu_C.fp8_gemm_w8a16(x, weight, weight_scale, bias)
+
+    def apply_scaled_mm(
+        self,
+        *,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        out_dtype: torch.dtype,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+        bias: torch.Tensor | None,
+        output_shape: list,
+    ) -> torch.Tensor:
+        pass
+
+
+class XPUFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_xpu():
+            return False, "XPUFp8BlockScaledMM only support on XPU"
+        return True, None
+
+    def apply_block_scaled_mm(
+        self,
+        A: torch.Tensor,
+        B: torch.Tensor,
+        As: torch.Tensor,
+        Bs: torch.Tensor,
+    ) -> torch.Tensor:
+        # Weight is [N, K]. Use .t() to create a [K, N] view without copying.
+        # Bs is [N/128, K/128] — transpose to [K/128, N/128] for oneDNN.
+        return torch.ops._xpu_C.fp8_gemm(
+            A,
+            B.t(),
+            self.config.out_dtype,
+            As,
+            Bs.t().contiguous(),
+            torch.Tensor(),
+        )

vllm/model_executor/kernels/linear/scaled_mm/zentorch.py

@@ -0,0 +1,98 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Zentorch dynamic-symmetric W8A8 int8 linear kernel for AMD Zen CPUs.
+
+Selected by ``choose_scaled_mm_linear_kernel`` ahead of the generic
+oneDNN-backed ``CPUInt8ScaledMMLinearKernel``. When ``is_supported`` or
+``can_implement`` rejects a layer, the selector falls through to the next
+kernel in ``_POSSIBLE_INT8_KERNELS[PlatformEnum.CPU]``.
+"""
+
+import torch
+
+from vllm.logger import init_logger
+from vllm.model_executor.kernels.linear.zentorch_utils import has_zentorch_op
+from vllm.model_executor.layers.quantization.utils import replace_parameter
+from vllm.platforms import current_platform
+
+from .ScaledMMLinearKernel import (
+    Int8ScaledMMLinearKernel,
+    Int8ScaledMMLinearLayerConfig,
+)
+
+logger = init_logger(__name__)
+
+
+class ZentorchInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
+    @classmethod
+    def is_supported(
+        cls, compute_capability: int | None = None
+    ) -> tuple[bool, str | None]:
+        if not current_platform.is_cpu():
+            return False, "requires CPU."
+        if not current_platform.is_zen_cpu():
+            return False, "requires AMD Zen CPU."
+        if not has_zentorch_op(["zentorch_dynamic_qlinear"]):
+            return (
+                False,
+                "torch.ops.zentorch.zentorch_dynamic_qlinear is not registered.",
+            )
+        return True, None
+
+    @classmethod
+    def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
+        if c.is_static_input_scheme:
+            return False, "requires dynamic activation quantization."
+        if not c.input_symmetric:
+            return False, "requires symmetric activation quantization."
+        if not c.is_channelwise:
+            return False, "requires per-channel weight quantization."
+        return True, None
+
+    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
+        """Prepare weights for ``zentorch_dynamic_qlinear``.
+
+        Keeps weight in [N, K] layout (int8, contiguous) and converts the
+        per-channel weight scale to bf16 with shape ``(N,)``.
+        """
+        w_q_name, w_s_name, _, _, _ = self.layer_param_names
+        weight = getattr(layer, w_q_name)
+        n = weight.shape[0]
+        replace_parameter(
+            layer,
+            w_q_name,
+            torch.nn.Parameter(weight.data.contiguous(), requires_grad=False),
+        )
+
+        weight_scale = getattr(layer, w_s_name)
+        ws = weight_scale.data
+        if ws.dim() == 2 and ws.shape[-1] == 1:
+            ws = ws.squeeze(-1)
+        ws = ws.to(torch.bfloat16).contiguous()
+        assert ws.shape == (n,), (
+            f"[zen_cpu] expected weight scale shape ({n},), got {tuple(ws.shape)}"
+        )
+
+        replace_parameter(
+            layer,
+            w_s_name,
+            torch.nn.Parameter(ws, requires_grad=False),
+        )
+        logger.info_once(
+            "[zen_cpu] Using zentorch_dynamic_qlinear for W8A8 (dynamic-symmetric)"
+        )
+
+    def apply_weights(
+        self,
+        layer: torch.nn.Module,
+        x: torch.Tensor,
+        bias: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        w_q_name, w_s_name, _, _, _ = self.layer_param_names
+        return torch.ops.zentorch.zentorch_dynamic_qlinear(
+            x,
+            getattr(layer, w_q_name),
+            getattr(layer, w_s_name),
+            bias,
+            zentorch_op_name="zentorch::zentorch_dynamic_qlinear",
+        )

vllm/model_executor/kernels/linear/zentorch_utils.py

@@ -0,0 +1,23 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Gates zentorch CPU linear dispatch on platform/op availability."""
+
+from __future__ import annotations
+
+import torch
+
+from vllm.platforms import current_platform
+
+__all__ = ["has_zentorch_op"]
+
+
+def has_zentorch_op(op_names: list[str]) -> bool:
+    """Return ``True`` when running on Zen CPU with all named ops registered."""
+    if not op_names:
+        raise ValueError("has_zentorch_op requires at least one op name")
+    if not current_platform.is_zen_cpu():
+        return False
+    ns = getattr(torch.ops, "zentorch", None)
+    if ns is None:
+        return False
+    return all(hasattr(ns, op_name) for op_name in op_names)

vllm/model_executor/kernels/mhc/__init__.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+from .aiter import *
+from .tilelang import *
+from .torch import *
+from .triton import *
+
+__all__ = [
+    "mhc_pre_cuda",
+    "mhc_post_cuda",
+    "mhc_fused_post_pre_cuda",
+    "hc_head_fused_kernel_cuda",
+    "mhc_pre_aiter",
+    "mhc_post_aiter",
+    "mhc_fused_post_pre_aiter",
+    "hc_head_fused_aiter",
+    "mhc_pre_tilelang",
+    "mhc_post_tilelang",
+    "mhc_fused_post_pre_tilelang",
+    "hc_head_fused_tilelang",
+    "mhc_pre_torch",
+    "mhc_post_torch",
+    "mhc_fused_post_pre_torch",
+    "hc_head_fused_torch",
+    "mhc_pre_triton",
+    "mhc_post_triton",
+    "mhc_fused_post_pre_triton",
+    "hc_head_fused_triton",
+]

vllm/model_executor/kernels/mhc/aiter.py

@@ -0,0 +1,138 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.utils.torch_utils import direct_register_custom_op
+
+
+def mhc_pre_aiter(
+    residual: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    n_splits: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Forward pass for mHC pre block.
+
+    Args:
+        residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
+        fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
+        hc_scale: shape (3,), dtype torch.float32
+        hc_base: shape (hc_mult3,), dtype torch.float32
+        rms_eps: RMS normalization epsilon
+        hc_pre_eps: pre-mix epsilon
+        hc_sinkhorn_eps: sinkhorn epsilon
+        hc_post_mult_value: post-mix multiplier value
+        sinkhorn_repeat: number of sinkhorn iterations
+        n_splits: split-k factor;
+
+    Returns:
+        post_mix: shape (..., hc_mult), dtype torch.float32
+        comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
+        layer_input: shape (..., hidden_size), dtype torch.bfloat16
+    """
+
+    hidden_size = residual.shape[-1]
+    assert hidden_size % 256 == 0
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    return rocm_aiter_ops.mhc_pre(
+        residual,
+        fn,
+        hc_scale,
+        hc_base,
+        rms_eps,
+        hc_pre_eps,
+        hc_sinkhorn_eps,
+        hc_post_mult_value,
+        sinkhorn_repeat,
+    )
+
+
+def _mhc_pre_aiter_fake(
+    residual: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    n_splits: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    hc_mult = residual.shape[-2]
+    hidden_size = residual.shape[-1]
+    outer_shape = residual.shape[:-2]
+
+    # Create empty tensors with correct shapes for meta device / shape inference
+    post_mix = torch.empty(
+        *outer_shape,
+        hc_mult,
+        1,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    comb_mix = torch.empty(
+        *outer_shape,
+        hc_mult,
+        hc_mult,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    layer_input = torch.empty(
+        *outer_shape,
+        hidden_size,
+        dtype=torch.bfloat16,
+        device=residual.device,
+    )
+
+    return post_mix, comb_mix, layer_input
+
+
+def mhc_post_aiter(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+) -> torch.Tensor:
+    hidden_size = residual.shape[-1]
+
+    assert hidden_size % 256 == 0
+    from vllm._aiter_ops import rocm_aiter_ops
+
+    return rocm_aiter_ops.mhc_post(
+        x,
+        residual,
+        post_layer_mix,
+        comb_res_mix,
+    )
+
+
+def _mhc_post_aiter_fake(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+) -> torch.Tensor:
+    return torch.empty_like(residual)
+
+
+direct_register_custom_op(
+    op_name="mhc_pre_aiter",
+    op_func=mhc_pre_aiter,
+    mutates_args=[],
+    fake_impl=_mhc_pre_aiter_fake,
+)
+direct_register_custom_op(
+    op_name="mhc_post_aiter",
+    op_func=mhc_post_aiter,
+    mutates_args=[],
+    fake_impl=_mhc_post_aiter_fake,
+)

vllm/model_executor/kernels/mhc/tilelang.py

@@ -0,0 +1,683 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.utils.torch_utils import direct_register_custom_op
+
+
+def _torch_hc_prenorm_gemm(
+    x: torch.Tensor,
+    fn: torch.Tensor,
+    out: torch.Tensor,
+    sqrsum: torch.Tensor,
+) -> None:
+    assert out.shape[0] == 1
+    assert sqrsum.shape[0] == 1
+    x_float = x.float()
+    out[0].copy_(x_float @ fn.t())
+    sqrsum[0].copy_(x_float.square().sum(dim=-1))
+
+
+def _tilelang_hc_prenorm_gemm(
+    x: torch.Tensor,
+    fn: torch.Tensor,
+    out: torch.Tensor,
+    sqrsum: torch.Tensor,
+    hidden_size: int,
+    hc_mult: int,
+    tile_n: int = 12,
+    n_thr: int = 512,
+    n_splits: int = 1,
+) -> None:
+    from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+        hc_prenorm_gemm_block_m_tilelang,
+        hc_prenorm_gemm_tilelang,
+    )
+
+    assert out.shape[0] == n_splits
+    assert sqrsum.shape[0] == n_splits
+    assert x.shape[1] == hc_mult * hidden_size
+    assert x.shape[1] % n_splits == 0
+    assert (x.shape[1] // n_splits) % n_thr == 0
+    use_default_config = tile_n == 12 and n_thr == 512
+    if n_splits == 1 and use_default_config and x.shape[0] >= 1024:
+        hc_prenorm_gemm_block_m_tilelang(
+            x,
+            fn,
+            out,
+            sqrsum,
+            hidden_size,
+            hc_mult,
+            fn.shape[0],
+            n_thr,
+            tile_n,
+            2,
+        )
+        return
+    if (
+        n_splits == 1
+        and use_default_config
+        and x.shape[0] < 128
+        and x.shape[1] % 1024 == 0
+    ):
+        hc_prenorm_gemm_tilelang(
+            x,
+            fn,
+            out,
+            sqrsum,
+            hidden_size,
+            hc_mult,
+            fn.shape[0],
+            1024,
+            4,
+            n_splits,
+        )
+        return
+    hc_prenorm_gemm_tilelang(
+        x,
+        fn,
+        out,
+        sqrsum,
+        hidden_size,
+        hc_mult,
+        fn.shape[0],
+        n_thr,
+        tile_n,
+        n_splits,
+    )
+
+
+def mhc_pre_tilelang(
+    residual: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    n_splits: int = 1,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Forward pass for mHC pre block.
+
+    Args:
+        residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
+        fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
+        hc_scale: shape (3,), dtype torch.float32
+        hc_base: shape (hc_mult3,), dtype torch.float32
+        rms_eps: RMS normalization epsilon
+        hc_pre_eps: pre-mix epsilon
+        hc_sinkhorn_eps: sinkhorn epsilon
+        hc_post_mult_value: post-mix multiplier value
+        sinkhorn_repeat: number of sinkhorn iterations
+        n_splits: split-k factor;
+        norm_weight: optional RMSNorm weight, shape (hidden_size,), dtype
+            torch.bfloat16. When provided, RMSNorm is fused into the
+            layer_input write path of the big_fuse kernel.
+        norm_eps: epsilon for the fused RMSNorm; only consulted when
+            norm_weight is given.
+
+    Returns:
+        post_mix: shape (..., hc_mult), dtype torch.float32
+        comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
+        layer_input: shape (..., hidden_size), dtype torch.bfloat16
+    """
+    from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+        compute_num_split,
+        mhc_pre_big_fuse_tilelang,
+        mhc_pre_big_fuse_with_norm_tilelang,
+    )
+    from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
+    from vllm.utils.math_utils import cdiv
+
+    assert residual.dtype == torch.bfloat16
+    assert fn.dtype == torch.float32
+    assert hc_scale.dtype == torch.float32
+    assert hc_base.dtype == torch.float32
+
+    hc_mult = residual.shape[-2]
+    hidden_size = residual.shape[-1]
+    hc_mult2 = hc_mult * hc_mult
+    hc_mult3 = hc_mult * 2 + hc_mult2
+
+    hc_hidden_size = hc_mult * hidden_size
+    assert fn.shape[0] == hc_mult3
+    assert fn.shape[1] == hc_hidden_size
+    assert hc_scale.shape == (3,)
+    assert hc_base.shape == (hc_mult3,)
+
+    if norm_weight is not None:
+        assert norm_weight.shape == (hidden_size,)
+        if norm_weight.dtype != torch.bfloat16:
+            norm_weight = norm_weight.to(torch.bfloat16)
+        if not norm_weight.is_contiguous():
+            norm_weight = norm_weight.contiguous()
+
+    outer_shape = residual.shape[:-2]
+
+    residual_flat = residual.view(-1, hc_mult, hidden_size)
+    num_tokens = residual_flat.shape[0]
+
+    from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+    use_deep_gemm = is_deep_gemm_supported()
+    if use_deep_gemm:
+        # these numbers are from deepgemm kernel impl
+        block_k = 64
+        block_m = 64
+        n_splits = compute_num_split(block_k, hc_hidden_size, cdiv(num_tokens, block_m))
+    else:
+        n_splits = 1
+
+    post_mix = torch.empty(
+        num_tokens, hc_mult, dtype=torch.float32, device=residual.device
+    )
+    comb_mix = torch.empty(
+        num_tokens, hc_mult2, dtype=torch.float32, device=residual.device
+    )
+    layer_input = torch.empty(
+        num_tokens, hidden_size, dtype=torch.bfloat16, device=residual.device
+    )
+
+    gemm_out_mul = torch.empty(
+        n_splits, num_tokens, hc_mult3, dtype=torch.float32, device=residual.device
+    )
+    gemm_out_sqrsum = torch.empty(
+        n_splits, num_tokens, dtype=torch.float32, device=residual.device
+    )
+
+    residual_2d = residual_flat.view(num_tokens, hc_mult * hidden_size)
+    if use_deep_gemm:
+        tf32_hc_prenorm_gemm(
+            residual_2d,
+            fn,
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            n_splits,
+        )
+    else:
+        _tilelang_hc_prenorm_gemm(
+            residual_2d,
+            fn,
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            hidden_size,
+            hc_mult,
+        )
+
+    if norm_weight is None:
+        mhc_pre_big_fuse_tilelang(
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            hc_scale,
+            hc_base,
+            residual_flat,
+            post_mix,
+            comb_mix,
+            layer_input,
+            hidden_size,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+            n_splits,
+            hc_mult,
+        )
+    else:
+        mhc_pre_big_fuse_with_norm_tilelang(
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            hc_scale,
+            hc_base,
+            residual_flat,
+            post_mix,
+            comb_mix,
+            layer_input,
+            norm_weight,
+            hidden_size,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+            norm_eps,
+            n_splits,
+            hc_mult,
+        )
+
+    return (
+        post_mix.view(*outer_shape, hc_mult, 1),
+        comb_mix.view(*outer_shape, hc_mult, hc_mult),
+        layer_input.view(*outer_shape, hidden_size),
+    )
+
+
+def _mhc_pre_tilelang_fake(
+    residual: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    n_splits: int = 1,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    hc_mult = residual.shape[-2]
+    hidden_size = residual.shape[-1]
+    outer_shape = residual.shape[:-2]
+
+    # Create empty tensors with correct shapes for meta device / shape inference
+    post_mix = torch.empty(
+        *outer_shape,
+        hc_mult,
+        1,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    comb_mix = torch.empty(
+        *outer_shape,
+        hc_mult,
+        hc_mult,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    layer_input = torch.empty(
+        *outer_shape,
+        hidden_size,
+        dtype=torch.bfloat16,
+        device=residual.device,
+    )
+
+    return post_mix, comb_mix, layer_input
+
+
+def mhc_post_tilelang(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+) -> torch.Tensor:
+    from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+        mhc_post_tilelang as _mhc_post_kernel,
+    )
+
+    out = torch.empty_like(residual)
+    _mhc_post_kernel(
+        comb_res_mix,
+        residual,
+        post_layer_mix.squeeze(-1),
+        x,
+        out,
+        residual.shape[-2],
+        residual.shape[-1],
+    )
+    return out
+
+
+def mhc_fused_post_pre_tilelang(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    n_splits: int = 1,
+    tile_n: int = 1,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Run one MHC post block followed by the next MHC pre block.
+
+    When ``norm_weight`` is provided, the layer_input_cur output is the
+    RMSNorm'd activation (fused into the kernel); otherwise it is the
+    raw pre-norm activation as before.
+
+    Returns:
+        residual_cur: post-mapped residual, shape (..., hc_mult, hidden_size)
+        post_mix_cur: shape (..., hc_mult, 1)
+        comb_mix_cur: shape (..., hc_mult, hc_mult)
+        layer_input_cur: shape (..., hidden_size)
+    """
+
+    from vllm.model_executor.kernels.mhc.tilelang_kernels import (
+        compute_num_split,
+        mhc_fused_tilelang,
+        mhc_post_tilelang,
+        mhc_pre_big_fuse_tilelang,
+        mhc_pre_big_fuse_with_norm_tilelang,
+    )
+    from vllm.utils.math_utils import cdiv
+
+    assert residual.dtype == torch.bfloat16
+    assert x.dtype == torch.bfloat16
+    assert post_layer_mix.dtype == torch.float32
+    assert comb_res_mix.dtype == torch.float32
+    assert fn.dtype == torch.float32
+    assert hc_scale.dtype == torch.float32
+    assert hc_base.dtype == torch.float32
+
+    hc_mult = residual.shape[-2]
+    hidden_size = residual.shape[-1]
+    hc_mult2 = hc_mult * hc_mult
+    hc_mult3 = hc_mult * 2 + hc_mult2
+    hc_hidden_size = hc_mult * hidden_size
+    outer_shape = residual.shape[:-2]
+
+    assert x.shape == (*outer_shape, hidden_size)
+    assert post_layer_mix.shape in (
+        (*outer_shape, hc_mult, 1),
+        (*outer_shape, hc_mult),
+    )
+    assert comb_res_mix.shape == (*outer_shape, hc_mult, hc_mult)
+    assert fn.shape == (hc_mult3, hc_hidden_size)
+    assert hc_scale.shape == (3,)
+    assert hc_base.shape == (hc_mult3,)
+
+    if norm_weight is not None:
+        assert norm_weight.shape == (hidden_size,)
+        if norm_weight.dtype != torch.bfloat16:
+            norm_weight = norm_weight.to(torch.bfloat16)
+        if not norm_weight.is_contiguous():
+            norm_weight = norm_weight.contiguous()
+
+    assert n_splits in (1, 2, 4, 8)
+    assert hidden_size % n_splits == 0
+
+    residual_flat = residual.view(-1, hc_mult, hidden_size)
+    num_tokens = residual_flat.shape[0]
+    x_flat = x.view(num_tokens, hidden_size)
+    post_layer_mix_flat = post_layer_mix.view(num_tokens, hc_mult)
+    comb_res_mix_flat = comb_res_mix.view(num_tokens, hc_mult, hc_mult)
+
+    from vllm.utils.deep_gemm import is_deep_gemm_supported
+
+    use_deep_gemm = is_deep_gemm_supported()
+    use_small_fma = num_tokens <= 16
+    if use_small_fma:
+        # TODO(gnovack): investigate autotuning these heuristics
+        tile_n = 2 if num_tokens < 8 else 3
+        n_splits = 8 if (num_tokens < 8 and hidden_size <= 4096) else 4
+    else:
+        if use_deep_gemm:
+            # these number are from deepgemm kernel impl
+            block_k = 64
+            block_m = 64
+            n_splits = compute_num_split(
+                block_k, hc_hidden_size, cdiv(num_tokens, block_m)
+            )
+        else:
+            n_splits = 1
+
+    gemm_out_mul = torch.empty(
+        n_splits,
+        num_tokens,
+        hc_mult3,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    gemm_out_sqrsum = torch.empty(
+        n_splits,
+        num_tokens,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    residual_cur = torch.empty_like(residual_flat)
+    post_mix_cur = torch.empty(
+        num_tokens,
+        hc_mult,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    comb_mix_cur = torch.empty(
+        num_tokens,
+        hc_mult2,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    layer_input_cur = torch.empty(
+        num_tokens,
+        hidden_size,
+        dtype=torch.bfloat16,
+        device=residual.device,
+    )
+
+    if use_small_fma:
+        mhc_fused_tilelang(
+            comb_res_mix_flat,
+            residual_flat,
+            post_layer_mix_flat,
+            x_flat,
+            fn.view(hc_mult3, hc_mult, hidden_size),
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            residual_cur,
+            hc_mult,
+            hidden_size,
+            hc_mult3,
+            tile_n=tile_n,
+            n_splits=n_splits,
+        )
+    else:
+        mhc_post_tilelang(
+            comb_res_mix_flat,
+            residual_flat,
+            post_layer_mix_flat,
+            x_flat,
+            residual_cur,
+            residual.shape[-2],
+            residual.shape[-1],
+        )
+
+        residual_cur_2d = residual_cur.view(num_tokens, hc_mult * hidden_size)
+        if use_deep_gemm:
+            from vllm.utils.deep_gemm import tf32_hc_prenorm_gemm
+
+            tf32_hc_prenorm_gemm(
+                residual_cur_2d,
+                fn,
+                gemm_out_mul,
+                gemm_out_sqrsum,
+                n_splits,
+            )
+        else:
+            _tilelang_hc_prenorm_gemm(
+                residual_cur_2d,
+                fn,
+                gemm_out_mul,
+                gemm_out_sqrsum,
+                hidden_size,
+                hc_mult,
+            )
+
+    if norm_weight is None:
+        mhc_pre_big_fuse_tilelang(
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            hc_scale,
+            hc_base,
+            residual_cur,
+            post_mix_cur,
+            comb_mix_cur,
+            layer_input_cur,
+            hidden_size,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+            n_splits,
+            hc_mult,
+        )
+    else:
+        mhc_pre_big_fuse_with_norm_tilelang(
+            gemm_out_mul,
+            gemm_out_sqrsum,
+            hc_scale,
+            hc_base,
+            residual_cur,
+            post_mix_cur,
+            comb_mix_cur,
+            layer_input_cur,
+            norm_weight,
+            hidden_size,
+            rms_eps,
+            hc_pre_eps,
+            hc_sinkhorn_eps,
+            hc_post_mult_value,
+            sinkhorn_repeat,
+            norm_eps,
+            n_splits,
+            hc_mult,
+        )
+
+    return (
+        residual_cur.view(*outer_shape, hc_mult, hidden_size),
+        post_mix_cur.view(*outer_shape, hc_mult, 1),
+        comb_mix_cur.view(*outer_shape, hc_mult, hc_mult),
+        layer_input_cur.view(*outer_shape, hidden_size),
+    )
+
+
+def _mhc_fused_post_pre_tilelang_fake(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    n_splits: int = 1,
+    tile_n: int = 1,
+    norm_weight: torch.Tensor | None = None,
+    norm_eps: float = 1e-6,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    hc_mult = residual.shape[-2]
+    hidden_size = residual.shape[-1]
+    outer_shape = residual.shape[:-2]
+
+    residual_cur = torch.empty_like(residual)
+    post_mix_cur = torch.empty(
+        *outer_shape,
+        hc_mult,
+        1,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    comb_mix_cur = torch.empty(
+        *outer_shape,
+        hc_mult,
+        hc_mult,
+        dtype=torch.float32,
+        device=residual.device,
+    )
+    layer_input_cur = torch.empty(
+        *outer_shape,
+        hidden_size,
+        dtype=torch.bfloat16,
+        device=residual.device,
+    )
+
+    return residual_cur, post_mix_cur, comb_mix_cur, layer_input_cur
+
+
+def _mhc_post_tilelang_fake(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+) -> torch.Tensor:
+    return torch.empty_like(residual)
+
+
+def hc_head_fused_kernel_tilelang(
+    hs_flat: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_eps: float,
+) -> torch.Tensor:
+    """Apply the fused hc_head kernel and return the (T, H) bf16 result."""
+    num_tokens, hc_mult, hidden_size = hs_flat.shape
+    out = torch.empty(
+        num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_flat.device
+    )
+    if num_tokens == 0:
+        return out
+    from vllm.model_executor.kernels.mhc.tilelang_kernels import hc_head_fuse_tilelang
+
+    hc_head_fuse_tilelang(
+        hs_flat,
+        fn,
+        hc_scale,
+        hc_base,
+        out,
+        hidden_size,
+        rms_eps,
+        hc_eps,
+        hc_mult,
+    )
+    return out
+
+
+def _hc_head_fused_kernel_tilelang_fake(
+    hs_flat: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_eps: float,
+) -> torch.Tensor:
+    num_tokens, _, hidden_size = hs_flat.shape
+    return torch.empty(
+        num_tokens, hidden_size, dtype=torch.bfloat16, device=hs_flat.device
+    )
+
+
+direct_register_custom_op(
+    op_name="mhc_pre_tilelang",
+    op_func=mhc_pre_tilelang,
+    mutates_args=[],
+    fake_impl=_mhc_pre_tilelang_fake,
+)
+direct_register_custom_op(
+    op_name="mhc_post_tilelang",
+    op_func=mhc_post_tilelang,
+    mutates_args=[],
+    fake_impl=_mhc_post_tilelang_fake,
+)
+
+direct_register_custom_op(
+    op_name="mhc_fused_post_pre_tilelang",
+    op_func=mhc_fused_post_pre_tilelang,
+    mutates_args=[],
+    fake_impl=_mhc_fused_post_pre_tilelang_fake,
+)
+
+direct_register_custom_op(
+    op_name="hc_head_fused_kernel_tilelang",
+    op_func=hc_head_fused_kernel_tilelang,
+    mutates_args=[],
+    fake_impl=_hc_head_fused_kernel_tilelang_fake,
+)

vllm/model_executor/kernels/mhc/tilelang_kernels.py

@@ -0,0 +1,811 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import math
+from functools import cache
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.platforms import current_platform
+from vllm.utils.import_utils import has_tilelang
+from vllm.utils.math_utils import cdiv
+
+# TileLang is used for MHC on CUDA and ROCm. Keep non-GPU imports cheap so
+# registering the Python wrapper modules does not require TileLang everywhere.
+if TYPE_CHECKING or current_platform.is_cuda_alike():
+    if not has_tilelang():
+        raise ImportError(
+            "tilelang is required for mhc but is not installed. Install it with "
+            "`pip install tilelang`."
+        )
+    import tilelang
+    import tilelang.language as T
+else:
+    tilelang = None  # type: ignore[assignment]
+    T = None  # type: ignore[assignment]
+
+ENABLE_PDL = current_platform.is_arch_support_pdl() and current_platform.is_cuda()
+
+
+@cache
+def compute_num_split(block_k: int, k: int | None, grid_size: int) -> int:
+    device_props = torch.cuda.get_device_properties(0)
+    n_sms = device_props.multi_processor_count
+    split_k = n_sms // grid_size
+    if k is not None:
+        # avoid split_k for small k
+        num_block_k = cdiv(k, block_k)
+        split_k = min(split_k, num_block_k // 4)
+    split_k = max(split_k, 1)
+    return split_k
+
+
+pass_configs: dict[tilelang.PassConfigKey, Any] = {
+    tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
+    tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
+}
+
+if current_platform.is_cuda():
+    pass_configs[tilelang.PassConfigKey.TL_PTXAS_REGISTER_USAGE_LEVEL] = 10
+
+
+@tilelang.jit(
+    pass_configs=pass_configs,
+)
+def mhc_pre_big_fuse_tilelang(
+    gemm_out_mul,
+    gemm_out_sqrsum,
+    hc_scale,
+    hc_base,
+    residual,
+    post_mix,
+    comb_mix,
+    layer_input,
+    hidden_size: int,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    n_splits: int = 16,
+    hc_mult: int = 4,
+):
+    """Deeply fused kernels, everything other than gemm & sqrsum in mHC pre block."""
+    num_tokens = T.dynamic("num_tokens")
+    hc_mult3 = hc_mult * (2 + hc_mult)
+    hidden_block = math.gcd(512, hidden_size)
+
+    gemm_out_mul: T.Tensor[[n_splits, num_tokens, hc_mult3], T.float32]  # type: ignore[no-redef, valid-type]
+    gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32]  # type: ignore[no-redef, valid-type]
+    hc_scale: T.Tensor[[3], T.float32]  # type: ignore[no-redef, valid-type]
+    hc_base: T.Tensor[[hc_mult3], T.float32]  # type: ignore[no-redef, valid-type]
+    residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
+    # outputs
+    post_mix: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
+    comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
+    layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
+
+    with T.Kernel(num_tokens, threads=96) as i:
+        if ENABLE_PDL:
+            T.pdl_sync()
+        ##################################################################
+        # _pre_norm_fn_fwd_norm
+        rms = T.alloc_fragment(1, T.float32)
+        mixes = T.alloc_fragment(hc_mult3, T.float32)
+        T.clear(mixes)
+        rms[0] = 0
+        for i_split in T.serial(n_splits):
+            rms[0] += gemm_out_sqrsum[i_split, i]
+        rms[0] = T.rsqrt(rms[0] / (hc_mult * hidden_size) + rms_eps)
+        for j in T.Parallel(hc_mult3):
+            mixes[j] = 0
+            for i_split in T.serial(n_splits):
+                mixes[j] += gemm_out_mul[i_split, i, j]
+            mixes[j] *= rms[0]
+        mixes_shared = T.alloc_shared(hc_mult3, T.float32)
+        T.copy(mixes, mixes_shared)
+
+        if T.get_thread_binding() < 32:
+            ##################################################################
+            # _pre_split_mixes_fwd (post & comb)
+            cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
+            for j in T.Parallel(hc_mult):
+                post_mix[i, j] = (
+                    T.sigmoid(
+                        mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
+                    )
+                    * hc_post_mult_value
+                )
+            for j, k in T.Parallel(hc_mult, hc_mult):
+                cm[j, k] = (
+                    mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2]
+                    + hc_base[j * hc_mult + k + hc_mult * 2]
+                )
+
+            ##################################################################
+            # _sinkhorn_fwd
+            row_sum = T.alloc_fragment(hc_mult, T.float32)
+            col_sum = T.alloc_fragment(hc_mult, T.float32)
+
+            # comb = comb.softmax(-1) + eps
+            row_max = T.alloc_fragment(hc_mult, T.float32)
+            T.reduce_max(cm, row_max, dim=1)
+            for j, k in T.Parallel(hc_mult, hc_mult):
+                cm[j, k] = T.exp(cm[j, k] - row_max[j])
+            T.reduce_sum(cm, row_sum, dim=1)
+            for j, k in T.Parallel(hc_mult, hc_mult):
+                cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps
+
+            # comb = comb / (comb.sum(-2) + eps)
+            T.reduce_sum(cm, col_sum, dim=0)
+            for j, k in T.Parallel(hc_mult, hc_mult):
+                cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
+
+            for _ in T.serial(sinkhorn_repeat - 1):
+                # comb = comb / (comb.sum(-1) + eps)
+                T.reduce_sum(cm, row_sum, dim=1)
+                for j, k in T.Parallel(hc_mult, hc_mult):
+                    cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps)
+
+                # comb = comb / (comb.sum(-2) + eps)
+                T.reduce_sum(cm, col_sum, dim=0)
+                for j, k in T.Parallel(hc_mult, hc_mult):
+                    cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
+
+            # save comb_mix to global memory
+            for j, k in T.Parallel(hc_mult, hc_mult):
+                comb_mix[i, j * hc_mult + k] = cm[j, k]
+        else:
+            ##################################################################
+            # _pre_split_mixes_fwd (pre)
+            pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
+            for j in T.Parallel(hc_mult):
+                pre_mix_shared[j] = (
+                    T.sigmoid(
+                        mixes_shared[j] * hc_scale[0] + hc_base[j],
+                    )
+                    + hc_pre_eps
+                )
+            ###################################################################
+            # _pre_apply_mix_fwd
+            for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
+                xs = T.alloc_shared((hc_mult, hidden_block), T.float32)
+                xl = T.alloc_fragment((hc_mult, hidden_block), T.float32)
+                T.copy(residual[i, 0, i0_h * hidden_block], xs)
+                T.copy(xs, xl)
+
+                ol = T.alloc_fragment(hidden_block, T.float32)
+                T.clear(ol)
+
+                for i_hc in T.serial(hc_mult):
+                    pre = pre_mix_shared[i_hc]
+                    for i1_h in T.Parallel(hidden_block):
+                        ol[i1_h] += pre * xl[i_hc, i1_h]
+
+                T.copy(ol, layer_input[i, i0_h * hidden_block])
+
+        if ENABLE_PDL:
+            T.pdl_trigger()
+
+
+# Copied from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/mhc.py#L478
+
+
+@tilelang.jit(
+    pass_configs=pass_configs,
+)
+def mhc_pre_big_fuse_with_norm_tilelang(
+    gemm_out_mul,
+    gemm_out_sqrsum,
+    hc_scale,
+    hc_base,
+    residual,
+    post_mix,
+    comb_mix,
+    layer_input,
+    norm_weight,
+    hidden_size: int,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    norm_eps: float,
+    n_splits: int = 16,
+    hc_mult: int = 4,
+    gemm_last_dim: int = -1,
+):
+    num_tokens = T.dynamic("num_tokens")
+    hc_mult3 = hc_mult * (2 + hc_mult)
+    if gemm_last_dim < 0:
+        gemm_last_dim = hc_mult3
+    hidden_block = math.gcd(1024, hidden_size)
+
+    gemm_out_mul: T.Tensor[[n_splits, num_tokens, gemm_last_dim], T.float32]  # type: ignore[no-redef, valid-type]
+    gemm_out_sqrsum: T.Tensor[[n_splits, num_tokens], T.float32]  # type: ignore[no-redef, valid-type]
+    hc_scale: T.Tensor[[3], T.float32]  # type: ignore[no-redef, valid-type]
+    hc_base: T.Tensor[[hc_mult3], T.float32]  # type: ignore[no-redef, valid-type]
+    residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
+    post_mix: T.Tensor[[num_tokens, hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
+    comb_mix: T.Tensor[[num_tokens, hc_mult * hc_mult], T.float32]  # type: ignore[no-redef, valid-type]
+    layer_input: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
+    norm_weight: T.Tensor[[hidden_size], T.bfloat16]  # type: ignore[no-redef, valid-type]
+
+    with T.Kernel(num_tokens, threads=96) as i:
+        rms = T.alloc_fragment(1, T.float32)
+        mixes = T.alloc_fragment(hc_mult3, T.float32)
+        T.clear(mixes)
+        rms[0] = 0
+
+        if ENABLE_PDL:
+            T.pdl_sync()
+
+        for i_split in T.serial(n_splits):
+            rms[0] += gemm_out_sqrsum[i_split, i]
+        rms[0] = T.rsqrt(rms[0] / (hc_mult * hidden_size) + rms_eps)
+        for j in T.Parallel(hc_mult3):
+            mixes[j] = 0
+            for i_split in T.serial(n_splits):
+                mixes[j] += gemm_out_mul[i_split, i, j]
+            mixes[j] *= rms[0]
+        mixes_shared = T.alloc_shared(hc_mult3, T.float32)
+        T.copy(mixes, mixes_shared)
+
+        if T.get_thread_binding() < 32:
+            cm = T.alloc_fragment((hc_mult, hc_mult), T.float32)
+            for j in T.Parallel(hc_mult):
+                post_mix[i, j] = (
+                    T.sigmoid(
+                        mixes_shared[j + hc_mult] * hc_scale[1] + hc_base[j + hc_mult]
+                    )
+                    * hc_post_mult_value
+                )
+            for j, k in T.Parallel(hc_mult, hc_mult):
+                cm[j, k] = (
+                    mixes_shared[j * hc_mult + k + hc_mult * 2] * hc_scale[2]
+                    + hc_base[j * hc_mult + k + hc_mult * 2]
+                )
+
+            row_sum = T.alloc_fragment(hc_mult, T.float32)
+            col_sum = T.alloc_fragment(hc_mult, T.float32)
+
+            row_max = T.alloc_fragment(hc_mult, T.float32)
+            T.reduce_max(cm, row_max, dim=1)
+            for j, k in T.Parallel(hc_mult, hc_mult):
+                cm[j, k] = T.exp(cm[j, k] - row_max[j])
+            T.reduce_sum(cm, row_sum, dim=1)
+            for j, k in T.Parallel(hc_mult, hc_mult):
+                cm[j, k] = cm[j, k] / row_sum[j] + hc_sinkhorn_eps
+
+            T.reduce_sum(cm, col_sum, dim=0)
+            for j, k in T.Parallel(hc_mult, hc_mult):
+                cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
+
+            for _ in T.serial(sinkhorn_repeat - 1):
+                T.reduce_sum(cm, row_sum, dim=1)
+                for j, k in T.Parallel(hc_mult, hc_mult):
+                    cm[j, k] = cm[j, k] / (row_sum[j] + hc_sinkhorn_eps)
+
+                T.reduce_sum(cm, col_sum, dim=0)
+                for j, k in T.Parallel(hc_mult, hc_mult):
+                    cm[j, k] = cm[j, k] / (col_sum[k] + hc_sinkhorn_eps)
+
+            for j, k in T.Parallel(hc_mult, hc_mult):
+                comb_mix[i, j * hc_mult + k] = cm[j, k]
+        else:
+            pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
+            for j in T.Parallel(hc_mult):
+                pre_mix_shared[j] = (
+                    T.sigmoid(
+                        mixes_shared[j] * hc_scale[0] + hc_base[j],
+                    )
+                    + hc_pre_eps
+                )
+
+            # Pass 1: stash unnormalized weighted-sum output in shared memory
+            # as bf16 (matches the rounding that RMSNorm would see) while
+            # accumulating the per-position squared sum.
+            output_shared = T.alloc_shared(hidden_size, T.bfloat16)
+            sumsq_per_pos = T.alloc_fragment(hidden_block, T.float32)
+            T.clear(sumsq_per_pos)
+
+            for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
+                xs = T.alloc_shared((hc_mult, hidden_block), T.bfloat16)
+                xl = T.alloc_fragment((hc_mult, hidden_block), T.float32)
+                T.copy(residual[i, 0, i0_h * hidden_block], xs)
+                T.copy(xs, xl)
+
+                ol = T.alloc_fragment(hidden_block, T.float32)
+                T.clear(ol)
+
+                for i_hc in T.serial(hc_mult):
+                    pre = pre_mix_shared[i_hc]
+                    for i1_h in T.Parallel(hidden_block):
+                        ol[i1_h] += pre * xl[i_hc, i1_h]
+
+                for i1_h in T.Parallel(hidden_block):
+                    sumsq_per_pos[i1_h] += ol[i1_h] * ol[i1_h]
+                    output_shared[i0_h * hidden_block + i1_h] = T.bfloat16(ol[i1_h])
+
+            sumsq = T.alloc_fragment(1, T.float32)
+            T.reduce_sum(sumsq_per_pos, sumsq, dim=0)
+            rsqrt_norm = T.alloc_fragment(1, T.float32)
+            rsqrt_norm[0] = T.rsqrt(sumsq[0] / hidden_size + norm_eps)
+
+            # Pass 2: scale by rsqrt * norm_weight and write the result to HBM.
+            for i0_h in T.Pipelined(hidden_size // hidden_block, num_stages=2):
+                w_shared = T.alloc_shared(hidden_block, T.bfloat16)
+                w_local = T.alloc_fragment(hidden_block, T.float32)
+                T.copy(norm_weight[i0_h * hidden_block], w_shared)
+                T.copy(w_shared, w_local)
+
+                ol = T.alloc_fragment(hidden_block, T.float32)
+                for i1_h in T.Parallel(hidden_block):
+                    ol[i1_h] = (
+                        output_shared[i0_h * hidden_block + i1_h]
+                        * rsqrt_norm[0]
+                        * w_local[i1_h]
+                    )
+
+                T.copy(ol, layer_input[i, i0_h * hidden_block])
+
+        if ENABLE_PDL:
+            T.pdl_trigger()
+
+
+@tilelang.jit(
+    pass_configs=pass_configs,
+)
+def mhc_fused_tilelang(
+    comb_mix,
+    residual_in,
+    post_mix,
+    x_in,
+    weight_t,
+    yp_out,
+    rp_out,
+    residual_out,
+    hc: int,
+    hidden: int,
+    n_out: int,
+    n_thr: int = 256,
+    h_blk: int = 256,
+    tile_n: int = 1,
+    split_k: int = 1,
+) -> tilelang.JITKernel:
+    """Fused mhc post-mapping + pre-norm GEMM FMA"""
+    m = T.dynamic("num_tokens")
+    split_k = T.dynamic("split_k")
+    h = hidden
+    h_blk = math.gcd(hidden, h_blk)
+    h_per_split = h // split_k
+    n_tiles = n_out // tile_n
+
+    comb_mix: T.Tensor((m, hc, hc), T.float32)  # type: ignore[no-redef, valid-type]
+    residual_in: T.Tensor((m, hc, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
+    post_mix: T.Tensor((m, hc), T.float32)  # type: ignore[no-redef, valid-type]
+    x_in: T.Tensor((m, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
+    weight_t: T.Tensor((n_out, hc, h), T.float32)  # type: ignore[no-redef, valid-type]
+    yp_out: T.Tensor((split_k, m, n_out), T.float32)  # type: ignore[no-redef, valid-type]
+    rp_out: T.Tensor((split_k, m), T.float32)  # type: ignore[no-redef, valid-type]
+    residual_out: T.Tensor((m, hc, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
+
+    h_iters = h_per_split // n_thr
+    num_warps = n_thr // 32
+
+    with T.Kernel(m, n_tiles, split_k, threads=n_thr) as (i_n, i_nt, i_ks):
+        tid = T.get_thread_binding()
+        warp_id = tid // 32
+        lane = tid % 32
+
+        s_warp = T.alloc_shared((num_warps, tile_n + 1), T.float32)
+        s_post = T.alloc_shared((hc,), T.float32)
+        s_comb = T.alloc_shared((hc, hc), T.float32)
+
+        pm = T.alloc_local((hc,), T.float32)
+        cm = T.alloc_local((hc, hc), T.float32)
+        acc = T.alloc_local((tile_n,), T.float32)
+        sqr = T.alloc_local((1,), T.float32)
+        new_r = T.alloc_local((hc,), T.float32)
+
+        T.clear(acc)
+        T.clear(sqr)
+        h_split_start = i_ks * h_per_split
+
+        if ENABLE_PDL:
+            T.pdl_sync()
+
+        T.copy(post_mix[i_n, 0], s_post)
+        T.copy(comb_mix[i_n, 0, 0], s_comb)
+
+        for j in T.unroll(hc):
+            pm[j] = s_post[j]
+        for j in T.unroll(hc):
+            for k in T.unroll(hc):
+                cm[k, j] = s_comb[k, j]
+
+        # Each thread owns h_iters elements of the k-split's h slice.
+        for it in T.serial(h_iters):
+            h_idx = h_split_start + it * n_thr + tid
+
+            # Compute new residual from layer output and past residual
+            for j in T.unroll(hc):
+                new_r[j] = pm[j] * x_in[i_n, h_idx]
+                for k in T.unroll(hc):
+                    new_r[j] += cm[k, j] * residual_in[i_n, k, h_idx]
+
+            # populate residual_out and compute sqr sum
+            if i_nt == 0:
+                for j in T.unroll(hc):
+                    residual_out[i_n, j, h_idx] = new_r[j]
+                    sqr[0] += new_r[j] * new_r[j]
+
+            # Per-thread FMA into acc[n]
+            for n in T.unroll(tile_n):
+                for j in T.unroll(hc):
+                    acc[n] += weight_t[i_nt * tile_n + n, j, h_idx] * new_r[j]
+
+        for n in T.unroll(tile_n):
+            acc[n] = T.warp_reduce_sum(acc[n])
+        if i_nt == 0:
+            sqr[0] = T.warp_reduce_sum(sqr[0])
+
+        # Cross-warp reduce via shared mem
+        if lane == 0:
+            for n in T.unroll(tile_n):
+                s_warp[warp_id, n] = acc[n]
+            if i_nt == 0:
+                s_warp[warp_id, tile_n] = sqr[0]
+        T.sync_threads()
+
+        # Warp 0 does the final cross-warp sum and writes outputs
+        if warp_id == 0:
+            if lane < tile_n:
+                v = T.alloc_var(T.float32, init=0.0)
+                for w in T.unroll(num_warps):
+                    v += s_warp[w, lane]
+                yp_out[i_ks, i_n, i_nt * tile_n + lane] = v
+
+            if i_nt == 0 and lane == 0:
+                v2 = T.alloc_var(T.float32, init=0.0)
+                for w in T.unroll(num_warps):
+                    v2 += s_warp[w, tile_n]
+                rp_out[i_ks, i_n] = v2
+
+        if ENABLE_PDL:
+            T.pdl_trigger()
+
+
+@tilelang.jit(
+    pass_configs=pass_configs,
+)
+def mhc_post_tilelang(
+    a,
+    b,
+    c,
+    d,
+    x,
+    hc: int,
+    hidden: int,
+    n_thr: int = 128,
+    h_blk: int = 1024,
+) -> tilelang.JITKernel:
+    # rename for shorter code
+    n = T.dynamic("num_tokens")
+    h = hidden
+
+    h_blk = math.gcd(hidden, h_blk)
+    a: T.Tensor((n, hc, hc), T.float32)  # type: ignore[no-redef, valid-type]
+    b: T.Tensor((n, hc, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
+    c: T.Tensor((n, hc), T.float32)  # type: ignore[no-redef, valid-type]
+    d: T.Tensor((n, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
+    x: T.Tensor((n, hc, h), T.bfloat16)  # type: ignore[no-redef, valid-type]
+    with T.Kernel(n, threads=n_thr) as i_n:
+        b_shared = T.alloc_shared((hc, h_blk), T.bfloat16)
+        d_shared = T.alloc_shared(h_blk, T.bfloat16)
+
+        x_local = T.alloc_fragment((hc, h_blk), T.float32)
+        b_local = T.alloc_fragment((hc, h_blk), T.float32)
+        d_local = T.alloc_fragment(h_blk, T.float32)
+
+        a_local = T.alloc_fragment((hc, hc), T.float32)
+        c_local = T.alloc_fragment(hc, T.float32)
+        if ENABLE_PDL:
+            T.pdl_sync()
+        T.copy(a[i_n, 0, 0], a_local)
+        T.copy(c[i_n, 0], c_local)
+
+        for i0_h in T.Serial(T.ceildiv(h, h_blk)):
+            T.copy(b[i_n, 0, i0_h * h_blk], b_shared)
+            T.copy(d[i_n, i0_h * h_blk], d_shared)
+
+            T.copy(b_shared, b_local)
+            T.copy(d_shared, d_local)
+            for i_hco, i1_h in T.Parallel(hc, h_blk):
+                x_local[i_hco, i1_h] = c_local[i_hco] * d_local[i1_h]
+                for i_hci in T.vectorized(hc):
+                    x_local[i_hco, i1_h] += a_local[i_hci, i_hco] * b_local[i_hci, i1_h]
+
+            T.copy(x_local, x[i_n, 0, i0_h * h_blk])
+        if ENABLE_PDL:
+            T.pdl_trigger()
+
+
+@tilelang.jit(
+    pass_configs=pass_configs,
+)
+def hc_prenorm_gemm_tilelang(
+    x,
+    fn,
+    out,
+    sqrsum,
+    hidden_size: int,
+    hc_mult: int = 4,
+    n_out: int = 24,
+    n_thr: int = 512,
+    tile_n: int = 12,
+    n_splits: int = 1,
+) -> tilelang.JITKernel:
+    num_tokens = T.dynamic("num_tokens")
+    hc_hidden_size = hc_mult * hidden_size
+    k_per_split = hc_hidden_size // n_splits
+    k_iters = k_per_split // n_thr
+    n_tiles = T.ceildiv(n_out, tile_n)
+
+    x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16)  # type: ignore[no-redef, valid-type]
+    fn: T.Tensor((n_out, hc_hidden_size), T.float32)  # type: ignore[no-redef, valid-type]
+    out: T.Tensor((n_splits, num_tokens, n_out), T.float32)  # type: ignore[no-redef, valid-type]
+    sqrsum: T.Tensor((n_splits, num_tokens), T.float32)  # type: ignore[no-redef, valid-type]
+
+    with T.Kernel(num_tokens, n_tiles, n_splits, threads=n_thr) as (
+        i_n,
+        i_t,
+        i_s,
+    ):
+        tid = T.get_thread_binding()
+        acc = T.alloc_local((tile_n,), T.float32)
+        sqr = T.alloc_local((1,), T.float32)
+        T.clear(acc)
+        T.clear(sqr)
+
+        if ENABLE_PDL:
+            T.pdl_sync()
+
+        for it in T.serial(k_iters):
+            i_k = i_s * k_per_split + it * n_thr + tid
+            x_val = x[i_n, i_k]
+            for i_o in T.unroll(tile_n):
+                out_idx = i_t * tile_n + i_o
+                if out_idx < n_out:
+                    acc[i_o] += x_val * fn[out_idx, i_k]
+            if i_t == 0:
+                sqr[0] += x_val * x_val
+
+        for i_o in T.unroll(tile_n):
+            acc[i_o] = T.warp_reduce_sum(acc[i_o])
+        if i_t == 0:
+            sqr[0] = T.warp_reduce_sum(sqr[0])
+
+        lane = tid % 32
+        warp_id = tid // 32
+        num_warps = n_thr // 32
+        warp_acc = T.alloc_shared((num_warps, tile_n), T.float32)
+        warp_sqr = T.alloc_shared(num_warps, T.float32)
+
+        if lane == 0:
+            for i_o in T.unroll(tile_n):
+                warp_acc[warp_id, i_o] = acc[i_o]
+            if i_t == 0:
+                warp_sqr[warp_id] = sqr[0]
+        T.sync_threads()
+
+        if warp_id == 0:
+            if lane < tile_n:
+                reduced_acc = T.alloc_var(T.float32, init=0.0)
+                for i_w in T.unroll(num_warps):
+                    reduced_acc += warp_acc[i_w, lane]
+                out_idx = i_t * tile_n + lane
+                if out_idx < n_out:
+                    out[i_s, i_n, out_idx] = reduced_acc
+            if lane == 0 and i_t == 0:
+                reduced_sqr = T.alloc_var(T.float32, init=0.0)
+                for i_w in T.unroll(num_warps):
+                    reduced_sqr += warp_sqr[i_w]
+                sqrsum[i_s, i_n] = reduced_sqr
+
+        if ENABLE_PDL:
+            T.pdl_trigger()
+
+
+@tilelang.jit(
+    pass_configs=pass_configs,
+)
+def hc_prenorm_gemm_block_m_tilelang(
+    x,
+    fn,
+    out,
+    sqrsum,
+    hidden_size: int,
+    hc_mult: int = 4,
+    n_out: int = 24,
+    n_thr: int = 512,
+    tile_n: int = 12,
+    block_m: int = 2,
+) -> tilelang.JITKernel:
+    num_tokens = T.dynamic("num_tokens")
+    hc_hidden_size = hc_mult * hidden_size
+    k_iters = hc_hidden_size // n_thr
+    n_tiles = T.ceildiv(n_out, tile_n)
+    m_tiles = T.ceildiv(num_tokens, block_m)
+
+    x: T.Tensor((num_tokens, hc_hidden_size), T.bfloat16)  # type: ignore[no-redef, valid-type]
+    fn: T.Tensor((n_out, hc_hidden_size), T.float32)  # type: ignore[no-redef, valid-type]
+    out: T.Tensor((1, num_tokens, n_out), T.float32)  # type: ignore[no-redef, valid-type]
+    sqrsum: T.Tensor((1, num_tokens), T.float32)  # type: ignore[no-redef, valid-type]
+
+    with T.Kernel(m_tiles, n_tiles, threads=n_thr) as (i_mt, i_t):
+        tid = T.get_thread_binding()
+        acc = T.alloc_local((block_m, tile_n), T.float32)
+        sqr = T.alloc_local((block_m,), T.float32)
+        T.clear(acc)
+        T.clear(sqr)
+
+        if ENABLE_PDL:
+            T.pdl_sync()
+
+        for it in T.serial(k_iters):
+            i_k = it * n_thr + tid
+            fn_val = T.alloc_local((tile_n,), T.float32)
+            for i_o in T.unroll(tile_n):
+                out_idx = i_t * tile_n + i_o
+                if out_idx < n_out:
+                    fn_val[i_o] = fn[out_idx, i_k]
+                else:
+                    fn_val[i_o] = 0.0
+            for i_m in T.unroll(block_m):
+                token_idx = i_mt * block_m + i_m
+                if token_idx < num_tokens:
+                    x_val = x[token_idx, i_k]
+                    for i_o in T.unroll(tile_n):
+                        acc[i_m, i_o] += x_val * fn_val[i_o]
+                    if i_t == 0:
+                        sqr[i_m] += x_val * x_val
+
+        for i_m in T.unroll(block_m):
+            for i_o in T.unroll(tile_n):
+                acc[i_m, i_o] = T.warp_reduce_sum(acc[i_m, i_o])
+            if i_t == 0:
+                sqr[i_m] = T.warp_reduce_sum(sqr[i_m])
+
+        lane = tid % 32
+        warp_id = tid // 32
+        num_warps = n_thr // 32
+        warp_acc = T.alloc_shared((num_warps, block_m, tile_n), T.float32)
+        warp_sqr = T.alloc_shared((num_warps, block_m), T.float32)
+
+        if lane == 0:
+            for i_m in T.unroll(block_m):
+                for i_o in T.unroll(tile_n):
+                    warp_acc[warp_id, i_m, i_o] = acc[i_m, i_o]
+                if i_t == 0:
+                    warp_sqr[warp_id, i_m] = sqr[i_m]
+        T.sync_threads()
+
+        if warp_id == 0:
+            for i_m in T.unroll(block_m):
+                token_idx = i_mt * block_m + i_m
+                if token_idx < num_tokens:
+                    if lane < tile_n:
+                        reduced_acc = T.alloc_var(T.float32, init=0.0)
+                        for i_w in T.unroll(num_warps):
+                            reduced_acc += warp_acc[i_w, i_m, lane]
+                        out_idx = i_t * tile_n + lane
+                        if out_idx < n_out:
+                            out[0, token_idx, out_idx] = reduced_acc
+                    if lane == 0 and i_t == 0:
+                        reduced_sqr = T.alloc_var(T.float32, init=0.0)
+                        for i_w in T.unroll(num_warps):
+                            reduced_sqr += warp_sqr[i_w, i_m]
+                        sqrsum[0, token_idx] = reduced_sqr
+
+        if ENABLE_PDL:
+            T.pdl_trigger()
+
+
+@tilelang.jit(
+    pass_configs=pass_configs,
+)
+def hc_head_fuse_tilelang(
+    residual,
+    fn,
+    hc_scale,
+    hc_base,
+    out,
+    hidden_size: int,
+    rms_eps: float,
+    hc_eps: float,
+    hc_mult: int = 4,
+    n_thr: int = 128,
+    h_blk: int = 1024,
+):
+    """Two-pass fused kernel for hc_head.
+
+    Pass 1: accumulate per-token squared sum and hc_mult dot-products
+            (projections onto fn rows) using cross-thread reducers.
+    Pass 2: apply sigmoid-gated weighted sum of residual channels to output.
+
+    Avoids materialising mixes / rsqrt / pre tensors to global memory.
+    """
+    num_tokens = T.dynamic("num_tokens")
+    hc_dim = hc_mult * hidden_size
+    h_block = math.gcd(h_blk, hidden_size)
+    n_h = hidden_size // h_block
+
+    residual: T.Tensor[[num_tokens, hc_mult, hidden_size], T.bfloat16]  # type: ignore[no-redef,valid-type]
+    fn: T.Tensor[[hc_mult, hc_dim], T.float32]  # type: ignore[no-redef,valid-type]
+    hc_scale: T.Tensor[[1], T.float32]  # type: ignore[no-redef,valid-type]
+    hc_base: T.Tensor[[hc_mult], T.float32]  # type: ignore[no-redef,valid-type]
+    out: T.Tensor[[num_tokens, hidden_size], T.bfloat16]  # type: ignore[no-redef,valid-type]
+
+    with T.Kernel(num_tokens, threads=n_thr) as i:
+        if ENABLE_PDL:
+            T.pdl_sync()
+
+        # ------------------------------------------------------------------
+        # Pass 1 – for each residual channel m_c and h_block:
+        #   • accumulate squared sum (for RMS norm denominator)
+        #   • accumulate hc_mult dot-products with fn rows
+        # ------------------------------------------------------------------
+        sqrsum_r = T.alloc_reducer((1,), T.float32, replication="all")
+        mixes_r = T.alloc_reducer((hc_mult,), T.float32, replication="all")
+        T.fill(sqrsum_r, 0.0)
+        T.fill(mixes_r, 0.0)
+
+        for m_c in T.serial(hc_mult):
+            for i_h in T.serial(n_h):
+                x_local = T.alloc_fragment(h_block, T.float32)
+                T.copy(residual[i, m_c, i_h * h_block], x_local)
+
+                for k in T.Parallel(h_block):
+                    sqrsum_r[0] += x_local[k] * x_local[k]
+
+                for m_m in T.unroll(hc_mult):
+                    fn_local = T.alloc_fragment(h_block, T.float32)
+                    T.copy(fn[m_m, m_c * hidden_size + i_h * h_block], fn_local)
+                    for k in T.Parallel(h_block):
+                        mixes_r[m_m] += x_local[k] * fn_local[k]
+
+        T.finalize_reducer(sqrsum_r)
+        T.finalize_reducer(mixes_r)
+
+        # ------------------------------------------------------------------
+        # Compute pre_mix = sigmoid(mix * rsqrt * scale + base) + eps
+        # ------------------------------------------------------------------
+        pre_mix_shared = T.alloc_shared(hc_mult, T.float32)
+        rsqrt_val = T.alloc_fragment(1, T.float32)
+        rsqrt_val[0] = T.rsqrt(sqrsum_r[0] / hc_dim + rms_eps)
+        for m in T.Parallel(hc_mult):
+            pre_mix_shared[m] = (
+                T.sigmoid(mixes_r[m] * rsqrt_val[0] * hc_scale[0] + hc_base[m]) + hc_eps
+            )
+
+        # ------------------------------------------------------------------
+        # Pass 2 – apply_mix: pipelined weighted sum over residual channels
+        # ------------------------------------------------------------------
+        for i0_h in T.Pipelined(n_h, num_stages=2):
+            xs = T.alloc_shared((hc_mult, h_block), T.bfloat16)
+            xl = T.alloc_fragment((hc_mult, h_block), T.float32)
+            T.copy(residual[i, 0, i0_h * h_block], xs, disable_tma=True)
+            T.copy(xs, xl)
+
+            ol = T.alloc_fragment(h_block, T.float32)
+            T.clear(ol)
+            for i_hc in T.serial(hc_mult):
+                pre = pre_mix_shared[i_hc]
+                for i1_h in T.Parallel(h_block):
+                    ol[i1_h] += pre * xl[i_hc, i1_h]
+
+            T.copy(ol, out[i, i0_h * h_block], disable_tma=True)
+
+        if ENABLE_PDL:
+            T.pdl_trigger()

vllm/model_executor/kernels/mhc/torch.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+
+def mhc_pre_torch(
+    residual: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    rms_eps: float,
+    hc_pre_eps: float,
+    hc_sinkhorn_eps: float,
+    hc_post_mult_value: float,
+    sinkhorn_repeat: int,
+    n_splits: int = 1,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """
+    Forward pass for mHC pre block.
+
+    Args:
+        residual: shape (..., hc_mult, hidden_size), dtype torch.bfloat16
+        fn: shape (hc_mult3, hc_mult * hidden_size), dtype torch.float32
+        hc_scale: shape (3,), dtype torch.float32
+        hc_base: shape (hc_mult3,), dtype torch.float32
+        rms_eps: RMS normalization epsilon
+        hc_pre_eps: pre-mix epsilon
+        hc_sinkhorn_eps: sinkhorn epsilon
+        hc_post_mult_value: post-mix multiplier value
+        sinkhorn_repeat: number of sinkhorn iterations
+        n_splits: split-k factor;
+
+    Returns:
+        post_mix: shape (..., hc_mult), dtype torch.float32
+        comb_mix: shape (..., hc_mult, hc_mult), dtype torch.float32
+        layer_input: shape (..., hidden_size), dtype torch.bfloat16
+    """
+
+    # Validate shapes
+    assert residual.dtype == torch.bfloat16
+    assert fn.dtype == torch.float32
+    assert hc_scale.dtype == torch.float32
+    assert hc_base.dtype == torch.float32
+
+    hc_mult = residual.shape[-2]
+    hidden_size = residual.shape[-1]
+    hc_mult2 = hc_mult * hc_mult
+    hc_mult3 = hc_mult * 2 + hc_mult2
+
+    hc_hidden_size = hc_mult * hidden_size
+    assert fn.shape[0] == hc_mult3
+    assert fn.shape[1] == hc_hidden_size
+    assert hc_scale.shape == (3,)
+    assert hc_base.shape == (hc_mult3,)
+
+    outer_shape = residual.shape[:-2]
+
+    residual_flat = residual.view(-1, hc_mult, hidden_size)
+    num_tokens = residual_flat.shape[0]
+    fn_flat = fn
+
+    x = residual_flat.view(num_tokens, hc_mult * hidden_size).to(torch.float32)
+    mixes = torch.matmul(x, fn_flat.t())
+    sqrsum = x.square().sum(dim=-1, keepdim=True)
+    mixes = mixes * torch.rsqrt(sqrsum / (hc_mult * hidden_size) + rms_eps)
+
+    pre_logits = mixes[:, :hc_mult] * hc_scale[0] + hc_base[:hc_mult]
+    pre_mix = torch.sigmoid(pre_logits) + hc_pre_eps
+
+    post_logits = (
+        mixes[:, hc_mult : 2 * hc_mult] * hc_scale[1] + hc_base[hc_mult : 2 * hc_mult]
+    )
+    post_mix = torch.sigmoid(post_logits) * hc_post_mult_value
+
+    comb_logits = mixes[:, 2 * hc_mult :].view(num_tokens, hc_mult, hc_mult) * hc_scale[
+        2
+    ] + hc_base[2 * hc_mult :].view(1, hc_mult, hc_mult)
+    comb_mix = torch.softmax(comb_logits, dim=-1) + hc_sinkhorn_eps
+    comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
+    for _ in range(sinkhorn_repeat - 1):
+        comb_mix = comb_mix / (comb_mix.sum(dim=-1, keepdim=True) + hc_sinkhorn_eps)
+        comb_mix = comb_mix / (comb_mix.sum(dim=-2, keepdim=True) + hc_sinkhorn_eps)
+
+    layer_input = torch.sum(
+        pre_mix.unsqueeze(-1) * residual_flat.to(torch.float32), dim=1
+    ).to(torch.bfloat16)
+    return (
+        post_mix.view(*outer_shape, hc_mult, 1),
+        comb_mix.view(*outer_shape, hc_mult, hc_mult),
+        layer_input.view(*outer_shape, hidden_size),
+    )
+
+
+def mhc_post_torch(
+    x: torch.Tensor,
+    residual: torch.Tensor,
+    post_layer_mix: torch.Tensor,
+    comb_res_mix: torch.Tensor,
+) -> torch.Tensor:
+    mixed_residual = torch.einsum(
+        "...ij,...ih->...jh",
+        comb_res_mix.to(torch.float32),
+        residual.to(torch.float32),
+    )
+    post_term = post_layer_mix.to(torch.float32) * x.unsqueeze(-2).to(torch.float32)
+    return (mixed_residual + post_term).to(residual.dtype)

vllm/model_executor/kernels/mhc/triton.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import direct_register_custom_op
+
+
+@triton.jit
+def _rmsnorm_nw_kernel(
+    x_ptr,
+    out_ptr,
+    stride_row,
+    D,
+    eps,
+    RBLOCK: tl.constexpr,
+):
+    """Weight-free RMSNorm Triton kernel: out = x * rsqrt(mean(x², -1) + eps)."""
+    row = tl.program_id(0)
+    cols = tl.arange(0, RBLOCK)
+    mask = cols < D
+
+    x = tl.load(
+        x_ptr + row * stride_row + cols,
+        mask=mask,
+        other=0.0,
+        eviction_policy="evict_first",
+    ).to(tl.float32)
+
+    var = tl.sum(x * x, 0) / D
+    rstd = tl.rsqrt(var + eps)
+
+    out = (x * rstd).to(out_ptr.dtype.element_ty)
+    tl.store(out_ptr + row * D + cols, out, mask=mask, eviction_policy="evict_first")
+
+
+def rmsnorm_nw(x: Tensor, eps: float) -> Tensor:
+    """Weight-free RMSNorm over the last dimension.
+
+    Treats *x* as ``[num_rows, D]`` where ``num_rows = product(shape[:-1])``.
+    Returns a contiguous tensor with the same shape and dtype as *x*.
+    """
+    orig_shape = x.shape
+    D = orig_shape[-1]
+    x_2d = x.reshape(-1, D)
+    num_rows = x_2d.shape[0]
+
+    out = torch.empty_like(x_2d)
+    RBLOCK = triton.next_power_of_2(D)
+
+    _rmsnorm_nw_kernel[(num_rows,)](
+        x_2d,
+        out,
+        x_2d.stride(0),
+        D,
+        eps,
+        RBLOCK=RBLOCK,
+        num_warps=1 if RBLOCK <= 512 else (4 if RBLOCK <= 4096 else 8),
+    )
+    return out.view(orig_shape)
+
+
+@triton.jit
+def _hc_head_reduce_store_kernel(
+    pre_ptr,
+    x_ptr,
+    out_ptr,
+    hidden_size: tl.constexpr,
+    hc_mult: tl.constexpr,
+    pre_stride_t: tl.constexpr,
+    pre_stride_m: tl.constexpr,
+    x_stride_t: tl.constexpr,
+    x_stride_m: tl.constexpr,
+    x_stride_h: tl.constexpr,
+    out_stride_t: tl.constexpr,
+    out_stride_h: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    block_idx = tl.program_id(1)
+    offsets = block_idx * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask = offsets < hidden_size
+
+    acc = tl.zeros((BLOCK_H,), dtype=tl.float32)
+    for mix_idx in tl.static_range(0, hc_mult):
+        pre = tl.load(pre_ptr + token_idx * pre_stride_t + mix_idx * pre_stride_m).to(
+            tl.float32
+        )
+        x = tl.load(
+            x_ptr
+            + token_idx * x_stride_t
+            + mix_idx * x_stride_m
+            + offsets * x_stride_h,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+        acc += pre * x
+
+    tl.store(
+        out_ptr + token_idx * out_stride_t + offsets * out_stride_h,
+        acc,
+        mask=mask,
+    )
+
+
+def hc_head_reduce_triton_kernel(
+    x: torch.Tensor,
+    hc_fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    out: torch.Tensor,
+    norm_eps: float,
+    hc_eps: float,
+) -> None:
+    x_flat = x.flatten(-2)
+    x_normed = rmsnorm_nw(x_flat, norm_eps)
+    mixes = F.linear(x_normed.float(), hc_fn)
+    pre = torch.sigmoid(mixes * hc_scale + hc_base) + hc_eps
+
+    hidden_size = x.shape[-1]
+    hc_mult = x.shape[-2]
+    block_h = 1024
+    _hc_head_reduce_store_kernel[(x.shape[0], (hidden_size + block_h - 1) // block_h)](
+        pre,
+        x,
+        out,
+        hidden_size,
+        hc_mult,
+        pre.stride(0),
+        pre.stride(1),
+        x.stride(0),
+        x.stride(1),
+        x.stride(2),
+        out.stride(0),
+        out.stride(1),
+        BLOCK_H=block_h,
+        num_warps=4,
+    )
+
+
+def _hc_head_triton(
+    hs_flat: torch.Tensor,
+    fn: torch.Tensor,
+    hc_scale: torch.Tensor,
+    hc_base: torch.Tensor,
+    out: torch.Tensor,
+    hidden_size: int,
+    rms_eps: float,
+    hc_eps: float,
+    hc_mult: int,
+) -> None:
+    """Fill pre-allocated `out` (T, H) in-place with the hc_head result."""
+    if hs_flat.shape[0] == 0:
+        return
+
+    hc_head_reduce_triton_kernel(
+        hs_flat,
+        fn,
+        hc_scale,
+        hc_base,
+        out,
+        rms_eps,
+        hc_eps,
+    )
+    return
+
+
+direct_register_custom_op(
+    op_name="hc_head_triton",
+    op_func=_hc_head_triton,
+    mutates_args=["out"],
+)