fix: preserve parameter attrs during weight reload (FP8 block + MXFP4 MoE)

[fergusfinn]

Summary

Fixes two weight reload crashes that break sleep/wake (L2) and RL weight reloading for quantized MoE models:

• FP8 block-wise MoE (e.g. stepfun-ai/Step-3.5-Flash-FP8): quant method must be one of ['tensor', 'channel', 'group', 'block']
• MXFP4 MoE (e.g. openai/gpt-oss-20b): default_weight_loader() got an unexpected keyword argument 'weight_name'

Both have the same root cause: custom parameter attributes are dropped when process_weights_after_loading replaces parameters with new nn.Parameter objects.

Root cause

replace_parameter() in utils.py creates a new nn.Parameter and only copies weight_loader from the old parameter. Other attributes like quant_method (needed by FP8 block-wise) are silently dropped. Additionally, marlin_utils_fp4.py bypasses replace_parameter() entirely, using raw setattr() which copies nothing.

After the initial load completes, these attributes are gone. When reload_weights() restores the model from saved metadata and tries to re-run weight loading, the missing attributes cause crashes:

1. FP8: FusedMoE.weight_loader reads getattr(param, "quant_method", None) → None → falls through all branches → ValueError
2. MXFP4: Params have no weight_loader → fall back to default_weight_loader(param, loaded_weight) which doesn't accept weight_name/shard_id/expert_id kwargs → TypeError

Changes

1. utils.py: replace_parameter() now copies all __dict__ attributes from the old parameter, not just weight_loader
2. marlin_utils_fp4.py: Replace raw nn.Parameter() + setattr() with replace_parameter() so MXFP4 Marlin gets the same attribute preservation

Upstream references

• Fixes vllm-project/vllm#23577
• replace_parameter was introduced in vllm-project/vllm#28480
• Layerwise reload system in vllm-project/vllm#32133
• Both bugs are still live on vllm main as of 2026-02-11

Test plan

• L2 sleep/wake cycle with stepfun-ai/Step-3.5-Flash-FP8 TP=2
• L2 sleep/wake cycle with openai/gpt-oss-20b MXFP4
• Verify initial model loading is unaffected (no regression)
• Run upstream test_reload.py (includes FP8_BLOCK and FP8_DYNAMIC reload tests)

🤖 Generated with Claude Code

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: a6e725f5f4

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Preserve weight_loader property when replacing params

replace_parameter now copies old_param.__dict__, but many quantized weights are BasevLLMParameter subclasses where weight_loader is exposed via a class property backed by _weight_loader (vllm/model_executor/parameter.py), not an instance weight_loader field. After replacement the new object is a plain torch.nn.Parameter, so copying __dict__ keeps _weight_loader but drops accessible param.weight_loader; this breaks reload paths that expect param.weight_loader (for example step3p5.py calls it directly with shard/expert kwargs) and can regress to the same reload failures this change is trying to fix.

Useful? React with 👍 / 👎.