nixl pd support qwen3.5#1340
Conversation
![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 adds support for transferring linear attention states (convolution and SSM states) alongside KV cache pages in NIXL, specifically for Qwen3Next models. It introduces a page_kind parameter to differentiate between KV cache and linear attention states, updates the NIXL transfer task structures, and implements the necessary state copying logic in the Qwen3Next memory manager. Feedback from the review highlights a critical inconsistency where NIXLChunckedTransTaskRet.get_key() was not updated to match the new key format of NIXLChunckedTransTask, which would cause task matching failures. Additionally, the reviewer suggests adding type annotations to the newly added methods in qwen3next_mem_manager.py to align with the base class signature.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def get_key(self) -> str: | ||
| return f"{self.request_id}_{self.start_kv_index}_{self.end_kv_index}" | ||
| return f"{self.request_id}_{self.req_idx}_{self.page_kind}_{self.start_kv_index}_{self.end_kv_index}" |
There was a problem hiding this comment.
The get_key() method of NIXLChunckedTransTask has been updated to include req_idx and page_kind to uniquely identify tasks and prevent collisions. However, NIXLChunckedTransTaskRet.get_key() (defined around line 389) was not updated and still uses the old format f"{self.request_id}_{self.start_kv_index}_{self.end_kv_index}". This inconsistency will cause task matching to fail when using the returned task objects. Please update NIXLChunckedTransTaskRet to include req_idx and page_kind fields and update its get_key() method accordingly.
| def write_mem_to_page_kv_move_buffer( | ||
| self, | ||
| mem_indexes, | ||
| page_index: int, | ||
| dp_index: int, | ||
| mem_managers, | ||
| dp_world_size: int, | ||
| page_kind: str = "kv", | ||
| req_idx: int = None, | ||
| ): |
There was a problem hiding this comment.
Add type annotations to mem_indexes and mem_managers to match the base class signature and improve type safety.
| def write_mem_to_page_kv_move_buffer( | |
| self, | |
| mem_indexes, | |
| page_index: int, | |
| dp_index: int, | |
| mem_managers, | |
| dp_world_size: int, | |
| page_kind: str = "kv", | |
| req_idx: int = None, | |
| ): | |
| def write_mem_to_page_kv_move_buffer( | |
| self, | |
| mem_indexes: List[int], | |
| page_index: int, | |
| dp_index: int, | |
| mem_managers: List["MemoryManager"], | |
| dp_world_size: int, | |
| page_kind: str = "kv", | |
| req_idx: int = None, | |
| ): |
| def read_page_kv_move_buffer_to_mem( | ||
| self, | ||
| mem_indexes, | ||
| page_index: int, | ||
| dp_index: int, | ||
| mem_managers, | ||
| dp_world_size: int, | ||
| page_kind: str = "kv", | ||
| req_idx: int = None, | ||
| ): |
There was a problem hiding this comment.
Add type annotations to mem_indexes and mem_managers to match the base class signature and improve type safety.
| def read_page_kv_move_buffer_to_mem( | |
| self, | |
| mem_indexes, | |
| page_index: int, | |
| dp_index: int, | |
| mem_managers, | |
| dp_world_size: int, | |
| page_kind: str = "kv", | |
| req_idx: int = None, | |
| ): | |
| def read_page_kv_move_buffer_to_mem( | |
| self, | |
| mem_indexes: List[int], | |
| page_index: int, | |
| dp_index: int, | |
| mem_managers: List["MemoryManager"], | |
| dp_world_size: int, | |
| page_kind: str = "kv", | |
| req_idx: int = None, | |
| ): |
1 participant
No description provided.