[DP/EP][GPTOSS] Use triton matmul-ogs kernels for GPTOSS DP/EP (#24588)

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>
2025-09-23 00:01:09 -04:00
parent fafbe11af4
commit e8db44f883
6 changed files with 275 additions and 76 deletions
--- a/vllm/model_executor/layers/fused_moe/modular_kernel.py
+++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py
@@ -76,7 +76,7 @@ def _moe_problem_size(
    """
    assert w1.dim() == 3 and w2.dim() == 3
    E, N, _ = w1.size()
-    K = w2.size(1)
+    K = a1.size(-1)

    if a1.dim() == 2:
        # Make sure we are using the correct a1 (pre-permute).