gmm tuning

Author: Prayer3th

python/sgl_jax/srt/layers/moe.py

@@ -164,8 +164,6 @@ def _get_tiling_from_configs(
         self, gmm_tiling_configs, m: int, k: int, n: int, num_groups: int
     ):
         key = (m, k, n, num_groups)
-        if gmm_tiling_configs is None:
-            return (8, 1024, 1024)  # Default fallback when configs not loaded
         return gmm_tiling_configs.get(key, (8, 1024, 1024))  # Default fallback
 
     def _detect_device_capabilities(self):
@@ -207,32 +205,6 @@ def __call__(self, inputs, router_logits=None, gmm_tiling_configs=None):
     def _expert_parallel_forward_with_shard_map(
         self, inputs, router_logits, gmm_tiling_configs
     ):
-        # 预先计算静态 tiling 参数（在 shard_map 外部）
-        total_tokens, hidden_dim = inputs.shape
-        m, k = total_tokens, hidden_dim
-        n_gate = self.intermediate_dim
-        n_down = hidden_dim
-
-        # 获取最优 tiling 配置
-        optimal_tiling_gate = self._get_tiling_from_configs(
-            gmm_tiling_configs, m, k, n_gate, self.num_experts
-        )
-        optimal_tiling_down = self._get_tiling_from_configs(
-            gmm_tiling_configs, m, n_gate, n_down, self.num_experts
-        )
-
-        # 转换为静态整数参数（在动态 m 值上使用最大值作为安全的静态值）
-        static_tiling_gate = (
-            min(optimal_tiling_gate[0], 16384),  # 设置合理的最大值
-            optimal_tiling_gate[1],
-            optimal_tiling_gate[2],
-        )
-        static_tiling_down = (
-            min(optimal_tiling_down[0], 16384),
-            optimal_tiling_down[1],
-            optimal_tiling_down[2],
-        )
-
         def _internal_moe_computation(
             hidden_states,
             router_logits,
@@ -279,16 +251,15 @@ def _internal_moe_computation(
             else:
                 local_group_sizes = group_sizes
 
-            # GMM (使用预先计算的静态 tiling)
-            intermediate_output = self._gmm_compute_with_static_tiling(
+            # GMM
+            intermediate_output = self._gmm_compute_with_sharded_weights(
                 x,
                 local_group_sizes,
                 selected_experts,
                 w0_weights,
                 w1_weights,
                 wo_weights,
-                static_tiling_gate,
-                static_tiling_down,
+                gmm_tiling_configs,
             )
 
             # EP Combine
@@ -328,26 +299,44 @@ def _internal_moe_computation(
             self.wo.value,
         )
 
-    def _gmm_compute_with_static_tiling(
+    def _gmm_compute_with_sharded_weights(
         self,
         x,
         local_group_sizes,
         selected_experts,
         w0_kernel,
         w1_kernel,
         wo_kernel,
-        static_tiling_gate,
-        static_tiling_down,
+        gmm_tiling_configs,
     ):
         if x.shape[0] == 0:
             empty_output = jnp.zeros(
                 (0, wo_kernel.shape[-1]), dtype=x.dtype
             )  # (0, hidden_dim)
             return empty_output
 
-        # 直接使用预先计算好的静态 tiling 参数
-        tiling_gate = static_tiling_gate
-        tiling_down = static_tiling_down
+        m, k = x.shape[0], x.shape[1]
+        n_gate = w0_kernel.shape[2]
+        n_down = wo_kernel.shape[2]
+
+        optimal_tiling_gate = self._get_tiling_from_configs(
+            gmm_tiling_configs, m, k, n_gate, self.num_experts
+        )
+        optimal_tiling_down = self._get_tiling_from_configs(
+            gmm_tiling_configs, m, n_gate, n_down, self.num_experts
+        )
+
+        # Convert to Python integers for static tiling parameters
+        tiling_gate = (
+            optimal_tiling_gate[0],
+            optimal_tiling_gate[1],
+            optimal_tiling_gate[2],
+        )
+        tiling_down = (
+            optimal_tiling_down[0],
+            optimal_tiling_down[1],
+            optimal_tiling_down[2],
+        )
         # gate
         layer_w0 = gmm(
             lhs=x,
@@ -390,33 +379,13 @@ def _single_device_forward(self, inputs, router_logits, gmm_tiling_configs):
 
         top_k_weights = top_k_weights / jnp.sum(top_k_weights, axis=-1, keepdims=True)
 
-        # 为单设备也预先计算静态 tiling 参数
-        total_tokens, hidden_dim = inputs.shape
-        m, k = total_tokens, hidden_dim
-        n_gate = self.intermediate_dim
-        n_down = hidden_dim
-
-        optimal_tiling_gate = self._get_tiling_from_configs(
-            gmm_tiling_configs, m, k, n_gate, self.num_experts
-        )
-        optimal_tiling_down = self._get_tiling_from_configs(
-            gmm_tiling_configs, m, n_gate, n_down, self.num_experts
+        return self._single_device_forward_impl(
+            inputs, top_k_indices, top_k_weights, gmm_tiling_configs
         )
 
-        static_tiling_gate = (
-            min(optimal_tiling_gate[0], 16384),
-            optimal_tiling_gate[1],
-            optimal_tiling_gate[2],
-        )
-        static_tiling_down = (
-            min(optimal_tiling_down[0], 16384),
-            optimal_tiling_down[1],
-            optimal_tiling_down[2],
-        )
-
-        return self._single_device_forward_impl(inputs, top_k_indices, top_k_weights)
-
-    def _single_device_forward_impl(self, inputs, top_k_indices, top_k_weights):
+    def _single_device_forward_impl(
+        self, inputs, top_k_indices, top_k_weights, gmm_tiling_configs
+    ):
         num_tokens = inputs.shape[0] * (inputs.shape[1] if inputs.ndim > 1 else 1)
         inputs_flat = inputs.reshape(num_tokens, -1)