[AMD] optimize Kimi K2.5 fused_moe_triton performance by tuning (#19228)

benchmark/kernels/fused_moe_triton/common_utils.py

@@ -38,6 +38,10 @@ def get_model_config(
 ) -> Dict:
     config = get_config(model_name, trust_remote_code=True)
 
+    # Replace config with text_config for encoder-decoder models after getting block_shape and architecture
+    if hasattr(config, "text_config"):
+        config = config.get_text_config()
+
     block_shape = None
     if (
         hasattr(config, "quantization_config")
@@ -46,11 +50,19 @@ def get_model_config(
         block_shape = config.quantization_config["weight_block_size"]
         assert len(block_shape) == 2
 
-    architecture = config.architectures[0]
+    if (
+        hasattr(config, "quantization_config")
+        and "config_groups" in config.quantization_config
+    ):
+        config_groups = config.quantization_config["config_groups"]
+        # Get group_size from the first group's weights config
+        first_group = next(iter(config_groups.values()), {})
+        weights_config = first_group.get("weights", {})
+        group_size = weights_config.get("group_size")
+        block_shape = [0, group_size]
+        assert len(block_shape) == 2
 
-    # Replace config with text_config for encoder-decoder models after getting block_shape and architecture
-    if hasattr(config, "text_config"):
-        config = config.get_text_config()
+    architecture = config.architectures[0]
 
     hidden_size = config.hidden_size
     if architecture == "DbrxForCausalLM":
@@ -223,6 +235,7 @@ def get_config_filename(
     use_fp8_w8a8: bool,
     use_int8_w8a8: bool,
     use_int8_w8a16: bool,
+    use_int4_w4a16: bool,
     per_channel_quant: bool,
     block_shape: List[int],
 ) -> str:
@@ -231,13 +244,18 @@ def get_config_filename(
         use_int8_w8a16=use_int8_w8a16,
         use_fp8_w8a8=use_fp8_w8a8,
         use_int8_w8a8=use_int8_w8a8,
+        use_int4_w4a16=use_int4_w4a16,
     )
 
     # NOTE(woosuk): The current naming convention uses w2.shape[2], which
     # is the intermediate size after silu_and_mul.
+    N = shard_intermediate_size // 2
+    if use_int4_w4a16:
+        N = N // 2
+
     filename = get_config_file_name(
         num_experts,
-        shard_intermediate_size // 2,
+        N,
         dtype_str,
         block_shape,
         per_channel_quant,