feat(exl3): update kernels from upstream

CMakeLists.txt

@@ -1192,26 +1192,29 @@ if(APHRODITE_GPU_LANG STREQUAL "CUDA" OR APHRODITE_GPU_LANG STREQUAL "HIP")
   endif()
 
   if(APHRODITE_GPU_LANG STREQUAL "CUDA")
+    file(GLOB APHRODITE_EXL3_GEMM_UNITS CONFIGURE_DEPENDS
+      "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_*_cb*.cu")
+    file(GLOB APHRODITE_EXL3_GEMV_INT8_UNITS CONFIGURE_DEPENDS
+      "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_gemv_int8_inst_*.cu")
+    file(GLOB APHRODITE_EXL3_MOE_UNITS CONFIGURE_DEPENDS
+      "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_moe_inst_*.cu")
     set(APHRODITE_EXL3_SRC
       "csrc/libtorch_stable/exl3_torch_bindings.cpp"
       "csrc/quantization/exl3/exllamav3_ext/graph.cu"
       "csrc/quantization/exl3/exllamav3_ext/hgemm.cu"
+      "csrc/quantization/exl3/exllamav3_ext/quant/coop_autotune.cu"
       "csrc/quantization/exl3/exllamav3_ext/quant/exl3_devctx.cu"
       "csrc/quantization/exl3/exllamav3_ext/quant/exl3_gemm.cu"
       "csrc/quantization/exl3/exllamav3_ext/quant/exl3_gemv.cu"
+      "csrc/quantization/exl3/exllamav3_ext/quant/exl3_gemv_int8.cu"
       "csrc/quantization/exl3/exllamav3_ext/quant/exl3_kernel_map.cu"
       "csrc/quantization/exl3/exllamav3_ext/quant/exl3_moe.cu"
       "csrc/quantization/exl3/exllamav3_ext/quant/hadamard.cu"
       "csrc/quantization/exl3/exllamav3_ext/quant/reconstruct.cu"
       "csrc/quantization/exl3/exllamav3_ext/quant/util.cu"
-      "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_1.cu"
-      "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_2.cu"
-      "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_3.cu"
-      "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_4.cu"
-      "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_5.cu"
-      "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_6.cu"
-      "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_7.cu"
-      "csrc/quantization/exl3/exllamav3_ext/quant/comp_units/exl3_comp_unit_8.cu")
+      ${APHRODITE_EXL3_GEMM_UNITS}
+      ${APHRODITE_EXL3_GEMV_INT8_UNITS}
+      ${APHRODITE_EXL3_MOE_UNITS})
 
     # EXL3 requires Ampere+ instructions. Preserve all configured SM80+
     # architectures, including native Thor SM110 cubins.

aphrodite/_custom_ops.py

@@ -187,6 +187,17 @@ def exl3_reconstruct(
     torch.ops._C.exl3_reconstruct(unpacked, packed, k, mcg, mul1)
 
 
+def exl3_reconstruct_slice(
+    unpacked: torch.Tensor,
+    packed: torch.Tensor,
+    k: int,
+    mcg: bool,
+    mul1: bool,
+    n_offset: int,
+) -> None:
+    torch.ops._C.exl3_reconstruct_slice(unpacked, packed, k, mcg, mul1, n_offset)
+
+
 def exl3_had_r_128(
     input: torch.Tensor,
     output: torch.Tensor,
@@ -197,6 +208,30 @@ def exl3_had_r_128(
     torch.ops._C.exl3_had_r_128(input, output, pre_scale, post_scale, scale)
 
 
+def exl3_had_r_128_dual(
+    input1: torch.Tensor,
+    output1: torch.Tensor,
+    pre_scale1: torch.Tensor | None,
+    post_scale1: torch.Tensor | None,
+    input2: torch.Tensor,
+    output2: torch.Tensor,
+    pre_scale2: torch.Tensor | None,
+    post_scale2: torch.Tensor | None,
+    scale: float = 1.0,
+) -> None:
+    torch.ops._C.exl3_had_r_128_dual(
+        input1,
+        output1,
+        pre_scale1,
+        post_scale1,
+        input2,
+        output2,
+        pre_scale2,
+        post_scale2,
+        scale,
+    )
+
+
 def exl3_hgemm(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> None:
     torch.ops._C.exl3_hgemm(a, b, c)
 
@@ -231,8 +266,10 @@ def exl3_moe(
     down_mcg: bool,
     down_mul1: bool,
     act_limit: float = 0.0,
+    num_active: int = -1,
 ) -> None:
-    torch.ops._C.exl3_moe(
+    op = torch.ops._C.exl3_moe if num_active < 0 else torch.ops._C.exl3_moe_active
+    args = (
         hidden_state,
         output_state,
         expert_count,
@@ -263,6 +300,10 @@ def exl3_moe(
         down_mul1,
         act_limit,
     )
+    if num_active < 0:
+        op(*args)
+    else:
+        op(*args, num_active)
 
 
 if hasattr(torch.ops, "_C") and hasattr(torch.ops._C, "exl3_gemm"):
@@ -312,6 +353,17 @@ def _exl3_reconstruct_fake(
     ) -> None:
         return None
 
+    @register_fake("_C::exl3_reconstruct_slice")
+    def _exl3_reconstruct_slice_fake(
+        unpacked: torch.Tensor,
+        packed: torch.Tensor,
+        k: int,
+        mcg: bool,
+        mul1: bool,
+        n_offset: int,
+    ) -> None:
+        return None
+
     @register_fake("_C::exl3_had_r_128")
     def _exl3_had_r_128_fake(
         input: torch.Tensor,

aphrodite/model_executor/layers/fused_moe/routed_experts.py

@@ -5,6 +5,7 @@
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Literal, cast, overload
 
+import regex as re
 import torch
 
 from aphrodite.distributed.eplb.eplb_state import EplbState
@@ -853,8 +854,12 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> Iterable[
         unpadded_hidden = self.moe_config.hidden_dim_unpadded
         for expert_name, loaded_weight in weights:
             qual_name = f"{self.layer_name}.{expert_name}"
-            # Fused expert weights can be identified by their 3D tensors
-            is_fused = loaded_weight.dim() == 3
+            # Most fused expert weights are 3D, but some quantization formats
+            # use 3D tensors for a single expert (for example EXL3 trellises).
+            # An explicit expert index in the checkpoint name takes priority
+            # over the tensor rank.
+            is_per_expert = re.search(r"(?:^|\.)experts\.\d+\.", qual_name) is not None
+            is_fused = loaded_weight.dim() == 3 and not is_per_expert
             matched = False
             for param_name, weight_name, expert_id, shard_id in expert_mapping:
                 if weight_name not in qual_name: