Repository navigation
[Cherry-pick to release/v0.5.12] [Fix] DSV4 cached_loc invalidated when SWA mapping is rebuilt (#25889) - #26079
Conversation
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request fixes a critical bug in the DeepSeekV4 memory pool where the SWA location cache was not invalidated when a new index mapping was registered. This stale cache could lead to incorrect memory writes or out-of-bounds access crashes. The fix introduces an invalidate_loc_cache method and ensures the cache is cleared during mapping registration. Several regression tests were added to cover the bug scenario, end-to-end behavior, and potential crashes. Review feedback suggests improving maintainability by calling the new invalidation method within register_mapping and adding a docstring to the public method for better documentation.
|
|
||
| def register_mapping(self, full_to_swa_index_mapping: torch.Tensor): | ||
| self.full_to_swa_index_mapping = full_to_swa_index_mapping | ||
| self.cached_loc = None # mapping replaced; discard any cached translation |
There was a problem hiding this comment.
To avoid logic duplication and improve maintainability, consider calling invalidate_loc_cache() here instead of directly setting self.cached_loc = None.
| self.cached_loc = None # mapping replaced; discard any cached translation | |
| self.invalidate_loc_cache() # mapping replaced; discard any cached translation |
| def invalidate_loc_cache(self) -> None: | ||
| self.cached_loc = None |
There was a problem hiding this comment.
Add a docstring to explain the purpose of this new public method. Additionally, consider adding this method to the BaseSWAKVPool abstract base class to provide a consistent interface for invalidating SWA-related caches across different pool implementations.
| def invalidate_loc_cache(self) -> None: | |
| self.cached_loc = None | |
| def invalidate_loc_cache(self) -> None: | |
| """Invalidate the cached SWA location translation.""" | |
| self.cached_loc = None |
Cherry-pick of commit
888a8794ef3d689c4cb9417639d3fcfcee23f0betorelease/v0.5.12.Source PR: #25889 (#25889)
Source commit:
888a8794ef3d689c4cb9417639d3fcfcee23f0beOriginal title: [Fix] DSV4 cached_loc invalidated when SWA mapping is rebuilt
This PR was automatically created by the cherry-pick workflow.
CI States
Latest PR Test (Base): ❌ Missing
run-cilabel -- add it to run CI tests.Latest PR Test (Extra): ❌ Blocked --
run-ciis required first.