EngineX-Hygon/sglang
5
0
Fork 0
Code Issues Pull Requests Actions 7 Projects Releases Wiki Activity

Change tensor alignment method to mn major (#9844)

Browse Source
This commit is contained in:
Mohammad Miadh Angkad
2025-09-02 16:23:13 +08:00
committed by GitHub
parent 1fbfdebe6b
commit b9eb0d9c2b
1 changed files with 2 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
python/sglang/srt/layers/moe/ep_moe/layer.py
View File

@@ -229,7 +229,7 @@ class EPMoE(FusedMoE):
( (
_cast_to_e8m0_with_rounding_up(gateup_input_scale) _cast_to_e8m0_with_rounding_up(gateup_input_scale)
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 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 gateup_input_scale
) )
), ),
@@ -286,9 +286,7 @@ class EPMoE(FusedMoE):
( (
down_input_scale down_input_scale
if deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0 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(down_input_scale)
down_input_scale
)
), ),
) )
down_output = torch.empty( down_output = torch.empty(