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

[ROCM][RL] Shuffle Weight In-Place to Preserve Parameter Attributes (#21825)

Browse Source
This commit is contained in:
Yuzhen Zhou
2026-04-02 23:43:55 -07:00
committed by GitHub
parent 75de479680
commit 6b876a7710
1 changed files with 5 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/quantization/unquant.py
+5 -7
View File
@@ -22,7 +22,7 @@ from sglang.srt.layers.quantization.base_config import (
LinearMethodBase, LinearMethodBase,
QuantizeMethodBase, QuantizeMethodBase,
) )
from sglang.srt.layers.utils import MultiPlatformOp from sglang.srt.layers.utils import MultiPlatformOp, copy_or_rebind_param
from sglang.srt.utils import ( from sglang.srt.utils import (
cpu_has_amx_support, cpu_has_amx_support,
get_bool_env_var, get_bool_env_var,
@@ -233,14 +233,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp):
# because aiter CK kernels don't support all GEMM dimensions # because aiter CK kernels don't support all GEMM dimensions
_should_use_aiter_moe = _use_aiter and get_moe_runner_backend().is_auto() _should_use_aiter_moe = _use_aiter and get_moe_runner_backend().is_auto()
if _should_use_aiter_moe: if _should_use_aiter_moe:
layer.w13_weight = torch.nn.Parameter( copy_or_rebind_param(
shuffle_weight(layer.w13_weight.data, (16, 16)), layer, "w13_weight", shuffle_weight(layer.w13_weight.data, (16, 16))
requires_grad=False,
) )
torch.cuda.empty_cache() torch.cuda.empty_cache()
layer.w2_weight = torch.nn.Parameter( copy_or_rebind_param(
shuffle_weight(layer.w2_weight.data, (16, 16)), layer, "w2_weight", shuffle_weight(layer.w2_weight.data, (16, 16))
requires_grad=False,
) )
torch.cuda.empty_cache() torch.cuda.empty_cache()