Skip to content

[MoonEP] Destroy communication buffers on runtime reset - #35834

Draft
musclemuller wants to merge 2 commits into
sgl-project:moonepfrom
wirybeaver:codex/moonep-runtime-finalizers
Draft

musclemuller wants to merge 2 commits into
sgl-project:moonepfrom
wirybeaver:codex/moonep-runtime-finalizers

Conversation

@musclemuller

@musclemuller musclemuller commented Aug 21, 2026

Copy link
Copy Markdown

Motivation

Part of #35783.

MoonEP communication buffers own NVLink/VMM resources and already expose explicit destruction, but reset_context() currently drops the runtime resource store without invoking that cleanup. This draft makes reset/test teardown deterministic for the communication-buffer path already present on the moonep branch.

This is intentionally the first narrow lifecycle PR. It does not cover the symmetric weight pool from #34874, observability, autotuning, model reload, production graceful shutdown, or GPU/multi-rank teardown validation.

Modifications

  • Add a named finalizer registry to runtime Resources.
  • Run finalizers once in LIFO order during reset_context().
  • Reject duplicate finalizer names.
  • Continue cleanup after individual failures, install fresh runtime resources, then raise an aggregate error chained from the first failure.
  • Register MoonEPBuffer.destroy_all_buffers() when the process-local MoonEP buffer state is first created.
  • Add unit coverage for finalizer ordering, uniqueness, failure handling, fresh reset state, and destruction of multiple MoonEP capacity buffers exactly once.

Accuracy Tests

Not applicable. This change does not alter model outputs or the dispatch/combine data path.

Speed Tests and Profiling

Not applicable. Cleanup runs only from reset_context(); no inference hot path was changed.

Tests

  • test/registered/unit/test_runtime_context.py: 72 passed, 35 subtests passed. The local environment used an import-only stub for the unrelated optional datasets dependency during collection.
  • test/registered/unit/layers/moe/test_moonep_buffer.py: 9 passed.
  • Black, isort, Ruff (F401,F821,UP037), git diff --check, and scripts/ci/check_registered_tests.py passed.

GPU and multi-rank teardown validation has not been run for this draft.

Checklist

  • Format code according to the repository pre-commit configuration.
  • Add focused unit tests.
  • Follow the SGLang code style guidance.
  • Documentation update — not applicable to this internal lifecycle seam.
  • Accuracy and speed benchmarks — not applicable to reset-only cleanup.

CI States

Latest PR Test (Base): ❌ Run #32463957071
Latest PR Test (Extra): ❌ Run #32463956581
Latest PR Test (AMD ROCm 10): ❌ Run #32463956855

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants