fix(gms): rebind non-parameter tensors to private memory after publish

[galletas1712]

Summary

Fixes fatal NVRM Xid 31 (FAULT_RO_VIOLATION ACCESS_TYPE_VIRT_WRITE) GPU faults when a GMS publisher engine writes to non-parameter tensors after publish — first hit in practice on all ranks of a Nemotron NVFP4-MoE (fp8 KV, DP8) deployment at the first wake after snapshot restore. The same fault is latent for plain sleep/wake on any --kv-cache-dtype fp8 model.

The problem

The publisher builds the entire model inside the GMS weights memory pool, so the committed allocations contain more than parameters: torch buffers (expert_map, Mamba conv_weights, ...) and tensor attributes (fp8 KV scales _k_scale/_v_scale/_q_scale/_prob_scale, quantization ranges) are swept into the same caching-allocator segments. That sweep is intentional — a fresh importer never runs load_weights, so everything the forward pass references must be materializable from GMS.

The bug is that publish semantics are uniform: after finalize_gms_write (and after any RO reconnect — wake, restore), the whole tag is remapped CU_MEM_ACCESS_FLAGS_PROT_READ. But some captured tensors are written after load:

• post_kv_cache_wake_up → init_fp8_kv_scales does fill_() on the fp8 scale attrs on every wake;
• fp8 range attrs are updated on the forward path.

A write through a PROT_READ mapping is a fatal MMU fault. Cold engines only avoid it by timing: every write lands inside the RW load window, and the late writers first run after a sleep/wake — or after a snapshot restore, which is how this was found. Fault forensics: Xid 31 on all ranks at first post-restore wake, fault addresses at the tail 4 KiB page of 2 MiB small-pool segments (fp8 scale blocks at the segment tail).

The fix

materialize_module_from_gms already gives importers the right binding: parameters bind to the shared read-only mapping; buffers/tensor attrs are detach().clone()d into ordinary CUDA memory. This PR gives the publisher the identical rule:

• rebind_nonparameter_tensors (new, client/torch/module.py): re-binds every GMS-resident non-parameter tensor (buffers, tensor attrs, list/tuple elements) on the publisher's model to a private detach().clone(); parameters stay on the shared read-only mapping. Returns bytes rebound.
• finalize_gms_write calls it after register_module_tensors + commit + RO remap (so GMS metadata still records the in-pool copies and importer materialization is unchanged), and logs the rebound MiB so per-model duplication stays visible.

Consequences:

• late writers hit ordinary CUDA memory — preserved byte-exact and VA-stable by cuda-checkpoint/CRIU in the snapshot flow;
• the weights tag is honestly PROT_READ for every consumer, including snapshot-restored publishers;
• the server's committed copy can never be mutated.

Notes / constraints

• Runs before CUDA graph capture (clones live at new addresses; captured graphs bake raw pointers). finalize_gms_write executes at model-load time in all three integrations (vLLM/TRT-LLM/SGLang), which precedes capture.
• Classification is deliberately blanket — every non-parameter is treated as mutable, because proving immutability per-buffer is impractical. The cost is the clone bytes, logged at finalize.
• Known residual caveat (same semantics importers already have): two attribute names aliasing one tensor become two independent clones after rebind.
• The medium-term direction (a misc pool for mutable non-parameter state, converged with KV sleep/wake semantics) is deliberately not in this PR; it will be proposed separately as an issue/DEP rather than a checked-in design doc.

Validation

• New end-to-end regression test test_finalize_gms_write_rebinds_nonparameter_tensors in lib/gpu_memory_service/tests/test_torch_integration.py: drives the real finalize_gms_write flow (register → commit → RO reconnect → remap → rebind) against a live GMS server, asserts the parameter keeps its shared GMS binding while buffer/tensor-attr move to private memory with values preserved, then writes to the rebound tensors after the RO remap — the exact operation that faulted.
  • Verified it fails (assert not _in_gms(gms_model.scale)) when the rebind_nonparameter_tensors call is stubbed out, and passes with the fix.
  • python -m pytest lib/gpu_memory_service/tests/test_torch_integration.py → 5 passed on a 2×RTX 5880 Ada box.
• pre-commit run --all-files → all hooks pass (black/isort/flake8/ruff/pytest-marker-report).
• Stack-level E2E (branch schwinns/vllm-snapshot-nightly-overlay-gms): post-restore inference returns 200 with finish_reason=stop, coherent output, no Xid 31, restored weights left read-only. This commit is patch-equivalent to experiment commit 8ef1ab30fb from that validated image.

AI assistance was used for this PR (review-feedback cleanup and test restructuring); all changes reviewed by the author.

[devin-ai-integration]

Devin Review found 4 potential issues.

[coderabbitai]

Walkthrough

Adds a rebind_nonparameter_tensors function that clones GMS-resident non-parameter tensors into private writable CUDA memory and rebinds them onto the model. Integrates this into finalize_gms_write after RO remap, updates logging/docstrings, documents the rationale, and adds tests.

Changes

Non-parameter tensor rebind

Layer / File(s)	Summary
Core rebind implementation lib/gpu_memory_service/client/torch/module.py	Adds rebind_nonparameter_tensors, which walks a model for CUDA non-parameter tensors, checks pointers against GMS mappings, clones in-pool tensors into private writable storage, rebinds buffers/attributes/list elements, and returns total duplicated bytes.
Publisher integration and docs lib/gpu_memory_service/integrations/common/utils.py, lib/gpu_memory_service/docs/mutable-state-and-memory-classes.md	Imports and calls rebind_nonparameter_tensors in finalize_gms_write after RO remap, updates docstring/log message with rebound byte counts, and adds documentation explaining the RO remap fault scenario, the rebind mitigation, and longer-term design plans.
Rebind test coverage lib/gpu_memory_service/tests/test_module_rebind.py	Adds CUDA-gated tests with _Layer fixture and _fake_manager helper validating selective tensor rebinding, value/writability preservation, untouched non-mapped tensors, and idempotency across repeated calls.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Related PRs: None specified.

Suggested labels: gpu-memory-service, enhancement, documentation, tests

Suggested reviewers: Reviewers familiar with GMS memory allocation and torch module internals.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description is detailed, but it does not follow the required template and is missing the required Related Issues section.	Reformat the PR body to match the template, adding Overview, Details, Where should the reviewer start?, and a completed Related Issues section.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title is concise and accurately summarizes the main change: rebinding non-parameter tensors to private memory after publish.

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

lib/gpu_memory_service/integrations/common/utils.py (1)

123-139: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider surfacing rebound_bytes in GMSCommittedMemoryStats.

The rebind byte count is only logged, not returned. Callers such as the TRT-LLM loader that capture finalize_gms_write(...).committed_bytes for metrics have no way to also track rebind overhead.

♻️ Optional: add field to stats

 class GMSCommittedMemoryStats:
     committed_bytes: int
     pruned_bytes: int
+    rebound_bytes: int

     return GMSCommittedMemoryStats(
         committed_bytes=int(total_bytes),
         pruned_bytes=int(pruned_bytes),
+        rebound_bytes=int(rebound_bytes),
     )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@lib/gpu_memory_service/integrations/common/utils.py` around lines 123 - 139,
The commit stats currently log rebound_bytes but do not expose it to callers, so
add a corresponding field to GMSCommittedMemoryStats and populate it from
finalize_gms_write in the common utils flow. Update the return path where
GMSCommittedMemoryStats is constructed so it includes the rebound byte count
alongside committed_bytes and pruned_bytes, and ensure any callers relying on
finalize_gms_write can access the new field for metrics.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@lib/gpu_memory_service/client/torch/module.py`:
- Around line 269-273: The tuple-handling branch in _iter_module_tensors
currently skips rebinding for immutable containers, leaving GMS-resident tensors
from tuple attributes read-only. Update the logic around the attr.isdigit() /
not isinstance(mod, torch.nn.Module) path so tuple elements are rebound to
private/writable storage just like list elements, and only keep the debug skip
for truly non-rebindable cases that are not tensor containers. Use the
_iter_module_tensors helper and its container rebinding branch to locate the
change.

In `@lib/gpu_memory_service/tests/test_module_rebind.py`:
- Around line 16-18: The tests in test_module_rebind.py only use the
requires_cuda skipif marker and are missing required repo markers. Update the
test definitions in this module to add a Lifecycle marker such as
pytest.mark.pre_merge, a test type marker such as pytest.mark.unit, and the
appropriate GPU hardware marker used by this repo (for example gpu_0), alongside
the existing requires_cuda guard. Ensure the marker additions are applied to
both test cases referenced in this file so they satisfy marker validation in CI.

---

Nitpick comments:
In `@lib/gpu_memory_service/integrations/common/utils.py`:
- Around line 123-139: The commit stats currently log rebound_bytes but do not
expose it to callers, so add a corresponding field to GMSCommittedMemoryStats
and populate it from finalize_gms_write in the common utils flow. Update the
return path where GMSCommittedMemoryStats is constructed so it includes the
rebound byte count alongside committed_bytes and pruned_bytes, and ensure any
callers relying on finalize_gms_write can access the new field for metrics.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: ff752cad-6d23-49d1-a37f-e8d4828427d0

📥 Commits

Reviewing files that changed from the base of the PR and between b9c5bd7 and 1900d35.

📒 Files selected for processing (4)

• lib/gpu_memory_service/client/torch/module.py
• lib/gpu_memory_service/docs/mutable-state-and-memory-classes.md
• lib/gpu_memory_service/integrations/common/utils.py
• lib/gpu_memory_service/tests/test_module_rebind.py

[github-actions]

🌿 Fern Docs Preview: https://nvidia-preview-019f2574-ffd7-735d-8925-0d1955932743.docs.buildwithfern.com/dynamo/dev

[galletas1712]

Pushed d782a99 addressing all review feedback:

• Tuple containers (CodeRabbit): now rebound by rebuilding the tuple on its owning attribute — previously a GMS-resident tensor in a tuple stayed on the read-only mapping.
• Clone-before-skip-checks (Devin): clone is now only allocated on branches that bind it.
• Test markers / _deps pattern (Devin, CodeRabbit): the standalone fake-manager test file is gone. Coverage moved into test_torch_integration.py as a true end-to-end test that drives finalize_gms_write against a live GMS server and then writes to the rebound buffer/tensor-attr after the RO remap — the exact operation that faulted with Xid 31. Verified locally that the test fails when the rebind call is stubbed out, and passes with it (5/5 in the file on GPU).
• Design doc: dropped docs/mutable-state-and-memory-classes.md; the problem/fix rationale now lives in the PR description, and the medium-term misc-pool plan will be filed separately as an issue rather than a checked-in doc that would go stale.
• rebound_bytes in GMSCommittedMemoryStats (CodeRabbit nitpick): declining for now — no caller consumes it and the finalize log line keeps the duplication observable per model; happy to add the field when a metrics consumer exists.
• pre-commit run --all-files passes (this was the red CI job — black reformats).