wip

Author: wenxindongwork

tpu_inference/layers/vllm/fused_moe.py

@@ -143,7 +143,6 @@ def _add_bias(gmm_result_local, rhs_bias_local, group_sizes_global):
             mesh=mesh,
             in_specs=(P("data", "model"), P(None, "model"), P("data")),
             out_specs=(P("data", "model")),
-            check_rep=False,
         )(gmm_result, rhs_bias, group_sizes)
 
     n_shards = mesh.shape["model"]
@@ -200,7 +199,6 @@ def _add_bias(gmm_result_local, rhs_bias_local, group_sizes_global):
             mesh=mesh,
             in_specs=(P("data"), P(), P("data")),
             out_specs=(P("data")),
-            check_rep=False,
         )(gmm_result, rhs_bias, group_sizes)
 
     return gmm_result