From 6b876a7710a1512f0f89097dd7d5821002913982 Mon Sep 17 00:00:00 2001 From: Yuzhen Zhou <82826991+zyzshishui@users.noreply.github.com> Date: Thu, 2 Apr 2026 23:43:55 -0700 Subject: [PATCH] [ROCM][RL] Shuffle Weight In-Place to Preserve Parameter Attributes (#21825) --- python/sglang/srt/layers/quantization/unquant.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 5c2f489e9..94f9a1375 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -22,7 +22,7 @@ from sglang.srt.layers.quantization.base_config import ( LinearMethodBase, QuantizeMethodBase, ) -from sglang.srt.layers.utils import MultiPlatformOp +from sglang.srt.layers.utils import MultiPlatformOp, copy_or_rebind_param from sglang.srt.utils import ( cpu_has_amx_support, get_bool_env_var, @@ -233,14 +233,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, MultiPlatformOp): # because aiter CK kernels don't support all GEMM dimensions _should_use_aiter_moe = _use_aiter and get_moe_runner_backend().is_auto() if _should_use_aiter_moe: - layer.w13_weight = torch.nn.Parameter( - shuffle_weight(layer.w13_weight.data, (16, 16)), - requires_grad=False, + copy_or_rebind_param( + layer, "w13_weight", shuffle_weight(layer.w13_weight.data, (16, 16)) ) torch.cuda.empty_cache() - layer.w2_weight = torch.nn.Parameter( - shuffle_weight(layer.w2_weight.data, (16, 16)), - requires_grad=False, + copy_or_rebind_param( + layer, "w2_weight", shuffle_weight(layer.w2_weight.data, (16, 16)) ) torch.cuda.empty_cache()