Skip to content

Fix DSA top-k pinned buffer growth - #1984

Open
jayhenry wants to merge 1 commit into
InternLM:mainfrom
jayhenry:fix/dsa-topk-pinned-buffer
Open

Fix DSA top-k pinned buffer growth#1984
jayhenry wants to merge 1 commit into
InternLM:mainfrom
jayhenry:fix/dsa-topk-pinned-buffer

Conversation

@jayhenry

@jayhenry jayhenry commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

Summary

DSA top-k offload used a monotonically increasing context ID for pinned storage, creating new pinned buffers every training step instead of reusing them.

This change:

  • assigns stable slots to concurrently active microbatches;
  • uses one (offload_slot, source_layer_idx) key for both the live SwapTensor and its reusable pinned storage;
  • rejects slot reuse while the previous runtime entry is still active;
  • propagates the slot through main and MTP sequence contexts;
  • removes the obsolete context ID counter;
  • adds multi-step micro1 and EP2/micro2 regression coverage.

Validation

  • pytest -q tests/model/test_glm52_mtp_checkpoint_repro.py: 2 passed
  • pytest -q tests/module/attention/test_dsa_mla.py::TestDSAAttention: 3 passed
  • pre-commit hooks: passed, including ruff, mypy, and pydantic checks

Result

Each slot/source pair has one live offload lifecycle and one reusable pinned buffer, so pinned memory remains bounded across optimizer steps.

@jayhenry
jayhenry force-pushed the fix/dsa-topk-pinned-buffer branch from e068496 to fab03d3 Compare July 29, 2026 03:05
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.

1 participant