chenchenghao/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-18 12: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

24
python/sglang/srt/layers/quantization/fp8.py
View File

@@ -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