How unregistered tensors break vLLM's RL weight reload

One bug class across five kernels, and why per-site patches won't hold.

Ryan Clark · July 2026 · PR #48438 · PR #48539 · RFC #48312

TL;DR

RL frameworks push new weights into a live vLLM engine every optimizer step. CUDA graphs bake device addresses at capture time, so vLLM's reload copies recomputed values back into the original storage. That copy-back only walks layer._parameters and layer._buffers, and kernels keep plenty of graph-visible tensors somewhere else.

I started on one row of the RL working group's audit and ended up with five confirmed kernels, two fix PRs, and a strong opinion about how to detect the rest. The obvious end-to-end test passes on broken code, so everything here rests on pointer identity instead.

RL frameworks push new weights into a running vLLM engine after every optimizer step. TRL, verl, and SkyRL all call some version of reload_weights on a live engine rather than tearing it down.

CUDA graphs make that hard. A captured graph doesn't run your Python. It replays a recorded sequence of kernel launches, and the capture baked every device address into that sequence. Move a tensor after capture and the graph keeps reading the address where the tensor used to live.

vLLM handles this. After reprocessing a layer, model_loader/reload/layerwise.py copies the recomputed values back into the original storage so the captured pointers stay valid. The copy-back walks layer._parameters and layer._buffers. Anything a kernel keeps outside those two dicts gets no protection.

Everything below comes out of that gap. I picked up one unclaimed row of RFC #48312, the working group's tracker for weight-reload correctness, and I've since fixed two kernels, confirmed three more instances, and spent most of my time arguing about how to find the ones nobody has looked at yet.

Marlin, and a test that passed when it shouldn't

MarlinLinearKernel.process_weights_after_loading creates two tensors that escape the copy-back. It allocates a scratch workspace for cross-block sync counters and stores it on the kernel object, which has no _parameters dict at all. For act-order GPTQ it computes a sort permutation with torch.argsort and assigns it to layer.g_idx_sort_indices. nn.Module.__setattr__ registers nn.Parameter assignments and ignores plain tensors, and argsort returns a plain tensor. One line away, the non-act-order path assigns the result of marlin_make_empty_g_idx, which returns a Parameter and registers fine. The two assignments look interchangeable in a diff, and a return type three calls away decides which one the registry ever hears about.

I rented a 4090 and wrote the obvious end-to-end check first. Load an act-order GPTQ TinyLlama with CUDA graphs at defaults, generate greedily, reload the same weights through the real RL path, generate again, compare token ids. It passed on unfixed code.

Three things were covering for the bug.

Output comparison can't catch a memory-identity bug whose contents are value-identical. So I checked pointer identity instead. llm.apply_model runs a function against the live model inside the engine, so I recorded data_ptr() for every workspace and every sort-indices tensor, reloaded, and recorded again. On main, 88 of 88 workspaces and 88 of 88 sort indices moved.

I added a second probe for the permutation. Reverse every live sort-indices tensor in place after the reload, then generate again. Correct code reads the storage the engine just mutated, so the output has to change. Broken code reads the old memory and the flip does nothing. The probe gives a positive signal either way, and it doesn't depend on winning a race with the allocator.

The flip also ended the run. tensor.flip(0) allocates a temporary per layer, 88 of them, and the caching allocator served those out of the freed workspace blocks. Marlin's kernel spins until its sync counters reach an expected value, and the counters now held permutation indices. A generation that had taken 1.2 seconds before the reload ran for ten minutes at 100% GPU utilization before I killed it. In a real RL run that's the difference between a metrics blip and a rollout fleet that stops making progress at full power draw, with nothing in any log.

On the fixed branch, 0 of 88 moved and the flip changed 8 of 8 generations.

The same pattern, eight more times

While reading Marlin I'd noticed what looked like the same unregistered-workspace pattern at ten other sites, and I posted that finding. Then I tested it, and two of the ten didn't survive. auto_awq.py and auto_gptq.py assign their workspace in create_weights, which runs once at construction and never again on reload. I had read the line without reading the function around it, and I retracted both in the same thread.

That left eight. The RFC author asked me to fold all of them into the PR with validation, so I ran every one through the same capture, reload, replay census. Every site rebound on the unfixed commit, and every site held at 0/N on the fix. One row read 18/30 rather than 30/30 on the bug side, because the allocator happened to hand back 12 same-size blocks at their old addresses, which is why the fixed side has to reach 0/N by construction instead of by luck.

Some rows had no exercising checkpoint. vLLM's own CI keeps tiny random-weight DeepSeek-V3 clones that covered three of them, and for the last one llm-compressor quantized a 10 MB random Qwen3-MoE clone, data-free, into the mxfp4-pack-quantized format that path demands. Lifecycle bugs don't need good models to reproduce.

A test over the whole registry

Every per-site check here has one shape. Build a minimal layer, run the kernel's post-load twice, replay the production copy-back in between, assert that no reachable tensor moved. The RFC author asked about systematic coverage, so I parametrized that shape over the mixed-precision kernel registry instead of writing it once per site. Kernels whose post-load is pure torch run on a laptop CPU, and two small stubs cover the device-only ops.

It failed on its first full run, on a kernel I had never audited.

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')}

Machete recomputes torch.argsort on every post-load call and captures the result inside self.act_perm, a functools.partial. Unregistered, on the kernel object, read at apply time. Machete is the preferred W4A16 kernel on Hopper, so act-order GPTQ plus RL reload plus CUDA graphs hits it on H100s in a mainstream configuration.

My first version of the walk missed it. The tensor hides in the partial's keyword arguments, and an attribute walk sees a callable and moves on. The walk has to unwrap partials and closure cells before this kind of capture becomes visible.

I rented an H100 to confirm it, and the failure mode turned out worse than Marlin's. 88 of 88 act_perm tensors rebound, and the allocator freed 0 of 88 capture-time storages, because compile artifacts still retain the old tensors. So the graphs keep permuting activations with the previous model's act order while the rest of the model updates. No crash, no allocation pressure that would produce one, and no same-weights comparison that could ever see it. The fix registers the permutation and reads it off the layer at apply time.

That probe also rejected my first attempt at the fix. I registered the permutation but still bound it into act_perm during post-load, which runs before the copy-back, so the callable held a detached alias. The registered parameter held at 0/88 while the callable's tensor moved 88/88 and the flip changed 0 of 8 generations. Only late binding through the layer reads the storage the engine refreshes.

Past the linear kernels

The same H100 rental turned up a fifth site with a different shape. FlashInferExperts allocates SwiGLU constants (gemm1_alpha, gemm1_beta, gemm1_clamp_limit) in its __init__, and post-load rebuilds the entire experts object. On gpt-oss-20b with moe_backend=flashinfer_cutlass, post-load rebuilt 24 of 24 experts objects, 72 of 72 tensors moved, and the allocator freed 72 of 72 capture-time storages. Dangling addresses rather than stale reads, so once something reuses that memory the captured graphs feed garbage constants into every expert GEMM.

These constants never change across reloads. No comparison of warm reload against cold load can catch them for any choice of weights, however you pick the two checkpoints. Only pointer identity or storage lifetime finds them. Three sibling classes in the TRT-LLM backends have the same shape and need SM100 hardware I don't have.

Reaching that configuration took four attempts. flashinfer's JIT gates -DENABLE_FP4 on a CUDA toolkit of 12.8 or newer, the system nvcc was 12.4, so fp4 compiled out silently and the engine died at startup with an unrelated-looking error.

Why per-site fixes won't hold

I stopped believing in per-site fixes somewhere around the fifth confirmed instance. Every one has the same root cause, every fix is the same three lines, and the audit keeps finding more. Three efforts upstream now aim at the class instead of the instances.

aoshen02's #48478 turns registration into a fail-closed contract, so a kernel that allocates a graph-visible tensor without registering it fails at load rather than silently at inference. new-TonyWang's #48902 walks each layer recursively, finds CUDA tensors on nested objects including partials and closures, and extends the copy-back to preserve their addresses. His #49789 goes much further and treats a weight update as a transaction, with a layer-owned ReloadArena holding graph-visible derived tensors and a manifest, recorded during the first real load, that decides whether a reload actually finished.

Reviewing the walker, I found that its path resolver getattrs each segment, so it silently skips any tensor reachable only through a dict, partial, or closure, and that its moe_kernel is None guard breaks the upstream layerwise path for kernels that rename parameters during post-load. A CPU repro against the real reload machinery shows both directions.

The transaction PR is a 5,800-line design, and I'd change three things in it. Its per-layer arena check runs five lines after the copy-back that publishes the layer, so a violation can't stop the publish. Move the check above the copy-back and refuse on findings and it becomes a gate, and I ran the reorder locally against their own per-layer arena tests to confirm the earlier position still catches everything. Its strict completion mode also can't work as written, because required_keys comes from observing a full load, so a deliberately partial update produces missing keys that are correct. You'd fix that with an explicit update scope rather than a stricter check.

The design doc asks who should own the transaction, and you can settle that from where the caches live. The model runner owns the encoder and multimodal caches and resets both at the end of reload. The scheduler owns the KV cache manager, so reload_weights can't reach prefix blocks at all, and nothing in any weight-update path invalidates prefix KV today. After an update the engine can still serve blocks that it computed under the previous policy. Whatever flushes those has to sit above the worker, which puts the coordinator on the caller side along with the per-rank report check.

My own argument is about the CI check rather than the repair. A lint built from the collector's own walker inherits the walker's blind spots, so the tensor the collector misses is also the tensor the lint misses. You want an oracle that can grade any repair mechanism instead of assuming one.

So I prototyped ReloadStorageManifest. It records a weak reference to every storage a graph could bake in, at the moment of capture, then after each reload checks that every recorded storage is still alive at the same address. Two recording modes cover each other's blind spots. An attribute walk that unwraps partials and closure cells, and a TorchDispatchMode that records every tensor argument flowing through an op while capture runs, which catches tensors reachable through no attribute path at all. I have a CPU test where a tensor held only in a module-level dict is invisible to any walk and the dispatch recorder finds it anyway. Storage lifetime behaves the same on any device, so the whole mechanism red/greens on CPU, and the per-reload check costs one weakref sweep.

Where it stands

The whole investigation ran on rented GPUs, a 4090 at $0.669/hr for the Marlin work and an H100 at $1.752/hr for Machete and FlashInfer. Both livelocks were on the bill.