Skip to content

GitLab

Unverified Commit b9eb0d9c authored by Mohammad Miadh Angkad's avatar Mohammad Miadh Angkad Committed by GitHub
Browse files

Change tensor alignment method to mn major (#9844)

parent 1fbfdebe
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 4 deletions
python/sglang/srt/layers/moe/ep_moe/layer.py
View file @ b9eb0d9c
...@@ -229,7 +229,7 @@ class EPMoE(FusedMoE): ...@@ -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): ...@@ -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(
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment