[CPU] expand the interface of shared_expert without scaling factor (#22933)

merge since this is CPU only change on sgl-kernel.

test/srt/cpu/utils.py

@@ -126,16 +126,28 @@ def native_w8a8_per_token_matmul(A, B, As, Bs, bias, output_dtype=torch.bfloat16
     return C.reshape(origin_C_shape).to(output_dtype)
 
 
-def torch_naive_moe(a, w1, w2, b, routed_scaling_factor):
+def torch_naive_moe(a, w1, w2, b, routed_scaling_factor, output_dtype=torch.bfloat16):
+
+    a = a.to(torch.float32)
+    w1 = w1.to(torch.float32)
+    w2 = w2.to(torch.float32)
+    b = b.to(torch.float32) if b is not None else None
 
     ic1 = torch.matmul(a, w1.transpose(0, 1))
     ic2 = SiluAndMul(ic1)
     ic3 = torch.matmul(ic2, w2.transpose(0, 1))
 
-    return ic3 + b * routed_scaling_factor
+    out = ic3 if b is None else ic3 + b * routed_scaling_factor
+
+    return out.to(output_dtype)
 
 
-def torch_w8a8_per_column_moe(a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_factor):
+def torch_w8a8_per_column_moe(
+    a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_factor, output_dtype=torch.bfloat16
+):
+
+    a = a.to(torch.float32)
+    b = b.to(torch.float32) if b is not None else None
 
     # Perform per-token quantization
     a_q, a_s = per_token_quant_int8(a)
@@ -150,7 +162,9 @@ def torch_w8a8_per_column_moe(a, w1_q, w2_q, w1_s, w2_s, b, routed_scaling_facto
         a1_q, w2_q, a1_s, w2_s, bias=None, output_dtype=torch.float32
     )
 
-    return ic3 + b * routed_scaling_factor
+    out = ic3 if b is None else ic3 + b * routed_scaling_factor
+
+    return out.to(output_dtype)
 
 
 def scaled_weight(weight, scales):