Skip to content

[mem_cache] Add free_full to release the full side of a tombstoned SWA node - #36637

Merged
hnyls2002 merged 6 commits into
mainfrom
lsyin/swa-free-full
Aug 28, 2026
Merged

hnyls2002 merged 6 commits into
mainfrom
lsyin/swa-free-full

Conversation

@hnyls2002

@hnyls2002 hnyls2002 commented Aug 27, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Add free_full to the allocator family: a plain free() on the base, full-pool-only on the hybrid SWA allocators, a no-op on the all-SWA one
  • Route the seven tombstoned-node releases through it, in both SWARadixCache and the unified tree
  • Collapse the legacy tree's three evict-path releases into one _free_node_value helper, so the allocator call and the two eviction counters stay on the same branch

Background

  • On a hybrid SWA allocator free() releases both halves: it gathers full_to_swa_index_mapping[idx], drops the entries that read as the padding slot, and returns the rest to the SWA pool
  • A tombstoned node's SWA peers went back earlier (an SWA evict, or dec_swa_lock_only), and that cleared their mapping entries -- so free() on such a node already resolves to "full side only" today, purely because the sentinel filter absorbs it
  • The caller knows which case it is. It just had no way to say so, and the filter is what stood between a mistake and the padding slot entering the SWA free list

Scope

  • No behavior change: free_full(x) releases a subset of what free(x) releases, and on a tombstoned node the two are already equal
  • What it buys is that a wrong caller becomes an assertable bug rather than silent SWA-pool corruption
  • PureSWATokenToKVPoolAllocator gets an explicit no-op override -- its full_attn_allocator is its swa_attn_allocator, so inheriting the base default would free one pool twice
  • All classes that define free_swa now define free_full as well; a hybrid subclass silently inheriting the base default is exactly the double free above

Also here

Free-group timing. _recover_tombstone_keeping_locked_full and the RecoverSWAWithLockedFull handler reached into alloc.full_attn_allocator.free(...) directly, which released immediately even inside a free group while every other path deferred to the barrier. They now go through free_full and defer like the rest.

Unified needs a distinct action type, not a flag read. FreeDeviceKV and friends are batched until a barrier, and the same batch's SWARebuild reassigns the node's full value before the batch drains -- by then the node's tombstone state no longer describes the value being freed. FreeDeviceKVFullOnly freezes the decision at emit time instead.

Test plan

  • test_swa_unittest.py: TestFreeFullPartition -- the partition adds up to free, the SWA peers stay allocated, the mapping is left intact, the release defers inside a free group, and a single-pool allocator falls back to free
  • test_unified_radix_cache_unittest.py: the RecoverSWAWithLockedFull case now asserts free_full instead of the inner allocator, so the free-group defer above is covered
  • test/registered/unit/mem_cache passes

CI States

Latest PR Test (Base): 🚫 Run #33136294822
Latest PR Test (Extra): 🚫 Run #33136294751
Latest PR Test (AMD ROCm 7.2): 🚫 Run #33136294714

@hnyls2002
hnyls2002 merged commit daf6317 into main Aug 28, 2026
37 of 58 checks passed
@hnyls2002
hnyls2002 deleted the lsyin/swa-free-full branch August 28, 2026 02:34
saturn-acc pushed a commit to saturn-acc/sglang that referenced this pull request Aug 31, 2026
nzr-niu pushed a commit to nzr-niu/sglang that referenced this pull request Sep 1, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant