[RFC] Weight Reload Correctness for RL

[aoshen02]

Related RFCs: #45585 (RL CI roadmap), #48310 (sleep/wake correctness), #48305 (training-inference consistency), and #48478 (production fail-closed storage contract for category 1).

Problem

Weight reload must produce correct inference output for all model architectures and parallelism configurations. The bugs below were discovered by downstream RL frameworks in production.

This RFC classifies failures by root cause rather than symptom. A single bug may violate more than one invariant, but each item is tracked under its primary category.

Scope

In scope: model/adapter weight load and reload, post-load transformations, reload state restoration, parameter routing, weight-name resolution, and cache invalidation caused by a model or adapter update.

This tracker owns the update half of the complete RL transition:

pause/drain -> sleep -> wake(weights) -> reload/update -> post-load/finalize
            -> wake(kv_cache) -> invalidate versioned caches -> resume

The sleep/wake half is tracked in #48310. A successful RL update must satisfy both RFCs; HTTP success alone is not a correctness oracle.

Failure taxonomy

#	Category	Required invariant	Minimum regression check
1	Storage identity	A tensor address captured by CUDA/HIP Graph or retained by a backend must not change across reload	capture → reload → pointer check → allocation pressure → replay
2	Runtime value refresh	Every runtime-derived tensor must be recomputed from the new source weight and refreshed in place	load A → post-load → load B → post-load → check derived values/output
3	Loader lifecycle	A weight loader must remain valid for the parameter's current layout and lifecycle phase; accounting must finalize each layer exactly once and only after every required tensor arrives	load A → layout/kernel setup → load B; vary transfer order/chunking and check exact completion accounting
4	Reload state preservation	Parameters, buffers, configuration, and backend state required by forward must survive or be rebuilt correctly	compare cold load with warm reload across all runtime state
5	Parameter routing/sharding	Every checkpoint tensor must be written to the correct parameter, TP rank, and shard	multi-rank load with rank-distinct sentinel values
6	Name/key mapping	Checkpoint, model, and adapter names must resolve completely and unambiguously	bidirectional key-coverage test; no silent missing/unexpected keys
7	Post-update cache coherence	Any cache derived from model/adapter version N must either be invalidated or keyed by a new identity before version N+1 serves requests	warm cache on A → update to B → reuse the same request → assert the cache path is exercised and the result matches cold-cache B

1. Storage identity

This category covers long-lived runtime tensors whose device addresses are captured by a graph or retained by a backend. The first post-load pass may bind storage; every subsequent reload must reuse compatible storage. Shape/dtype/device/layout incompatibility must fail explicitly instead of silently replacing it.

Confirmed

Tracking	Component	Failure	Status
#48251	Generic MLA	W_UV/W_UK_T are rebound after graph capture	Open
#40390	ROCm AITER unquantized MoE	Weight shuffle replaced runtime parameters after an RL weight update; captured graphs continued using old addresses	Merged
#46009	ROCm unquantized MoE padding	Post-load padding assigns through .data and can replace parameter storage; replay validation checks both data_ptr and output	Open follow-up
#41670	CUTLASS grouped-GEMM FP8 MoE	Layerwise reload rebuilds moe_kernel; the old captured graph retains the replaced experts object’s ab_strides*/c_strides* storage, while the rebuilt quant config can retain staged scale tensors that are detached from the restored layer parameters	Confirmed on H200 with GLM-4.5-Air-FP8; open, no direct fix PR
#48438	Marlin linear workspace + act-order GPTQ sort indices	Kernel-owned workspace and unregistered g_idx_sort_indices rebind during post-load. On unfixed main, 88/88 pointers for each tensor moved and reclaimed workspace memory caused a graph-replay livelock	Open
All eight sibling-site A/B matrix	FP8/NVFP4/MXFP8/MXFP4 dense and MoE Marlin fallbacks, plus compressed-tensors WNA16 MoE	Every site now has live capture/reload/replay evidence on RTX 4090: the bug side rebinds workspaces, while the #48438 branch preserves every pointer; MXFP8 MoE also reproduced the graph-replay livelock	Fixed in #48438; live A/B validated
#48539	Machete act-order permutation	act_perm captured a fresh unregistered permutation during every post-load pass. H100 validation found 88/88 pointers rebound; capture artifacts retained the old storage, producing a silent stale read rather than a dangling-pointer crash	Open fix; live red/green validated
Live H100 evidence	FlashInfer CUTLASS MoE SwiGLU constants	Rebuilding the experts object replaced gemm1_alpha, gemm1_beta, and gemm1_clamp_limit; 72/72 graph-visible tensors moved and all capture-time storages expired on gpt-oss-20b	Confirmed; fix required

High-risk candidates

Audited against upstream main@83762b77. Static inspection confirms that these paths create or rebind long-lived tensors used by forward; supported-hardware replay is still required before classifying them as reproduced bugs.

Candidate	Risk	Needed validation
AITER MLA FP4/FP8 derived weights	Every post-load rebuilds W_K, W_K_scale, W_V, and W_V_scale; decode passes them directly to AITER BMM kernels	AMD HIP Graph capture → level-2 sleep/wake → reload → replay
FlashInfer B12x MoE runtime scales	_fc2_input_scale, w1_sf_mma, and w2_sf_mma are reallocated on the backend object and are outside layerwise parameter/buffer copy-back	SM120/SM121 CUDA Graph reload replay
Compressed-tensors NVFP4 CUTLASS MoE derived scales (quant-config construction)	PWAL stores weight/input scale intermediates as bare layer attributes; the CUTLASS quant config retains the fused g1/g2_alphas storage and additionally allocates 1 / a13_scale and 1 / a2_scale for graph-visible a1/a2_gscale. Rebuilding moe_kernel replaces these storages while captured CUTLASS FP4 calls retain the old addresses. Primary category 1, with category 2 value-refresh impact	Blackwell capture → reload with value-changing scales → pointer check → allocation pressure → old-graph replay against cold B
CUTLASS W4A8 FP8 MoE stride state	PWAL replaces b_strides1/2 on the quant method and rebuilds the experts object; CutlassExpertsW4A8Fp8.__init__ also allocates a_strides*, c_strides*, and s_strides*, all passed to grouped GEMM. This is the W4A8 sibling of confirmed #41670	Supported-CUDA capture → reload → pointer census → allocation pressure → replay
RDNA3 WNA16 MoE decode scratch	PWAL reallocates unregistered rdna3_w1_buf, rdna3_act_buf, and rdna3_empty_tw; the decode path passes them to HIP ops. (rdna3_out_buf is currently unused and is not part of the claim.)	gfx1100 HIP Graph capture → reload → pointer check → replay
Compressed-tensors WNA8O8 scale copies	Post-load clones activation scales, deletes the registered sources, and retains raw tensor copies on the quant method	CUDA Graph reload test with static input/output activation scales
TRT-LLM FP8/MXFP4/NVFP4 MoE constants	The SM100 expert classes allocate graph-visible gemm1_alpha/beta/clamp_limit; NVFP4 also derives g1_scale_c. Rebuilding the experts object replaces those storages. Static source tracing and CPU lifecycle probes confirm the ownership/rebind path; supported SM100 replay is still outstanding	SM100/SM120 capture → reload → storage lifetime check → replay; register the constants and g1_scale_c explicitly
ModelOpt/other NVFP4 reciprocal-scale chains	Quant configs can create 1 / a1_scale and 1 / a2_scale tensors which are retained by a rebuilt experts object. This is the same category-1/category-2 shape as the compressed-tensors row and must not be considered covered by testing only that frontend	For every advertised NVFP4 frontend/backend pair, use different A/B scales and compare old-graph warm B with cold B

The CUTLASS FP8 member is already a reproduced bug rather than a candidate: #41670 reports an H200 illegal-memory-access failure after identity reload. Machete and FlashInfer CUTLASS were promoted after live H100 validation. The NVFP4, W4A8, and TRT-LLM rows are sibling sites established by ownership/rebuild/forward data flow, but still require supported-hardware red/green replay before promotion to confirmed. These additions were rechecked against main@75fe92a and the evidence in comment 4964191093.

CPU lifecycle probes already show that two B12x post-load passes change all three runtime addresses and that WNA8O8 creates distinct scale storage while refreshing values. Backend graph replay remains the correctness gate.

A same-weights output comparison is insufficient: stale allocations can retain bit-identical contents until reclaimed. Tests for this category need pointer instrumentation or allocation pressure in addition to output comparison.

2. Runtime value refresh

A converted, packed, cast, shuffled, or otherwise derived runtime tensor must remain connected to its reloadable source. Reload must recompute its value and copy the result into the existing runtime storage.

Tracking	Component	Failure	Status
#48251	FlashInfer attention sinks	BF16/FP16 checkpoint sinks leave behind a stale FP32 runtime copy after the source parameter changes	Open

This category needs value-changing tests. Reloading identical weights can pass even when the derived copy is never refreshed.

Cross-category high-risk candidate: NVFP4 frontend paths (including compressed-tensors and ModelOpt/model-specific frontends) can recompute global/activation scale chains from the new checkpoint without refreshing the graph-visible g1/g2_alphas and a1/a2_gscale values into capture-time storage. They are tracked primarily under category 1 above. Each advertised frontend/backend pair needs a Blackwell regression with different A/B scale values and a warm-reload-B versus cold-B comparison; testing one frontend does not establish the others.

3. Loader lifecycle

Layout conversion, kernel setup, or post-load replacement must not leave a parameter with a loader that interprets its current storage using an obsolete layout.

Tracking	Failure	Impact	Status
#42821 (fix #42823)	A stale weight_loader remains after _setup_kernel() converts the layout; the second model.load_weights silently corrupts forward	Unquantized MoE on FlashInfer; root cause of SkyRL#1680	Open
#44814 (same accounting family: #37334, #38746)	Layerwise reload counted the same copied elements twice or counted/skipped tensors inconsistently, so a layer finalized too early or never finalized. In #44814, Mamba2 mixer.D arrived after premature finalization and remained uninitialized, producing NaN logits	Silent dropped weights, NaN output, or OOM depending on which side of the accounting error fires	Merged; canonical regression family

Category 3 subtypes: completion accounting and post-load coverage

Category 3 covers the complete loader lifecycle, not only the aggregate numel counter:

1. Application accounting — identify each required source application by its effective destination and loader selectors. Duplicate, missing, unexpected, or ambiguous applications must not advance completion silently.
2. Logical-numel accounting — the denominator must represent checkpoint-backed logical/loadable elements, not physical storage size. Padding, aliases, skipped buffers, and runtime-only or derived tensors need explicit treatment; otherwise finalization can happen too early or never happen.
3. Post-load dispatch coverage — every layer requiring deferred post-load processing must be in the expected set and processed exactly once after prerequisite quantization/layout work. Validate expected post-load layers == processed post-load layers == finalized layers; missing dispatch must fail closed.

A layer can therefore report a successful reload while still being incorrect: the numeric threshold may be reached while a required parameter was never applied or a required post-load hook was never invoked.

The existing reports map to these subtypes:

• #37334 and #38746: logical-numel denominators included elements that were never loaded, so finalization never fired.
• #44814: a composed loader advanced the counter twice, so finalization fired early and a later parameter was dropped.
• #42821 / #42823: a loader remained incompatible with the post-load layout on a subsequent load.
• #49519: deferred attention implementations, including MMEncoderAttention, were not consistently covered by the post-load lifecycle; first-load and reload ordering must both be validated.
• #48251: the post-load hook ran, but derived runtime tensors were not refreshed in a storage-safe way. This is primarily Category 1/2, and is the downstream failure that missing Category 3 dispatch can expose.

The minimum regression matrix should check expected versus observed application identities, logical loadable numel versus physical padded/aliased storage numel, expected versus processed post-load layer sets and invocation counts, different tensor orders/chunk boundaries/aliases/empty updates, and value-changing A → B reloads followed by derived-state and graph replay checks.

For every transfer transaction, the loader must fail closed unless all required tensors are loaded exactly once, no unexpected tensor is silently consumed, no duplicate advances completion, and no layer retains buffered tensors at commit. Tests must vary tensor order, chunk boundaries, empty updates, aliases, and composed loaders; a single canonical checkpoint order is insufficient.

4. Reload state preservation

Warm reload must preserve or correctly rebuild every piece of forward-visible state, not only persistent parameters.

Tracking	Failure	Impact	Status
#42481	Layerwise copy-back did not preserve a parent/child parameter-buffer storage alias (conv_weights versus child conv1d.weight)	Hybrid/Mamba warm reload corrupted the forward-visible alias relationship	Merged
#44371	Unloaded non-persistent buffers were copied back from materialized/meta state and corrupted after warm layerwise reload	Silent KL divergence for Gemma4 MoE and models with runtime buffers	Merged
#44613	Reload rebuilt a backend using an implicit global VllmConfig value that was no longer present; the fix snapshots max_cudagraph_capture_size into FusedMoEConfig	Reload crash or backend/config drift	Merged
#45989	enforce_eager is not snapshotted into FusedMoEConfig, so reload rebuilds the backend with the wrong configuration	MoE FP8 reload produces wrong output	Open
#48284	Checkpoint-format reload applies layerwise post-load processing to an unquantized hybrid model even though no repacking is required; the processing is not idempotent for Mamba/SSM state	Reloading identical weights changes NemotronH greedy output	Open

#40647 is historical umbrella evidence, not the canonical fix: its global-config branch is represented by #44613/#45989, and its alias/unloaded-state branch by #42481/#44371.

Partial-update policy is part of this invariant. The closed, unmerged #45835 exposed a fail-closed gap for FP8 updates that send weights but omit scale_inv: if the protocol requires the scale, classify and reject it as a category-6 missing key; if the protocol explicitly allows partial updates, category 4 must preserve the authoritative old scale. Inferring policy from whether a materialized value is NaN is not acceptable.

5. Parameter routing/sharding

Weight loaders must honor the layer-supplied rank and sharding metadata; correctness on one GPU does not establish TP correctness.

Tracking	Failure	Impact	Status
#48025	_ColumnvLLMParameter loaders ignore the layer-supplied tp_rank	Replicated parameters can receive the wrong TP slice	Open

6. Name/key mapping

Reload must resolve checkpoint, model, and adapter keys completely. Prefix insertion/removal and LoRA naming must not cause crashes, silent omissions, or writes to the wrong parameter.

Tracking	Failure	Impact	Status
#39935	With enable_lora=True, level-2 sleep/wake/reload sees LoRA-prefixed model names that do not match checkpoint names; the same flow also has LoRA-state preservation requirements	Reload crash; cross-classified as category 4 + 6 and #48310 sleep/wake	Open
#42577	Qwen3 MoE reload did not recognize the Hugging Face fused expert keys experts.gate_up_proj/experts.down_proj	Fused expert weights were not routed into vLLM's packed expert parameters	Merged; historical key-mapping regression

The update protocol must define an exact key manifest: required, loaded, explicitly_preserved, allowed_missing, and unexpected. Commit fails if any required key is absent, duplicated, ambiguous, or silently skipped. Model-specific/PP missing-key allowlists must be explicit and must not hide a real omission such as the FP8-scale case described in category 4.

7. Post-update cache coherence

This category is specifically about cache coherence across a model or adapter version change. Any affected cache must either be explicitly invalidated at update commit or include a model/adapter generation in its identity. General cache-correctness bugs without a weight update are out of scope.

Tracking	Failure	Status
#44950	A same-name tower/connector LoRA reload can reuse stale multimodal encoder-cache embeddings because the identifier lacks adapter-version identity; path fingerprinting fixes different-path reloads but not same-path content replacement	Open
#42125	Reloading a same-name LoRA reuses prefix-cache blocks from the previous adapter contents	Open

The cache inventory is feature-dependent and includes prefix KV, remote/offloaded KV, multimodal encoder/processor caches, LoRA caches, speculative draft hidden/KV state, and hybrid Mamba/SSM state. Supporting a new model means declaring which of these exist and proving that update commit either advances the cache generation or invalidates every affected entry.

Downstream reports

These are framework-side reports rather than vLLM core bugs. They remain evidence for API and CI requirements; resource-only issues are not promoted into an eighth correctness category.

Issue	Framework	Primary mapping	Bug
trl#5312	trl	Cross-RFC; vLLM root cause pending	Weights are not synchronized when vllm_enable_sleep_mode=True; also tracked by #48310 and must not be asserted as category 2 before a vLLM-side root cause is established
trl#6028	trl	6. Name/key mapping	AsyncGRPO VLM weight-sync key mismatch
verl#6256	verl	6. Name/key mapping	Qwen3-MoE LoRA KeyError during IPC weight sync
SkyRL#1776	SkyRL	Needs root-cause classification	LoRA training/inference mismatch for MoE models
verl#6468	verl	Related resource issue	CPU memory leak during FSDP2 rollout weight sync

Proposed CI architecture and coverage

A storage-address checker is intentionally narrow: by itself it directly covers only category 1. The broader proposal is a shared reload-correctness harness with five layers:

1. Offline source audit: enumerate post-load/kernel registries, run process_weights_after_loading twice with production-like copy-back, and inspect attributes, containers, callables, closures, backend objects, and expert objects. This discovers candidates; reachability is not proof that a tensor is graph-visible.
2. Production storage contract: use the single explicit registration design in #48478. Backends register graph-visible logical slots; capture snapshots their storage/layout; the common ModelRunner completion gate checks them after direct reload and WeightTransferEngine completion. A missing, expired, rebound, or layout-incompatible slot fails the update before resume. Production does not use TorchDispatchMode discovery or an arbitrary object walk.
3. Differential reload oracle: compare cold-load B against warm A → capture → reload B → replay the existing graph, using different deterministic weights. This detects stale values, obsolete loaders, premature finalization, and forward-visible reload-state corruption that pointer checks cannot see.
4. Specialized lanes: TP/PP/EP rank-distinct sentinels for routing, exact checkpoint/LoRA key accounting for name mapping, and warm-cache A → update B → cold-cache-B comparison for post-update cache coherence.
5. Composite RL transition lane: execute pause/drain → sleep → wake(weights) → update → post-load/finalize → wake(kv_cache) → cache invalidation → resume, then compare fixed-token logits/logprobs and liveness with a cold-B engine.

Prototype references: registry-wide post-load test and ReloadStorageManifest. They remain useful offline audit prototypes, but #48478 is the production design boundary.

The CUTLASS MoE family is a concrete audit driving case: tensors can live under layer.quant_method.moe_kernel.fused_experts, which is not an nn.Module. Offline scans must descend such backend/expert objects. Production support does not depend on finding them automatically; the owning backend must explicitly register them or warm reload with live graphs is unsupported and fails closed.

#	Coverage from the proposed system	Required lane
1. Storage identity	Direct; offline audit + explicit #48478 registry check + hardware replay	Generic identity lane
2. Runtime value refresh	Direct once A and B contain different sentinel values and derived values/output are compared	Differential reload lane
3. Loader lifecycle	Direct when the differential path uses the real loader before and after layout/kernel setup	Differential reload lane
4. Reload state preservation	Partial from output comparison; complete coverage requires explicit parameters/buffers/config/backend-state inventory	Differential state lane
5. Parameter routing/sharding	Not covered by a single-device generic run	TP/distributed sentinel lane
6. Name/key mapping	Not covered by pointer or output checks alone	Key-accounting + LoRA/checkpoint fixture lane
7. Post-update cache coherence	Not covered unless affected caches are deliberately warmed under model/adapter A and cache engagement is asserted	Cache-generation/invalidation lane

Coverage summary:

• The category-1 storage contract directly covers 1/7 categories; offline audit may provide incidental signal for category 4 but is not a production oracle.
• The common single-device differential harness can directly exercise categories 1–4 (4/7).
• With the TP, key-accounting, and cache-version lanes, the full CI system can cover 7/7 categories; this is a test framework with category-specific oracles, not one universal test case.

Exit criteria

• 1 — Storage identity: [Bugfix][Attention] Preserve post-load tensors across weight reloads #48251, [Bugfix][MoE] Preserve unquantized weight storage on ROCm #46009, Layerwise reload crashes with CUDA illegal memory access on compressed-tensors channel-wise FP8 MoE #41670, [Bugfix] Preserve Marlin runtime tensor storage across weight reload #48438, and [Bugfix] Preserve Machete act-order permutation storage across weight reload #48539 resolved; confirmed regressions (including FlashInfer CUTLASS expert constants) covered by graph replay CI; AITER MLA, B12x, NVFP4 scale chains, CUTLASS W4A8 strides, RDNA3 WNA16 scratch, TRT-LLM expert constants, and WNA8O8 candidates reproduced and fixed or cleared on supported hardware. Production enforcement follows [RFC] Fail-Closed Graph Storage Contract for Weight Reload #48478.
• 2 — Runtime value refresh: value-changing reload CI covers derived tensors, including FlashInfer attention sinks.
• 3 — Loader lifecycle: [Bug]: model.load_weights silently corrupts MoE forward on second call #42821/[Bugfix][Model Loader] Make model.load_weights safe to invoke on already-initialized model #42823 resolved with a second-load-after-layout-conversion regression test; [Bugfix] Fix layerwise reload dropping params after a composed weight loader #44814/[BUG] Exclude SKIP_TENSORS from get_layer_size() + new weight sync example for dpep #37334/[Bug] Add e_score_correction_bias to SKIP_TENSORS #38746 completion-accounting regressions have order/chunk/duplicate/empty-update coverage.
• 4 — Reload state preservation: [Bugfix] Fix layerwise reload alias-buffer corruption #42481, [Bugfix] Preserve unloaded non-persistent buffers during layerwise reload #44371, [Bugfix][MoE] Snapshot max_cudagraph_capture_size into FusedMoEConfig #44613, [Bugfix][RL] Set vLLM config during weight reload #45989, and Skip layerwise reload for non-quantized models #48284 covered by parameter, buffer, alias, configuration, and backend-state comparisons. Partial updates have an explicit preserve-or-reject policy.
• 5 — Parameter routing/sharding: [Bugfix] Re-sync parameter tp_rank after process_weights_after_loading (fix replicated / disable_tp weight reload) #48025 fixed with multi-rank sentinel tests.
• 6 — Name/key mapping: [Bugfix] Fix level-2 sleep/wake/reload with enable_lora=True #39935 fixed; checkpoint/model/LoRA keys and required auxiliary tensors have exact bidirectional accounting with explicit allowlists only.
• 7 — Post-update cache coherence: [Bugfix][V1] Invalidate encoder cache on finish_weight_update #48762 or an equivalent non-reverted fix lands; model, adapter, target, and draft updates either invalidate every affected cache or advance the generation used in cache identity.
• The composite sleep → update → wake → resume lane passes for each advertised model feature combination, jointly satisfying [RFC] Sleep/Wake Correctness for RL #48310 and this RFC.
• Weight-update regression coverage runs in the RL CI matrix for 30 days with no new correctness failure.

[RyanClark2k]

Re: high-risk candidate "Marlin workspace and sort indices"

I took this item. PR #48438 is ready for review with a fix, a CPU regression test, and completed GPU capture/reload/replay validation.

What I confirmed for this row, against main at 4c81772:

• MarlinLinearKernel.process_weights_after_loading rebinds self.workspace on every call (kernels/linear/mixed_precision/marlin.py). It lives on the kernel object, not the layer, so the parameter/buffer copy-back in model_loader/reload/layerwise.py never sees it. After a reload, captured CUDA graphs keep launching with the freed pointer. This applies to all Marlin models under CUDA graphs, not just act-order.
• For act-order GPTQ, layer.g_idx_sort_indices is a plain argsort tensor, never registered, so it has the same gap. The kernel reads it as a permutation, so stale memory means wrong outputs rather than scratch corruption.
• Checked and benign: the empty g_idx/w_zp/g_idx_sort_indices tensors in the non-act-order paths are zero-numel nn.Parameters and auto-register, and input_global_scale is registered and covered by copy-back.

The fix follows #48251: replace_parameter(prefer_copy=True) for the sort indices, and a new existing argument on marlin_make_workspace_new that reuses compatible storage.

Validation, RTX 4090, act-order GPTQ TinyLlama, CUDA graphs at defaults:

• Unfixed main: 88/88 workspace and 88/88 sort-indices pointers moved after reload_weights, and the engine then livelocked (100% GPU utilization for over 10 minutes on a generation that took 1.2s before the reload) once reclaimed workspace memory fed the Marlin kernels garbage sync counters.
• Fixed branch: 0/88 pointers moved, and flipping the live sort indices in place changed 8/8 generations, confirming the captured graphs read the storage the engine mutates.

Full method, logs, and the validation script are on the PR. One methodological note that may matter for the other rows in this table: a naive same-weights reload comparison passes on unfixed main, because the stale memory still holds bit-identical values until something reclaims it. Detecting this bug class needs pointer instrumentation or allocation pressure.

One finding beyond the row itself: the workspace pattern is not specific to MPLinearKernel. I audited the rest of the Marlin family and ran the same live pointer check on the reachable paths, on the same 4090 and on this PR's branch (which does not cover these sites):

• Live-confirmed rebinding on reload: the dense FP8 Marlin fallback (prepare_fp8_layer_for_marlin, marlin_utils_fp8.py:135, via MarlinFP8ScaledMMLinearKernel, 88/88 workspaces moved), the dense NVFP4 W4A16 Marlin fallback (prepare_fp4_layer_for_marlin, marlin_utils_fp4.py:243, via MarlinNvFp4LinearKernel, 88/88 moved), and the FP8 Marlin MoE path (prepare_fp8_moe_layer_for_marlin, marlin_utils_fp8.py:283, 6/6 moved on a small Mixtral-architecture model under online FP8 quantization with VLLM_TEST_FORCE_FP8_MARLIN=1). The dense fallbacks are the default paths for FP8 and NVFP4 checkpoints on GPUs without native support, so RL reload on Ampere and Ada hits them broadly.
• Static-confirmed, same pattern, not yet demonstrated live (the exercising checkpoints are large or rare): marlin_utils_fp8.py:466/552 (MXFP8 dense and MoE), marlin_utils_fp4.py:399/484 (NVFP4 and FP4 MoE), and a conditional rebind in compressed_tensors_moe_wna16_marlin.py:490. All sit in functions called from process_weights_after_loading.

Scope question: would you prefer the confirmed sibling sites folded into #48438, handled as a follow-up sweep, or tracked as their own row in the RFC? Happy to do whatever you prefer.

The investigation and implementation were AI-assisted. I reviewed the diff and the reload machinery myself and can defend the change end to end.

[aoshen02]

Hi hi, thanks for the thorough investigation! I think it makes sense to fold the confirmed sibling sites into #48438, together with the corresponding capture/reload/replay validation results.
More broadly, I’m considering adding systematic runtime check coverage for runtime tensors whose addresses are captured by CUDA Graphs, so we can catch storage-identity/address-drift bugs across weight reloads before they recur, any ideas about it?

[RyanClark2k]

Done — all eight sibling sites are fixed in #48438, and rather than leaving five of them at static analysis, every site now has live capture/reload/replay evidence on both sides of the fix (bug side rebinding, fixed side pointer-stable), on a rented 4090. Full per-site matrix, models, and methods: #48438 (comment).

On systematic coverage, I prototyped one concrete option: a unit test that runs each kernel's process_weights_after_loading twice — simulating what reload does, including the copy-back of registered params/buffers into original storage — and asserts that no tensor held by the layer or the kernel object changed address. It's parametrized over the kernel registry instead of written per site, so it needs no GPU or checkpoints for kernels whose post-load runs on CPU. It may be a useful addition, since on its first full run it caught a live instance of the bug class that my audit for #48438 missed:

FAILED test_post_load_runtime_tensors_stable[MacheteLinearKernel-act_order]
AssertionError: runtime tensors escape reload copy-back protection:
  moved={'MacheteLinearKernel.act_perm.kw[perm]': ('0x149b07fc0', '0x149b09e80')}

The cause: MacheteLinearKernel.process_weights_after_loading computes perm = torch.argsort(g_idx) fresh on every call and captures it in self.act_perm (machete.py:75-83) — unregistered, on the kernel object, used at apply time. Same violation as the g_idx_sort_indices case in #48438, but Machete is the preferred W4A16 kernel on Hopper, so act-order GPTQ + RL reload + CUDA graphs on H100s would hit it in a mainstream configuration. I'd suggest adding it as a row in the high-risk table (evidence: this test plus reading the code; it still wants a live capture/reload/replay confirmation on SM90+, which my Ada card can't provide). One implementation note for any check in this space: the first version of my pointer walk missed this because the tensor hides inside a partial/lambda rather than an attribute — the walk has to look inside callables too.

The test: https://github.com/RyanClark2k/vllm/blob/reload-storage-stability-test/tests/model_executor/model_loader/test_post_load_storage_stability.py (demo branch, not a PR). Coverage is partial and I want to be upfront about that: on GPU it would run every registry kernel that supports the canonical config; on CPU, kernels run natively only when their post-load is pure torch, or via small per-kernel stubs (Marlin needs 2, Machete 1), and skip otherwise — current CPU tally is 5 passed, 1 failed (the Machete catch), 22 skipped. Happy to iterate on it or turn it into a PR if you think it's useful.

[RyanClark2k]

Re-reading your question, the unit test above only half answers it: it's CI-shaped coverage over one kernel registry, and you asked about runtime checks for tensors whose addresses CUDA graphs captured. I prototyped that too: ReloadStorageManifest records a weak reference to every storage that could be baked into a graph, at the moment of capture, then after each reload asserts that no recorded storage was freed (a dangling graph address, always a bug) and no slot rebound (stale reads once weights change). Two recording modes cover each other's blind spots: an attribute walk of the model and kernel objects that also looks inside functools.partial bindings and closure cells, and a TorchDispatchMode that records every tensor argument ops touch while capture runs, which catches tensors reachable through no attribute path at all.

Storage lifetime is device-independent, so the mechanism is red/green tested on CPU: green on current Marlin, red when the pre-#48438 rebinding behavior is restored on the real kernel (both workspace and g_idx_sort_indices reported at the reload boundary), and the dispatch recorder catches a tensor referenced only from a module-level dict, which no walk can find. It also independently reports the unfixed Machete act_perm capture from my previous comment, encoded as a strict xfail asserting the desired invariant: it shows up as a known failure today, errors on the day a Machete fix lands so the marker has to come off, and then guards the fix. Both halves are cheap where it matters: recording overhead lands at capture time, which happens once at startup, and the per-reload check is a weakref sweep. So it could run raise-mode in CI's existing reload tests, where it would make every one of them discriminating for this bug class, and warn-mode in production RL loops.

What I deliberately haven't built is the engine integration: recording during the actual capture passes and checking inside reload_weights. The known gap there is that dispatch doesn't see inside compiled regions, so the walk-plus-dispatch union is the coverage story, and that needs validating against real capture on a GPU. Since you mentioned you're considering building something in this space, I'd rather align with your design before going further; happy to take the integration on if this shape looks right to you.

[aoshen02]

Re-reading your question, the unit test above only half answers it: it's CI-shaped coverage over one kernel registry, and you asked about runtime checks for tensors whose addresses CUDA graphs captured. I prototyped that too: ReloadStorageManifest records a weak reference to every storage that could be baked into a graph, at the moment of capture, then after each reload asserts that no recorded storage was freed (a dangling graph address, always a bug) and no slot rebound (stale reads once weights change). Two recording modes cover each other's blind spots: an attribute walk of the model and kernel objects that also looks inside functools.partial bindings and closure cells, and a TorchDispatchMode that records every tensor argument ops touch while capture runs, which catches tensors reachable through no attribute path at all.

Storage lifetime is device-independent, so the mechanism is red/green tested on CPU: green on current Marlin, red when the pre-#48438 rebinding behavior is restored on the real kernel (both workspace and g_idx_sort_indices reported at the reload boundary), and the dispatch recorder catches a tensor referenced only from a module-level dict, which no walk can find. It also independently reports the unfixed Machete act_perm capture from my previous comment, encoded as a strict xfail asserting the desired invariant: it shows up as a known failure today, errors on the day a Machete fix lands so the marker has to come off, and then guards the fix. Both halves are cheap where it matters: recording overhead lands at capture time, which happens once at startup, and the per-reload check is a weakref sweep. So it could run raise-mode in CI's existing reload tests, where it would make every one of them discriminating for this bug class, and warn-mode in production RL loops.

What I deliberately haven't built is the engine integration: recording during the actual capture passes and checking inside reload_weights. The known gap there is that dispatch doesn't see inside compiled regions, so the walk-plus-dispatch union is the coverage story, and that needs validating against real capture on a GPU. Since you mentioned you're considering building something in this space, I'd rather align with your design before going further; happy to take the integration on if this shape looks right to you.

Nice, we can talk about it on slack.

[github-actions]

CC @hongxiayang @tjtanaa @vllmellm for ROCm-related issue