gmm tuning

Author: Prayer3th

python/sgl_jax/srt/layers/moe.py

@@ -328,32 +328,31 @@ def _gmm_compute_with_sharded_weights(
             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),
-        # )
+        static_tiling_gate = (
+            int(optimal_tiling_gate[0]),
+            int(optimal_tiling_gate[1]),
+            int(optimal_tiling_gate[2]),
+        )
+        static_tiling_down = (
+            int(optimal_tiling_down[0]),
+            int(optimal_tiling_down[1]),
+            int(optimal_tiling_down[2]),
+        )
         # gate
         layer_w0 = gmm(
             lhs=x,
             rhs=w0_kernel,
             group_sizes=local_group_sizes,
             preferred_element_type=self.dtype,
-            tiling=optimal_tiling_gate,
+            tiling=static_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=static_tiling_gate,
         )
 
         # activation
@@ -366,7 +365,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=static_tiling_down,
         )
 
         return intermediate_output

python/sgl_jax/srt/managers/tp_worker_overlap_thread.py

@@ -248,5 +248,8 @@ def run_precompile(self):
     def run_gmm_auto_tune(self):
         self.worker.run_gmm_auto_tune()
 
+    def get_gmm_tiling_configs(self):
+        return self.worker.get_gmm_tiling_configs()
+
     def __delete__(self):
         self.input_queue.put((None, None, None, None))