diff --git a/python/sglang/srt/layers/quantization/fp8.py b/python/sglang/srt/layers/quantization/fp8.py index 1e12baff1..c7874830d 100644 --- a/python/sglang/srt/layers/quantization/fp8.py +++ b/python/sglang/srt/layers/quantization/fp8.py @@ -998,23 +998,15 @@ class Fp8MoEMethod(FusedMoEMethodBase): w2_weight_scale, requires_grad=False ) layer.w2_input_scale = None - if _use_aiter: - # add this section for MI300 - # Pre-shuffle weights - layer.w13_weight.data = shuffle_weight( - layer.w13_weight.contiguous(), (16, 16) - ) - layer.w2_weight.data = shuffle_weight( - layer.w2_weight.contiguous(), (16, 16) - ) - elif _use_aiter: + + if _use_aiter: # Pre-shuffle weights - layer.w13_weight.data = shuffle_weight( - layer.w13_weight.contiguous(), (16, 16) - ) - layer.w2_weight.data = shuffle_weight( - layer.w2_weight.contiguous(), (16, 16) - ) + t = shuffle_weight(layer.w13_weight, (16, 16)) + layer.w13_weight.copy_(t) + del t + t = shuffle_weight(layer.w2_weight, (16, 16)) + layer.w2_weight.copy_(t) + del t elif _is_cpu: assert ( _is_cpu_amx_available