minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

rocm: fix oom when loading fp8 weights close to size of available vram (#19941)

Browse Source
This commit is contained in:
Evgueni Petrov
2026-03-17 22:44:19 -07:00
committed by GitHub
parent 0acc1d3c9a
commit 2e860233ca
1 changed files with 8 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/fp8.py
+8 -16
View File
@@ -998,23 +998,15 @@ class Fp8MoEMethod(FusedMoEMethodBase):
w2_weight_scale, requires_grad=False w2_weight_scale, requires_grad=False
) )
layer.w2_input_scale = None layer.w2_input_scale = None
if _use_aiter:
# add this section for MI300 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)
)
elif _use_aiter:
# Pre-shuffle weights # Pre-shuffle weights
layer.w13_weight.data = shuffle_weight( t = shuffle_weight(layer.w13_weight, (16, 16))
layer.w13_weight.contiguous(), (16, 16) layer.w13_weight.copy_(t)
) del t
layer.w2_weight.data = shuffle_weight( t = shuffle_weight(layer.w2_weight, (16, 16))
layer.w2_weight.contiguous(), (16, 16) layer.w2_weight.copy_(t)
) del t
elif _is_cpu: elif _is_cpu:
assert ( assert (
_is_cpu_amx_available _is_cpu_amx_available