Change tensor alignment method to mn major (#9844)

b9eb0d9c · Mohammad Miadh Angkad · GitHub · 1fbfdebe · b9eb0d9c
Unverified Commit b9eb0d9c authored Sep 02, 2025 by Mohammad Miadh Angkad Committed by GitHub Sep 02, 2025
Hide whitespace changes
Inline Side-by-side

Showing with 2 additions and 4 deletions

python/sglang/srt/layers/moe/ep_moe/layer.py python/sglang/srt/layers/moe/ep_moe/layer.py +2 -4

No files found.
--- a/python/sglang/srt/layers/moe/ep_moe/layer.py
+++ b/python/sglang/srt/layers/moe/ep_moe/layer.py
@@ -229,7 +229,7 @@ class EPMoE(FusedMoE):
            (
                _cast_to_e8m0_with_rounding_up(gateup_input_scale)
                if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
-                else deep_gemm_wrapper.get_col_major_tma_aligned_tensor(
+                else deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(
                    gateup_input_scale
                )
            ),
@@ -286,9 +286,7 @@ class EPMoE(FusedMoE):
            (
                down_input_scale
                if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
-                else deep_gemm_wrapper.get_col_major_tma_aligned_tensor(
-                    down_input_scale
-                )
+                else deep_gemm_wrapper.get_mn_major_tma_aligned_tensor(down_input_scale)
            ),
        )
        down_output = torch.empty(