[sgl-kernel][2/N]Support Expert Specialization Grouped GEMM (#11534)

2025-10-14 07:24:48 +08:00
parent 6dc9ca8c85
commit dc48c4c0e3
4 changed files with 112 additions and 106 deletions
--- a/sgl-kernel/python/sgl_kernel/expert_specialization.py
+++ b/sgl-kernel/python/sgl_kernel/expert_specialization.py
@@ -0,0 +1,27 @@
+import torch
+
+
+def es_fp8_blockwise_scaled_grouped_mm(
+    output,
+    a,
+    b,
+    scales_a,
+    scales_b,
+    stride_a,
+    stride_b,
+    stride_d,
+    problem_sizes,
+    expert_offsets,
+):
+    torch.ops.sgl_kernel.es_fp8_blockwise_scaled_grouped_mm.default(
+        output,
+        a,
+        b,
+        scales_a,
+        scales_b,
+        stride_a,
+        stride_b,
+        stride_d,
+        problem_sizes,
+        expert_offsets,
+    )