[JIT Kernel] Migrate dsv3_router_gemm from AOT sgl-kernel to JIT kernel (#21531)

Co-authored-by: Guohao Shao <shao.gh.98@gmail.com>
Co-authored-by: Brayden Zhong <brayden@radixark.ai>

sgl-kernel/python/sgl_kernel/__init__.py

@@ -57,7 +57,6 @@ else:
         awq_dequantize,
         bmm_fp8,
         dsv3_fused_a_gemm,
-        dsv3_router_gemm,
         fp8_blockwise_scaled_mm,
         fp8_scaled_mm,
         gptq_gemm,

sgl-kernel/python/sgl_kernel/gemm.py

@@ -191,25 +191,6 @@ def qserve_w4a8_per_group_gemm(
     return out_feats
 
 
-def dsv3_router_gemm(
-    hidden_states: torch.Tensor,
-    router_weights: torch.Tensor,
-    out_dtype: torch.dtype = torch.bfloat16,
-) -> torch.Tensor:
-    output = torch.empty(
-        hidden_states.shape[0],
-        router_weights.shape[0],
-        device=hidden_states.device,
-        dtype=out_dtype,
-    )
-    torch.ops.sgl_kernel.dsv3_router_gemm(
-        output,
-        hidden_states,
-        router_weights,
-    )
-    return output
-
-
 def shuffle_rows(input_tensor, dst2src_map, output_tensor_shape):
     output_tensor = torch.empty(
         output_tensor_shape,