From 2e860233ca375aa0e6e76378f31d0a1b278cf9d1 Mon Sep 17 00:00:00 2001 From: Evgueni Petrov Date: Wed, 18 Mar 2026 12:44:19 +0700 Subject: [PATCH] rocm: fix oom when loading fp8 weights close to size of available vram (#19941) --- python/sglang/srt/layers/quantization/fp8.py | 24 +++++++------------- 1 file changed, 8 insertions(+), 16 deletions(-) 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