Debugging pointer identity in vLLM's RL weight-reload path: why the obvious test passes on broken code, the instrument that caught it, and the livelock that confirmed it harder than I intended.
RL training frameworks push fresh weights into a live vLLM engine every step. CUDA graphs bake device pointers at capture time. vLLM's reload machinery preserves storage addresses for registered parameters and buffers, and silently protects nothing else. Marlin's quantized-GEMM kernel keeps two runtime tensors outside that protection: a workspace on the kernel object and, for act-order GPTQ, the g_idx_sort_indices permutation. After one reload, every captured graph is reading freed memory. The obvious A/B test (generate, reload identical weights, generate, compare) passes on the broken code, because stale memory still holds bit-identical values until something reclaims it. An instrumented test caught it: 88/88 tensor addresses move on reload on main, and a write-to-live-state probe shows the graphs ignore the tensors the engine believes it has. Then the probe's own temporary allocations reclaimed the freed workspace blocks and the next generation, 1.2 s before the reload, spun the GPU at 100% for ten minutes.
Fix is ~30 lines; the same audit found the pattern at eight more sites and retracted two of my own claims along the way. All eight are now fixed and live-validated, and a generalized version of the test caught the same class in a second kernel on its first run. Total GPU bill: about $3 of rented 4090.
vLLM has a code path that RL frameworks exercise constantly and almost nobody else touches. During policy training, the rollout engine has to receive updated weights after every optimizer step without tearing the engine down: TRL, verl, SkyRL and the rest all call some flavor of reload_weights on a running engine. This is also the path where CUDA graphs make everything harder. A captured graph does not run your Python. It replays a recorded sequence of kernel launches with every device address baked in at capture time. Your Python objects can rebind, reallocate, and drift; the captured graph keeps launching against the addresses it saw on capture day.
vLLM's reload machinery (model_loader/reload/layerwise.py) knows this. After reprocessing a layer's weights it copies the recomputed values back into the original tensor storage, so the pointers the graphs captured stay valid. The docstring on the helper spells out the invariant: preserving data_ptr is required for captured CUDA graphs to remain valid across weight updates. There is one catch, and the whole bug lives inside it: the copy-back walks layer._parameters and layer._buffers. A tensor that is not registered in one of those two dicts gets no protection at all.
The vLLM RL working group had been auditing exactly this class of bug (RFC #48312, "Weight Reload Correctness for RL"), and their table of unverified candidates listed the Marlin kernel's workspace as a suspect. I picked up that row. Marlin is the mixed-precision GEMM kernel behind GPTQ, AWQ, and several FP8/FP4 fallbacks, so anyone doing RL on a quantized policy with CUDA graphs enabled (the default) is standing on this path.
Reading MarlinLinearKernel.process_weights_after_loading against the reload machinery, two tensors escape the protected set.
The first is the workspace. Marlin allocates a small scratch buffer for cross-block synchronization counters and stores it as self.workspace on the kernel object. A kernel object is a plain Python class. It has no _parameters dict, no module registry, no path by which the copy-back could ever learn the tensor exists. Reprocessing allocates a fresh workspace, the old one's refcount drops to zero, and the allocator gets the block back. The captured graphs keep the address.
The second is act-order specific. GPTQ with desc_act=true sorts weight groups by activation order, and the kernel needs the sort permutation at inference time. The code computes it with torch.argsort and assigns it with layer.g_idx_sort_indices = g_idx_sort_indices. Here is the detail that makes this whole bug class slippery: whether that line registers anything depends on the type of the right-hand side. nn.Module.__setattr__ intercepts nn.Parameter assignments and registers them; a plain tensor lands as an ordinary object attribute and the registry never hears about it. argsort returns a plain tensor. One line away, the non-act-order path assigns marlin_make_empty_g_idx(device), which returns an nn.Parameter, so that path auto-registers and is fine. Two assignments that look interchangeable in a diff, with different runtime lifecycle semantics, decided by a return type three calls away.
The workspace is write-only scratch, so its failure mode is corrupting whatever tensor the allocator hands the block to next. The permutation is worse in a quieter way: the kernel reads it, so stale or reclaimed memory there means wrong outputs with no error anywhere.
The fix follows the invariant the machinery already states: recompute values into the same storage. About 30 lines total, following the pattern of a sibling fix in the same RFC campaign. The sort indices go through replace_parameter(..., prefer_copy=True), which both copies in place and registers the tensor so the copy-back protects it from then on. The workspace gets an existing argument on its factory function: reuse compatible storage, zero it, return it. A CPU unit test pins both invariants with data_ptr assertions, needs no GPU (the two CUDA-only ops are monkeypatched), and fails on main exactly where it should. So far, a tidy afternoon.
Static analysis plus a unit test is an argument. I wanted the bug demonstrated in a real engine with real captured graphs before claiming it, so I rented a 4090 on Vast for $0.669/hr and wrote the obvious end-to-end check: load an act-order GPTQ model (TinyLlama, desc_act=true, all 88 linear modules on the Marlin act-order path), CUDA graphs at defaults, generate greedily on 8 prompts, reload identical weights through the real RL path, generate again, compare token ids. Broken code should diverge.
[ok ] prompt 0: 'The capital of France is'
[ok ] prompt 1: 'In machine learning, overfitting means'
...
RESULT: PASS (all generations identical after reload)That is the unfixed code passing the test that its bug is supposed to fail. It took me a minute to believe it, and most of an hour to fully explain it. Three separate effects were covering for the bug:
The general statement, and the thing I will carry to other systems: a memory-identity bug with value-identical contents is invisible to output comparison by construction. Absence of divergence tells you nothing. This is presumably exactly why the bug survived in tree: any reasonable smoke test of the reload path reloads the same weights and compares, and it passes every time.
The fix for a test that can be fooled is to test the invariant instead of the symptom. Two probes, neither of which depends on allocator luck:
Pointer census. vLLM's llm.apply_model runs a function against the live model inside the engine. Walk the modules, record data_ptr() for every Marlin workspace and every sort-indices tensor, reload, record again. The graphs baked the before-pointers; if the after-pointers differ, the graphs are provably reading memory the engine no longer owns, whatever the contents happen to be.
The flip. Proving the graph reads freed memory directly is awkward. Proving the graph ignores the live tensor is easy: write to it and see if anything changes. After the reload, reverse every live sort-indices permutation in place (reversed indices are still in range, so this is memory-safe, just numerically wrong) and generate again. On correct code the graphs read the storage the engine just mutated, so generations must change. On broken code the graphs read the old memory and the flip is a no-op. The probe is symmetric: it produces a positive signal in both the healthy and the broken case, just opposite ones. There is no analogous probe for the workspace, and the reason is instructive: sort indices have legitimately-shaped wrong values you can write, while any nonzero value in a sync-counter buffer is garbage, and on the fixed code the graph reads the live workspace by design, so poisoning it breaks the healthy configuration too. The workspace gets the pointer census only.
On main:
pointer identity after reload: 88/88 workspaces moved, 88/88 sort indices moved
flipped 88 live sort-indices tensors in placeEvery single tracked tensor rebound. And then the run stopped producing output.
The generation after the flip had taken 1.2 seconds before the reload. Ten minutes later the log had not advanced and nvidia-smi read 100% utilization, 22 GB resident, no errors. I killed it after confirming it was not merely slow.
The failure signature does the diagnosis for you, which is the part of this story I enjoy most. If the graphs had been reading garbage permutations, the outcome would be wrong tokens or an illegal memory access, and either of those terminates. An infinite spin at full utilization is a thread waiting on a condition that will never come true. Marlin's workspace holds cross-block synchronization counters; the kernel spins until a counter reaches its expected value. Garbage counters, infinite spin. And where did the garbage come from? The flip itself. tensor.flip(0) allocates a temporary per layer, 88 of them, and the caching allocator served those temporaries out of the freed workspace blocks, filling the memory the captured graphs were about to read with permutation indices. The probe I built to detect the sort-indices half of the bug supplied the allocation pressure that detonated the workspace half.
I will take it. In an actual RL run this is the difference between a metrics blip and a rollout fleet that stops making progress at full power draw, with nothing in any log. The naive test says the code is fine. The instrumented test hangs the GPU.
On the fixed branch, the same script, same model, same everything:
pointer identity after reload: 0/88 workspaces moved, 0/88 sort indices moved
flipped 88 live sort-indices tensors in place
generations changed after live flip: 8/8
VERDICT: FIXED - storage identity preserved across reload; captured graphs read live storageAddresses stable, and the flip lands: the graphs are demonstrably coupled to the storage the engine mutates. One commit of difference between this block and the livelock.
While reading the code I had noticed the same unregistered-workspace pattern at what looked like ten more sites across the Marlin family: the FP8 and FP4 fallback paths, several MoE variants, AWQ and GPTQ MoE. I posted that as a finding. Then, since the 4090 was still rented, I tested it, and the claim did not survive contact intact.
| site | verdict | how |
|---|---|---|
FP8 dense Marlin (marlin_utils_fp8.py:135) | rebinds, 88/88 | live: online-FP8 TinyLlama; this is the default path on GPUs without native FP8 |
NVFP4 dense fallback (marlin_utils_fp4.py:243) | rebinds, 88/88 | live: NVFP4 TinyLlama; default path on pre-Blackwell hardware |
FP8 Marlin MoE (marlin_utils_fp8.py:283) | rebinds, 6/6 | live: a tiny random-weight Mixtral. A pointer census doesn't care if the model is any good; it needs the architecture, not the training |
| MXFP8 + FP4 MoE + one conditional (5 sites) | static only, at the time; see day two below | the exercising checkpoints are tens of GB or barely exist publicly |
AWQ MoE (auto_awq.py:664) | retracted | workspace assigned in create_weights, which never re-runs on reload |
GPTQ MoE (auto_gptq.py:651) | retracted | same |
The retractions came from a distinction my first pass glossed over: where the assignment lives decides everything. A workspace assigned in create_weights is assigned once, at model construction, and reload never touches it; the identical-looking line inside process_weights_after_loading re-runs on every reload. Ten shrank to eight, three of the eight are now demonstrated in a live engine, and I posted the corrections in the same thread as the original claim. Testing your own findings before someone else does is cheap. Doing it while the GPU is already rented is nearly free.
One trick from this phase worth keeping: VLLM_TEST_FORCE_FP8_MARLIN=1 plus on-the-fly FP8 quantization of a small bf16 model reaches the FP8 Marlin paths on hardware that would otherwise route elsewhere, and a random-weight architecture clone (dacorvo/Mixtral-tiny) reaches the MoE machinery for pennies. Lifecycle bugs don't need good models to reproduce.
The maintainer's answer came back within hours: fold all eight sites into the PR, with the corresponding validation. That turned the five static rows into a checkpoint-hunting problem. Three fell to artifacts that already existed; vLLM's own test suite keeps tiny random-weight DeepSeek-V3 clones with real quantized variants, and online mxfp8 quantization reaches the MXFP8 paths with no checkpoint at all. The last two rows had no exercising checkpoint anywhere public.
The random-weights trick generalizes further than I first used it. llm-compressor will quantize a 10 MB random Qwen3-MoE clone, data-free, into the exact mxfp4-pack-quantized format the unreachable path demands. Fifteen seconds on CPU produced a 2 MB checkpoint, and the engine log printed the familiar fallback warning as it routed into prepare_moe_fp4_layer_for_marlin. If the exercising artifact doesn't exist, you can manufacture it.
Final matrix, all on one 4090: eight sites, every one rebinding on the unfixed commit, every one at 0/N on the fix. Two rows earned footnotes. The MXFP4 MoE census read 18/30 moved on the bug side; the caching allocator had handed 12 same-size blocks back at their old addresses. Any movement proves the rebind, but this is the masking story again in miniature: bug-side address stability can be allocator luck, so the fixed side has to show 0/N by construction, and it does. And the MXFP8 MoE path reproduced the livelock. Post-reload generation pinned the GPU at 100% again, and this time the wedged process shrugged off SIGTERM and had to be killed by PID. The failure signature replicates across kernels, which is what you want if the diagnosis is right.
Every per-site check in this story reduces to one shape: build a minimal layer, run the kernel's post-load processing 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 per site. Twenty-eight cases fall out of fourteen kernels. Kernels whose post-load is pure torch run on a laptop CPU, and two small stub adapters cover the device-only ops in Marlin and Machete. On its first full run it failed 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's post-load computes a fresh torch.argsort permutation on every call and captures it in self.act_perm, inside a functools.partial. Unregistered, on the kernel object, read at apply time: the sort-indices bug in a different kernel family. Machete is the preferred W4A16 kernel on Hopper, so act-order GPTQ reloads on H100s sit right on it. It is unfixed as I write this; I reported it as a new row on the RFC, and confirming it live needs SM90 hardware my rented Ada card can't impersonate.
The detail I keep thinking about is that the first version of the test missed it. act_perm is a callable, and the tensor hides in the partial's keyword arguments; an attribute walk sees a function and moves on. The walk has to unwrap partials and closure cells before this kind of capture becomes visible. Day one's lesson about registration semantics, recurring one level of indirection deeper. On the 4090 the same test runs five kernels against their real device ops (Marlin passes there without any monkeypatching, which retires the last asterisk on the fix), and Machete stays skipped on that card until someone points a Hopper at it.
module.attr = tensor and module.attr = Parameter(tensor) diverge silently in every system that walks the module registry: reload, state_dict, device moves. The type of the right-hand side is doing invisible work.