[Oleg-Goncharov]

Add explicit CUDA state lifecycle APIs to prevent reuse of cached TE resources across in-process restarts. Release resources while CUDA remains usable, or invalidate supported native resources after externally managed context loss, and recreate them lazily.

Framework cleanup preserves model tensors and FP8 numerical history. Unsupported framework recovery and unsafe communication/JIT teardown fail explicitly.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Add release/reset APIs, per-device generations, a global atomic notification, and resource registration callbacks.
• Integrate streams/events, library handles, RTC modules, normalization/attention plans, and context-local kernel initialization. Retain handles from finished threads in a reusable ownership pool.
• Invalidate Workspace plans, TMA descriptors, and persistent contents; release current/retired PyTorch workspace allocations and rebuild them on reuse.
• Integrate PyTorch/JAX cleanup, device-scoped tensor caches, FP8 registrations, and CuTeDSL/TVM-FFI executor retirement.
• Integrate communication cleanup and reject unsafe teardown or stale EP buffers.
• Document lifecycle ordering and limitations; add native, frontend, JIT, and two-rank communication tests.

Validation and remaining issues

• Focused validation: 154 passed, 2 skipped.
• Eight existing Flex Attention cases fail with both the updated and original HEAD tests, without requiring cleanup: softcap/post-scale-bias numerical and layout errors.
• NCCL EP and cuBLASMp syntax checks pass; runtime validation remains outstanding.
• Full L0/CI and restart-orchestrator integration have not been run.
• Initialized PyTorch/XLA context-loss recovery and leak-free post-loss cleanup of opaque vendor host state remain unsupported.
• Dependency/compiler warnings and cuDNN nanobind shutdown leak diagnostics remain visible.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

Unchecked items require full L0/CI validation, resolution or triage of the reported failures/warnings, and review of the remaining recovery limitations. The contributing guidelines have been read; their full validation requirements are not yet satisfied

[timmoon10]

The scope of this PR is a bit larger than what we'd like. From what I've seen here, I don't trust Codex's code quality and I'd like a human to properly go through and vouch that they fully understand it. I wonder if we could undo all the API changes by making a minor tweak: #3779 (comment)

The core infrastructure for resource invalidation is pretty nice though. I have some suggestions to make it even more cute, especially #3779 (comment).

[timmoon10]

Is there value in forcing all workspaces to use the same infrastructure? Workspace requirements are kernel-dependent and alignment management is not that challenging.

[timmoon10]

Nit: These size_t overflow checks scattered through the code don't really accomplish much. We'll never approach sizes that large unless we are already in a weird buggy state.

[timmoon10]

Nit: Why does alignment need to be a power of 2? It's not worth putting extra effort to support non-powers-of-2, but also it's not worth putting extra effort disallowing it.

[timmoon10]

Why not allow smaller than 128 byte alignment? If this is general workspace infrastructure, I could imagine situations like storing a few floats or ints in the workspace.

[timmoon10]

Why does T need to be trivially copyable? Isn't it valid to initialize the workspace by calling a constructor within a kernel? In fact, don't we do something like that for TMA descriptors?

[timmoon10]

Is it worth the effort to support per-device resets? I don't see why we couldn't have resources that are shared between devices, e.g. the cuDNN plan cache could have kernels for multiple devices, or the NVRTC kernel could have compiled code on CPU and per-device function pointers. It seems easier to support the desired use-case by blindly destroying everything that depends on the CUDA context.

Suggested change

	void nvte_reset_cuda_state(int device);
	void nvte_reset_cuda_state();

[timmoon10]

Is there value in tracking the resources on each device?

Suggested change

	struct DeviceState {
	std::atomic<uint64_t> generation{1};
	std::set<Resource *> resources;
	};
	namespace {
	struct Registry {
	std::mutex mutex;
	std::map<int, DeviceState> devices;
	std::atomic<uint64_t> generation{1};
	};
	namespace {
	struct Registry {
	std::mutex mutex;
	std::vector<Resource *> resources;
	std::atomic<uint64_t> generation{1};
	};

[timmoon10]

We could make registration automatic by putting this in the Resource constructor.

[timmoon10]

I wonder if there could be problems from resetting resources in a non-deterministic order. It doesn't seem like we need to support removing resources, so perhaps it would be safer to reset resources in reverse the order they were registered.

Suggested change

	std::set<Resource *> resources;
	std::vector<Resource *> resources;

[timmoon10]

Handling invalidation logic with a mixin class is a nice approach. As I see it, this class has the following contract:

• Class instances register themselves with a global registry with bind_context.
• The internal state of a Resource can be invalidated at any time. Any time you access a Resource, you need to check whether the state is invalidated and reinitialize if needed.

I wonder if we can make this more automatic:

• Class instances automatically register themselves in global registry (https://github.com/NVIDIA/TransformerEngine/pull/3779/changes#r4418682371).
• Have virtual functions for initialization and invalidation.

class Resource {
 public:
  Resource() {
    auto &state = registry();
    std::lock_guard<std::mutex> lock(state.mutex);
    state.resources.push_back(this);
  }
  void ensure_valid() {
    if (!valid_.load(std::memory_order_acquire)) {
      initialize();
      valid_.store(true, std::memory_order_release);
    }
  }

 private:
  friend void reset();
  virtual void initialize() = 0;
  virtual void invalidate() noexcept = 0;
  std::atomic<bool> valid_{false};
};

void reset() {
  auto &state = registry();
  std::lock_guard<std::mutex> lock(state.mutex);
  for (auto it = state.resources.rbegin(); it != state.resources.rend(); ++it) {
    if (it->valid_.load(std::memory_order_acquire)) {
      it->invalidate();
      it->valid_.store(false, std::memory_order_release);
    }
  }
  state.generate.fetch_add(1, std::memory_order_release);
}

The only requirement for subclasses is to implement the virtual methods and to call ensure_valid() before accessing state.