[RFC] Recursive Tensor Collector: Fix CUDA Graph Invalidation After Weight Reload

[new-TonyWang]

Related Issue: vllm#48312 (Category 1: Storage Identity)
Related PRs: #48251, #48438, #48539, #41670

Summary

We propose a generic, infrastructure-level fix for CUDA graph invalidation caused by unmanaged tensor address drift during reload_weights(). The solution recursively discovers all CUDA tensors reachable from each nn.Module layer -- including those on nested Python objects like kernel instances, quant methods, and functools.partial -- and extends the existing copy-back mechanism to preserve their device addresses across reloads.

This eliminates an entire class of bugs (13 confirmed/high-risk sites) without requiring changes to any individual kernel, quantization scheme, or attention implementation.

Background: Lifecycle of Model Loading, CUDA Graph Capture, and Weight Reload

End-to-end timeline

The following diagram shows the chronological order of all key events. Understanding this order is essential to understanding the bug:

Time ─────────────────────────────────────────────────────────────────────►

[1] Initial Load              [2] CUDA Graph Capture        [3] Normal Inference
load_model()                  (lazy, on first inference)    (graph replay)
├─ load_weights()             ├─ CUDAGraphWrapper           ├─ .replay()
├─ process_weights_after_     │  .__call__()                ├─ .replay()
│  loading()                  │  torch.cuda.graph():        ├─ .replay()
│  creates workspace @addr_A  │   records forward pass      │  ...
│  creates W_UV    @addr_B    │   burns in addr_A, B, C     │
│  creates strides @addr_C    │   into the graph            │
│                             │                             │

                                                    RL framework updates weights
                                                            │
                                                            ▼

[4] reload_weights()                                [5] Resume Inference
├─ initialize_layerwise_reload()  ← inside reload   ├─ .replay()
│  ├─ snapshot old tensors (addr_A, B, C)            │  still uses addr_A, B, C
│  ├─ collect extra_tensor_slots                     │  but now with new values
│  ├─ move params to meta device                     │
│  └─ wrap weight_loaders                            │
├─ model.load_weights(new_weights)                   │
│  └─ per-layer _layerwise_process():                │
│     ├─ materialize (allocate new GPU memory)        │
│     ├─ load new weights                             │
│     ├─ process_weights_after_loading()              │
│     │  creates NEW workspace @addr_D                │
│     │  creates NEW W_UV     @addr_E                 │
│     ├─ copy-back:                                   │
│     │  addr_A.copy_(new workspace val)  ← existing  │
│     │  addr_B.copy_(new W_UV val)       ← OUR FIX   │
│     │  re-attach old tensors to layer                │
│     │  (addr_D, addr_E tensors discarded)           │
├─ finalize_layerwise_reload()                        │
└─ reset_encoder_cache()                              │

Key chronological order:

[1] load_model  →  [2] CUDA Graph Capture  →  [3] Inference  →  [4] reload_weights  →  [5] Inference

• [1] happens once at engine startup. process_weights_after_loading() creates derived tensors (workspace, W_UV, strides, etc.) at addresses A, B, C.
• [2] happens lazily on the first inference request for each batch size. torch.cuda.graph() records the forward pass and burns in addresses A, B, C -- the graph will always read/write these exact addresses on replay.
• [3] is the steady state. Each inference call just does entry.cudagraph.replay() -- no Python, no tensor lookups, just replaying recorded GPU operations at recorded addresses.
• [4] is triggered by the RL framework after training updates weights. initialize_layerwise_reload() is called inside reload_weights(), not as a separate step. After reload, process_weights_after_loading() creates new tensors at addresses D, E -- but the copy-back mechanism writes new values into old addresses A, B, C and discards D, E.
• [5] resumes inference. Graph replay reads addresses A, B, C. If copy-back worked correctly, these contain the new values. If any address was missed, the graph reads stale/freed/garbage memory.

Where CUDA graphs are captured

CUDA graphs are captured lazily by CUDAGraphWrapper.__call__() (compilation/cuda_graph.py:233-361). The first time a given batch size is seen during inference, the wrapper records the entire forward pass:

# compilation/cuda_graph.py:265-337 (simplified)
class CUDAGraphWrapper:
    def __call__(self, *args, **kwargs):
        entry = self.concrete_cudagraph_entries[batch_descriptor]

        if entry.cudagraph is None:
            # First time for this batch size -> capture
            cudagraph = torch.cuda.CUDAGraph()
            with torch.cuda.graph(cudagraph, pool=self.graph_pool, stream=...):
                output = self.runnable(*args, **kwargs)  # record all CUDA ops
            entry.cudagraph = cudagraph
            return output

        # Subsequent calls -> replay the recorded ops directly on GPU
        entry.cudagraph.replay()
        return entry.output

During capture, torch.cuda.graph() records every CUDA kernel launch and burns in the data_ptr() of every tensor accessed by those kernels. On replay, the GPU executes the recorded kernel sequence directly -- no Python code runs, no tensor lookups happen. The GPU just reads/writes the exact memory addresses that were recorded.

Critical design decision: reload_weights() does NOT clear or re-capture CUDA graphs. The reload_weights() code path (gpu_model_runner.py:5514-5518) has no call to CUDAGraphWrapper.clear_all_graphs() or any re-capture logic. This is intentional -- re-capturing graphs is expensive and would negate the performance benefit of hot-reload in RL loops.

This means the entire burden of correctness falls on the copy-back mechanism: after reload, every tensor address that a CUDA graph references must still be valid and contain the correct (updated) values. If any address drifts, the consequences depend on what happened to the old memory:

Scenario	Symptom
Old address reused by PyTorch allocator for another tensor	Silent wrong output (reads unrelated data)
Old address still allocated but holds stale values	Silent wrong output (reads old weights)
Old address freed back to OS / CUDA driver	illegal memory access (process crash)
Old address was Marlin workspace (atomic counter)	Garbage counter value -> kernel infinite wait -> livelock (confirmed: 100% GPU for >10 min on RTX 4090)

What does layerwise reload do? (step [4] in detail)

reload_weights() (gpu_model_runner.py:5468) calls three functions in sequence:

initialize_layerwise_reload(model)          # prepare
loaded_weights = model.load_weights(iter)   # load
finalize_layerwise_reload(model, config)    # finalize

The name "layerwise" means it processes one layer at a time to minimize peak GPU memory (avoids having both old and new weights for the entire model simultaneously).

initialize_layerwise_reload() (layerwise.py:84-126) -- For each nn.Module layer:

1. Snapshot kernel tensors (line 112): Save references to current _parameters and _buffers -- these hold the addresses that CUDA graphs burned in.
2. Collect extra tensor slots (line 119): Recursively find unmanaged CUDA tensors on the layer (our addition).
3. Restore to meta device (line 122): Move all params/buffers to meta device (zero GPU memory), freeing space for new weights.
4. Wrap weight loaders (line 125): Replace each param's weight_loader with a wrapper that buffers incoming weights.

model.load_weights() -- Weights stream in one tensor at a time. The wrapper buffers them. When all weights for a layer are ready (load_numel >= load_numel_total), it triggers _layerwise_process() (layerwise.py:361-401):

1. Materialize (line 372): Allocate real GPU memory for this layer
2. Load (line 384-387): Copy buffered weights into the new GPU tensors
3. Post-process (line 390-392): Run process_weights_after_loading() -- this creates new derived tensors at new addresses
4. Copy-back (line 396-398): Copy new values into old addresses, restore references

finalize_layerwise_reload() (layerwise.py:234-306) -- Handles Attention/MLA layers (deferred until all linear layers are done) and layers with padding.

The copy-back mechanism and its gap

The critical function is _copy_and_restore_kernel_tensors() (layerwise.py:417-433):

def _copy_and_restore_kernel_tensors(layer, info):
    parameters, buffers = info.kernel_tensors  # saved in Phase 1

    for name, param in parameters.items():
        param.data.copy_(getattr(layer, name))  # new values -> old address
    for name, buffer in buffers.items():
        buffer.data.copy_(getattr(layer, name))  # new values -> old address

    _place_kernel_tensors(layer, info)  # re-attach old tensors to layer

This ensures that after reload, layer.qweight.data_ptr() is the same address CUDA graph captured -- but with updated values. The problem is that it only covers layer._parameters and layer._buffers.

Motivation

The Problem

process_weights_after_loading() creates derived tensors -- workspaces, stride descriptors, fused scales, permutation indices -- stored as plain Python attributes on nn.Module, kernel objects, or quant method instances. These are invisible to _parameters/_buffers and therefore invisible to the copy-back mechanism:

layer (nn.Module)
|-- _parameters: {qweight, scales, g_idx}        <- copy-back covers these
|-- _buffers: {}                                   <- copy-back covers these
|-- g_idx_sort_indices (bare Tensor)               <- MISSED
|-- W_UV, W_UK_T (bare Tensor)                     <- MISSED
`-- quant_method (Python object)
    `-- kernel (Python object)
        `-- workspace (Tensor)                     <- MISSED
    `-- moe_kernel (Python object)
        `-- fused_experts (Python object)
            |-- ab_strides1 (Tensor)               <- MISSED
            |-- ab_strides2 (Tensor)               <- MISSED
            `-- c_strides1 (Tensor)                <- MISSED

After reload, these tensors get new addresses. CUDA graph replay reads the old (now stale or freed) addresses, causing:

• Illegal memory access (segfault/CUDA error)
• Silent wrong outputs (reading stale data at old addresses)
• Graph replay livelock (Marlin workspace used as thread-block synchronization counter; garbage value causes infinite wait -- confirmed on RTX 4090: 100% GPU utilization for >10 minutes)

Scale of the Problem

We identified 13 affected sites across the codebase through static audit + live verification:

#	Component	Affected Tensors	Severity
1	Generic MLA (#48251)	W_UV, W_UK_T	Wrong output
2	AITER MoE shuffle (#40390)	w13_weight, w2_weight	Wrong output (fixed)
3	ROCm MoE padding (#46009)	.data storage swap	Wrong output
4	CUTLASS FP8 MoE (#41670)	ab_strides*, c_strides*	Illegal memory access
5	Marlin workspace + sort indices (#48438)	workspace, g_idx_sort_indices	Livelock / wrong output
6	8 Marlin siblings (FP8/FP4/MoE variants)	workspace in each	Livelock
7	AITER MLA FP4/FP8	W_K, W_V, *_scale	Wrong output (audit)
8	FlashInfer B12x MoE	_fc2_input_scale, w*_sf_mma	Wrong output (audit)
9	Compressed-tensors NVFP4 CUTLASS MoE	g1/g2_alphas, a*/gscale	Wrong output (audit)
10	CUTLASS W4A8 FP8 MoE	a_strides*, c_strides*, s_strides*	Illegal memory access (audit)
11	RDNA3 WNA16 MoE	rdna3_w1_buf, rdna3_act_buf	Wrong output (audit)
12	Machete act-order	act_perm	Wrong output (audit)
13	Compressed-tensors WNA8O8	cloned scales	Wrong output (audit)

All 13 sites share the same root cause: tensors created by process_weights_after_loading() that live outside _parameters/_buffers. Fixing them one by one is a game of whack-a-mole; any new kernel or quant method can introduce the same bug. We need a systematic solution.

Proposed Solution

Design

Extend the layerwise reload pipeline with a recursive tensor collector that discovers all CUDA tensors reachable from each layer's attribute tree, and a copy-back extension that restores their addresses after reload.

                     Current Pipeline              Enhanced Pipeline
                  +------------------+          +------------------+
Snapshot          | _parameters      |          | _parameters      |
(before reload)   | _buffers         |          | _buffers         |
                  |                  |          | + bare attrs     |
                  |                  |          | + nested objects  |
                  +------------------+          +------------------+
                         |                            |
PWAL creates new tensors (same)                      (same)
                         |                            |
Copy-back         | params & buffers |          | params & buffers |
(after reload)    | only             |          | + ALL tracked    |
                  |                  |          |   tensors        |
                  +------------------+          +------------------+

Key Components

1. tensor_collector.py -- Recursive Discovery

def collect_extra_tensors(layer: nn.Module) -> list[TensorSlot]:
    """
    Recursively walk layer's attribute tree, collecting all CUDA tensors
    NOT already in _parameters/_buffers.

    Traversal covers:
    - Plain tensor attributes on the module (Type A)
    - Tensors on nested Python objects: kernel, quant_method, experts (Type B)
    - Tensors inside dicts, lists, tuples
    - Tensors captured in functools.partial args/kwargs
    - Tensors captured in function closures

    Does NOT recurse into child nn.Modules (they have their own reload cycle).
    Deduplicates by storage id (handles shared-storage aliases).
    Max depth = 8 (safety bound; real paths are ~4-5 levels deep).
    """

Each discovered tensor is recorded as a TensorSlot:

@dataclass
class TensorSlot:
    path: str              # e.g. "quant_method.kernel.workspace"
    tensor: torch.Tensor   # strong ref (prevents GC of old address)
    data_ptr: int          # address at capture time
    storage_id: int        # for deduplication

2. path_resolver.py -- Object Graph Navigation

def resolve_path(root, path: str) -> torch.Tensor | None:
    """Navigate 'quant_method.kernel.workspace' from root object."""

def set_by_path(root, path: str, value: torch.Tensor) -> bool:
    """Set attribute at path to value, returning success."""

Handles dotted paths, dict/list indexing (some_dict['key'], some_list[0]).

3. Integration into layerwise.py

Three surgical modifications:

# 1. Snapshot phase: collect extra tensors alongside params/buffers
def initialize_layerwise_reload(model):
    for layer in model.modules():
        info = get_layerwise_info(layer)
        info.kernel_tensors = get_layer_params_buffers(layer)
        info.extra_tensor_slots = collect_extra_tensors(layer)  # NEW
        ...

# 2. Restore phase: copy-back extra tensors after PWAL
def _copy_and_restore_kernel_tensors(layer, info):
    # ... existing param/buffer copy-back ...
    _copy_back_extra_tensors(layer, info)  # NEW

# 3. New function: extra tensor copy-back
def _copy_back_extra_tensors(layer, info):
    for slot in info.extra_tensor_slots:
        new_tensor = resolve_path(layer, slot.path)
        if compatible(slot.tensor, new_tensor):  # shape + dtype check
            slot.tensor.data.copy_(new_tensor)           # new values -> old address
            set_by_path(layer, slot.path, slot.tensor)   # old tensor -> object graph

Why This Works Even When Entire Objects Are Replaced

Consider CUTLASS FP8 MoE where process_weights_after_loading() creates an entirely new CutlassExpertsFp8 instance:

self.moe_kernel = make_fp8_moe_kernel(...)  # brand new object

The copy-back still works because:

1. slot.tensor holds a strong reference to the old tensor -- it won't be GC'd
2. resolve_path() navigates to the new object's attribute (new stride tensor)
3. old_tensor.data.copy_(new_tensor) writes new values into the old address
4. set_by_path() sets the new object's attribute to point to the old tensor

Result: the new CutlassExpertsFp8 object's ab_strides1 now points to the old address with new values. CUDA graph replay reads correct data from the unchanged address.

Properties

• Zero changes to kernel/quant/attention code: Pure infrastructure fix in reload/
• Forward-compatible: New kernels that create derived tensors are automatically covered
• Safe defaults: Shape/dtype mismatches are logged and skipped (no crash), with a signal for graph recapture
• Minimal memory overhead: Extra tensor slots hold strong refs to small tensors (workspaces, strides, indices); negligible vs model weights
• Storage deduplication: Shared-storage aliases (e.g., c_strides2 = ab_strides1) are handled correctly

Validation Status

Verified: H200 End-to-End Model Reload Tests (20 PASS)

These tests perform real model loading, weight reload, and generation on H200 GPUs (rl_learning cluster). They are the ground truth for correctness.

Category 1: Storage Identity -- Address Preservation After Reload

Test	Model	Result
FP8 MoE strides	DeepSeek-V3-debug, FP8_DYNAMIC	PASS (19 extras tracked, 0 drifted)
MLA W_UV/W_UK_T	DeepSeek-V3-debug, BF16	PASS (6 extras tracked, 0 drifted)
Machete W4A16	Qwen3-0.6B-W4A16-G128	PASS (114 extras tracked, 0 drifted)
BF16 baseline	Qwen3-0.6B, BF16	PASS (30 extras tracked, 0 drifted)
FP8 online quant	Qwen3-0.6B, FP8	PASS (114 extras tracked, 0 drifted)

Category 2: Runtime Value Refresh -- Correct Generation After Reload

Test	Model	Result
BF16 mul/add perturbation	Qwen3-0.6B	PASS
W4A16 mul/add perturbation	Qwen3-0.6B-W4A16-G128	PASS
FP8 mul/add perturbation	DeepSeek-V3-debug FP8	PASS
BF16 mul/add perturbation	DeepSeek-V3-debug BF16	PASS

Category 3: Loader Lifecycle -- MoE Kernel Identity Across 2 Consecutive Reloads

Test	Model	Result
FP8 MoE 2-reload + generation	DeepSeek-V3-debug FP8	PASS
BF16 MoE 2-reload + generation	DeepSeek-V3-debug BF16	PASS

Category 4: Reload State Preservation -- Upstream vLLM Test Suite

Test	Result
9 upstream test_reload.py CPU tests (via pytest on H200)	9/9 PASS

Ongoing: Unit Tests with Mock Layers (13 tests, preliminary PASS)

These tests use mock nn.Module layers with synthetic tensors to validate the collector/copy-back logic in isolation. They verify the mechanism's correctness at the code level, but do not constitute end-to-end validation with real models and CUDA graphs. Results are preliminary; end-to-end validation with real models is currently in progress.

Category 1: Storage Identity (mock)

Test	What It Validates	Preliminary	Blocker
W4A8 stride collection	collect_extra_tensors finds stride tensors on mock CutlassExperts-like objects	PASS	No compressed-tensors W4A8 MoE checkpoint on HuggingFace for E2E
W4A8 stride copy-back	copy_back_extra_tensors preserves address on mock layer	PASS	Same
Stride drift detection (negative)	Detects address change when copy-back is NOT applied	PASS	Same
Config mismatch (negative)	Handles different quant scheme objects gracefully	PASS	Same

Category 2: WNA8O8 Scale Preservation (mock)

Test	What It Validates	Preliminary	Blocker
WNA8O8 scale discovery	Collector finds _input_scale/_output_scale on mock scheme	PASS	No WNA8O8 model on HuggingFace for E2E
WNA8O8 scale copy-back	Address preserved after simulated reload	PASS	Same
No-copy-back drift (negative)	Confirms drift without the fix	PASS	Same
Wrong-dtype detection (negative)	Rejects shape/dtype mismatch safely	PASS	Same

Category 3: Loader Lifecycle (mock)

Test	What It Validates	Preliminary
Guard preserves kernel id	MoE kernel rebuild guard skips rebuild on reload	PASS
Guard bypass rebuilds kernel	Guard allows rebuild when flag is cleared	PASS
WITH vs WITHOUT guard monkeypatch	Side-by-side comparison of guarded vs unguarded behavior	PASS

Category 4: State Preservation (mock)

Test	What It Validates	Preliminary
Reload-cycle buffer corruption (negative)	Detects NaN injection in non-persistent buffers	PASS
Unloaded buffer metadata preservation	Buffer shape/dtype survives meta-device round-trip	PASS

Pending: Hardware / Model Checkpoint Not Available

These sites are confirmed by static code audit to have the same root-cause pattern. They require specific hardware or model checkpoints not currently available for end-to-end testing:

Site	Blocker	Current Coverage
FlashInfer CUTLASS MoE	Requires Blackwell SM100+	Static audit; same collect_extra_tensors path
FlashInfer B12x scales	Requires SM120+	Static audit
AITER MLA FP4/FP8	Requires AMD ROCm + HIP Graph	Static audit; same attribute pattern
RDNA3 WNA16 MoE scratch	Requires gfx1100 + HIP Graph	Static audit
Machete act-order (CUDA graph)	Live graph capture/replay on SM90 pending	Collection verified on H200; graph replay not yet tested
ROCm MoE padding (#46009)	.data storage swap; requires ROCm	Existing param copy-back may suffice; needs live verification

Summary

Tier	Count	Confidence
H200 end-to-end (real model reload + generation)	20 PASS	High
Mock unit tests (mechanism validation, includes W4A8/WNA8O8)	13 preliminary PASS	Medium -- pending E2E with real models
Hardware/model not available	6 sites	Low -- static audit only

What This Does NOT Fix

Category	Issue	Why Not Covered
Cat 5	Parameter routing/sharding (TP rank mismatch)	Logic bug in weight_loader, not an address issue
Cat 6	Name/key mapping (LoRA prefix mismatch)	Naming convention issue
Cat 7	Post-update cache coherence	Requires cache invalidation, orthogonal to tensor addresses
Cat 2 (partial)	Non-tensor derived state (e.g., Python scalars, config flags)	Collector only tracks torch.Tensor; scalar state needs separate handling

Files Changed

vllm/model_executor/model_loader/reload/
  tensor_collector.py   # NEW: recursive tensor discovery
  path_resolver.py      # NEW: object graph navigation
  types.py              # MODIFIED: +TensorSlot, +extra_tensor_slots field
  layerwise.py          # MODIFIED: 3 integration points (~40 lines added)

No changes to:

• Any process_weights_after_loading() implementation
• Any kernel code (Marlin, CUTLASS, FlashInfer, etc.)
• CUDA graph capture/replay logic
• Model definition code
• FX pass system

Alternatives Considered

Approach	Pros	Cons
register_buffer for each site	Simple per-fix	Whack-a-mole; doesn't cover kernel objects; every new quant method can regress
StorageManifest (detect-only)	Good for CI	Detects but doesn't fix; still need per-site patches
Recursive collector + copy-back (this RFC)	Generic, forward-compatible, fixes all sites	Slightly more complex; path resolution adds code
Disable CUDA graph on reload	Simplest	Unacceptable perf regression for RL workloads
Full model rebuild on reload	Correct by construction	Defeats the purpose of hot-reload; GPU memory churn

Open Questions

1. Should we add a needs_recapture flag? When _copy_back_extra_tensors encounters shape/dtype mismatches, it currently logs a warning. Should it also set a flag that tells the engine to recapture the CUDA graph?

2. Closure traversal depth: Currently we traverse function.__closure__ cells. This handles known cases (Machete act_perm), but aggressive closure walking could have edge cases with third-party libraries. Should we make this opt-in?

3. Upstream CI integration: The collect_extra_tensors() function can also serve as a CI lint -- run it after initial model loading and after reload, and assert zero address drift. Should we propose this as a standard CI check for new kernels?

Discussion: Pre-Capture Pinning vs Per-Reload Collection

The idea

Instead of recursively traversing the model on every reload, we could traverse once before CUDA graph capture, copy all unmanaged tensors into a fixed memory region, and update all references to point there. Then reload would only need to copy new values into the known fixed addresses -- no traversal, no path resolution.

Current approach:                    Pre-capture pinning:

[1] load_model                       [1] load_model
[2] CUDA graph capture               [1.5] traverse all tensors -> pin to fixed pool -> update refs
[3] inference                        [2] CUDA graph capture (burns in fixed pool addresses)
[4] reload                           [3] inference
    |- traverse & collect slots      [4] reload
    |- load + PWAL                       |- load + PWAL
    |- copy-back via path resolve        |- copy new values into fixed pool (no traversal)
    |- traverse & write-back

Why reload still needs copy-back even with pre-pinning

process_weights_after_loading() runs during every reload and creates new tensors on new objects:

# During reload, PWAL still executes:
self.moe_kernel = make_fp8_moe_kernel(...)  # brand new Python object
# new_object.ab_strides1 -> addr_D (new address, not the fixed pool)

After PWAL, the new objects don't know about the fixed pool. So we still need to:

1. Copy new values from addr_D into the fixed pool at addr_A
2. Update the new object's reference to point back to addr_A

This is fundamentally the same copy-back operation. The saving is that we skip the recursive traversal -- we already know which paths to look at.

What CAN be optimized: persistent slot mapping

The model structure doesn't change between reloads. The paths (quant_method.kernel.workspace, W_UV, etc.) are the same every time. So the real optimization is:

# Collect once (after first load, before first capture):
info.extra_tensor_slots = collect_extra_tensors(layer)  # O(N) object graph walk

# Reuse on every reload (skip re-collection):
def initialize_layerwise_reload(model):
    for layer in model.modules():
        info = get_layerwise_info(layer)
        # DON'T re-collect: info.extra_tensor_slots already populated
        # Just snapshot the current tensor references for copy-back
        for slot in info.extra_tensor_slots:
            slot.tensor = resolve_path(layer, slot.path)  # O(1) per slot

This reduces reload cost from O(N) traversal to O(k) path lookups, where k = number of tracked slots (typically 10-100 per layer).

Fixed memory pool: further optimization

Going further, we could allocate a contiguous CUDA buffer and pack all unmanaged tensors into it:

class TensorPinPool:
    def __init__(self, model):
        all_slots = []
        for layer in model.modules():
            all_slots.extend(collect_extra_tensors(layer))

        total_bytes = sum(s.tensor.nbytes for s in all_slots)
        self.pool = torch.empty(total_bytes, dtype=torch.uint8, device='cuda')

        offset = 0
        for slot in all_slots:
            pool_view = self.pool[offset:offset+slot.tensor.nbytes]
                .view(slot.tensor.dtype).reshape(slot.tensor.shape)
            pool_view.copy_(slot.tensor)
            set_by_path(root, slot.path, pool_view)  # replace original with pool view
            slot.pinned_view = pool_view
            offset += slot.tensor.nbytes

    def copy_back_after_reload(self, model):
        for slot in self.slots:
            new_tensor = resolve_path(root, slot.path)
            slot.pinned_view.copy_(new_tensor)          # new values -> fixed address
            set_by_path(root, slot.path, slot.pinned_view)  # re-attach pool view

Practical challenges for the fixed pool approach

Challenge	Detail
No pre-capture hook	vLLM captures CUDA graphs lazily on first inference (CUDAGraphWrapper.__call__). There's no explicit "about to capture" event to hook into. The pool setup would need to happen after first process_weights_after_loading() but before first forward pass.
Non-contiguous tensors	transpose/permute produce tensors with non-trivial strides. Packing them into a contiguous buffer changes their layout, which may break kernels that depend on specific memory layout.
Mixed with existing copy-back	_parameters/_buffers still use the existing copy-back mechanism. Running two parallel address-preservation systems increases complexity.
Dynamic shapes	If a different quantization config produces different-shaped derived tensors, the pool size is wrong. (Rare in practice, since quant config doesn't change between reloads.)

Comparison

Approach	Traversal per reload	Path resolution per reload	Complexity	When to adopt
Current: per-reload collection	O(N)	O(k)	Low	Now (this RFC)
Persistent slot mapping	0	O(k)	Low	Follow-up optimization
Fixed memory pool	0	O(k)	Medium	High-frequency RL loops
Skip PWAL on reload	0	0	High (change all kernels)	Long-term ideal

Recommendation

Ship the current per-reload collection approach first (this RFC). Then pursue persistent slot mapping as a follow-up -- it's a small change (don't clear extra_tensor_slots in info.reset(), update tensor refs instead of re-collecting) that eliminates the recursive traversal cost. The fixed memory pool is a further optimization that requires a pre-capture hook in vLLM's compilation pipeline.

References

• #48312 -- [RFC] Weight Reload Correctness for RL (RyanClark2k)
• #48251 -- MLA W_UV/W_UK_T fix
• #48438 -- Marlin workspace fix (all 8 siblings)
• #48539 -- Related reload fixes
• #41670 -- CUTLASS FP8 MoE stride drift
• #40390 -- AITER MoE shuffle fix
• #46009 -- ROCm MoE padding

[github-actions]

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

[RyanClark2k]

On open question 3: I'd argue the CI check shouldn't be built from collect_extra_tensors itself. A lint that uses the same walker as the fix inherits the same blind spots, so the tensor the collector misses is also the tensor the lint misses. The oracle you want is a record of what the graphs actually captured, taken at capture time. I prototyped that shape on reload-storage-manifest: record storage weakrefs of graph-visible tensors at capture, then after reload assert every one is still alive at the same address. Since it's independent of any repair mechanism, it can grade the collector (and the registry, and per-site fixes) instead of assuming them. Happy to adapt it for the CI piece if that's useful.

On the alternatives table: how does this compose with #48478? The registry contract isn't in the table, but your copy-back warning already points migrations at it, so I assume the intent is collector as safety net and migration detector, registry as the eventual fail-closed contract. Worth stating in the RFC either way, since a reader today can't tell whether the two RFCs stack or compete.

On open question 1: yes to needs_recapture, and I'd go further. A shape or dtype mismatch means the graph is reading dead values, so warn-and-continue shouldn't be a terminal state; recapture or raise, same philosophy as the completion gate in #48478.

The guard discussion on #48902 also bears on the "Why This Works Even When Entire Objects Are Replaced" section, since if that argument holds the PWAL guards shouldn't be needed at all, but that's better continued on the PR.

[new-TonyWang]

This RFC is now outdated. Going forward, we will refer to #48902 (comment) and #48902 (comment).