perf: exl3 decode kernel optimization experiments

Signed-off-by: AlpinDale <alpindale@gmail.com>

Author: AlpinDale

aphrodite/config/model.py

@@ -585,6 +585,7 @@ def __post_init__(
             if self.pooler_config.tok_pooling_type is None:
                 self.pooler_config.tok_pooling_type = default_tok_pooling_type
 
+        requested_dtype = self.dtype
         self.dtype: torch.dtype = _get_and_verify_dtype(
             self.model,
             self.hf_config,
@@ -667,6 +668,15 @@ def __post_init__(
         self.config_updated = False
         self._try_verify_and_update_model_config()
         self._verify_quantization()
+        if (
+            self.quantization == "exl3"
+            and isinstance(requested_dtype, str)
+            and requested_dtype.lower() == "auto"
+            and self.dtype != torch.float16
+            and "moe" in self.hf_config.model_type.lower()
+        ):
+            logger.info("Defaulting EXL3 activation dtype from %s to torch.float16.", self.dtype)
+            self.dtype = torch.float16
         self._verify_cuda_graph()
         self._verify_bnb_config()
 

aphrodite/model_executor/layers/quantization/exl3.py

@@ -184,6 +184,8 @@ def _exl3_gate_up(
         -1,
         0,
     )
+    if x.shape[0] == 1:
+        return output.view(1, out_features * 2)
     return torch.cat([output[0], output[1]], dim=-1)
 
 

csrc/quantization/exl3/exllamav3_ext/quant/exl3_gemm_kernel.cuh

@@ -223,4 +223,4 @@ __global__ __launch_bounds__(EXL3_GEMM_BASE_THREADS* TILESIZE_K /
       }
     }
   }
-}
+}

csrc/quantization/exl3/exllamav3_ext/quant/exl3_kernel_map.cu

@@ -168,6 +168,9 @@ fp_exl3_mgemm_kernel select_exl3_mgemm_kernel(
   } else if (cc == CC_BLACKWELL && K == 4 && size_m == 1 && size_k == 1024 &&
              size_n == 3072 && bszm_out == 2) {
     shape_idx = 2;
+  } else if (cc == CC_BLACKWELL && K == 4 && size_m == 1 && size_k == 256 &&
+             size_n == 1024 && bszm_out <= 32) {
+    shape_idx = 2;
   } else {
     shape_idx = select_gemm_shape(cc, size_m, size_k, size_n, K, true, bszm_in,
                                   bszm_out);