gmm tuning

Author: Prayer3th

python/sgl_jax/srt/layers/moe.py

@@ -207,6 +207,32 @@ 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,
@@ -253,15 +279,16 @@ def _internal_moe_computation(
             else:
                 local_group_sizes = group_sizes
 
-            # GMM
-            intermediate_output = self._gmm_compute_with_sharded_weights(
+            # GMM (使用预先计算的静态 tiling)
+            intermediate_output = self._gmm_compute_with_static_tiling(
                 x,
                 local_group_sizes,
                 selected_experts,
                 w0_weights,
                 w1_weights,
                 wo_weights,
-                gmm_tiling_configs,
+                static_tiling_gate,
+                static_tiling_down,
             )
 
             # EP Combine
@@ -301,59 +328,41 @@ def _internal_moe_computation(
             self.wo.value,
         )
 
-    def _gmm_compute_with_sharded_weights(
+    def _gmm_compute_with_static_tiling(
         self,
         x,
         local_group_sizes,
         selected_experts,
         w0_kernel,
         w1_kernel,
         wo_kernel,
-        gmm_tiling_configs,
+        static_tiling_gate,
+        static_tiling_down,
     ):
         if x.shape[0] == 0:
             empty_output = jnp.zeros(
                 (0, wo_kernel.shape[-1]), dtype=x.dtype
             )  # (0, hidden_dim)
             return empty_output
 
-        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
-        )
-
-        # Use JAX operations for tiling parameters (cannot use int() on tracers)
-        # tiling_gate = (
-        #     jnp.minimum(optimal_tiling_gate[0], m),
-        #     jnp.minimum(optimal_tiling_gate[1], k),
-        #     jnp.minimum(optimal_tiling_gate[2], n_gate),
-        # )
-        # tiling_down = (
-        #     jnp.minimum(optimal_tiling_down[0], m),
-        #     jnp.minimum(optimal_tiling_down[1], n_gate),
-        #     jnp.minimum(optimal_tiling_down[2], n_down),
-        # )
+        # 直接使用预先计算好的静态 tiling 参数
+        tiling_gate = static_tiling_gate
+        tiling_down = static_tiling_down
         # gate
         layer_w0 = gmm(
             lhs=x,
             rhs=w0_kernel,
             group_sizes=local_group_sizes,
             preferred_element_type=self.dtype,
-            tiling=optimal_tiling_gate,
+            tiling=tiling_gate,
         )
         # up
         layer_w1 = gmm(
             lhs=x,
             rhs=w1_kernel,
             group_sizes=local_group_sizes,
             preferred_element_type=self.dtype,
-            tiling=optimal_tiling_gate,
+            tiling=tiling_gate,
         )
 
         # activation
@@ -366,7 +375,7 @@ def _gmm_compute_with_sharded_weights(
             rhs=wo_kernel,
             group_sizes=local_group_sizes,
             preferred_element_type=self.dtype,
-            tiling=optimal_tiling_down,
+            tiling=tiling_down,
         )
 
         return intermediate_output
@@ -381,13 +390,33 @@ 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)
 
-        return self._single_device_forward_impl(
-            inputs, top_k_indices, top_k_weights, gmm_tiling_configs
+        # 为单设备也预先计算静态 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
+        )
+
+        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 _single_device_forward_impl(
-        self, inputs, top_k_indices, top_k_weights, gmm_tiling_configs
-    ):
+        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):
         num_tokens = inputs.shape[0] * (inputs.shape[1] if inputs.ndim > 1 else 1)
         inputs_flat = inputs.reshape(num_tokens, -1)