[Fix] Repair CI fixtures and ROCm speculative tree device checks (#40325)

This commit is contained in:
Xiaoyu Zhang
2026-09-19 18:16:05 +08:00
committed by GitHub
parent 567d5925fe
commit 0b0d2c257a
4 changed files with 16 additions and 11 deletions
@@ -78,7 +78,7 @@ class _TokenToKVPool:
extra_key_buffer if extra_key_buffer is not None else swa_key_buffer
)
self.full_to_swa_index_mapping = full_to_swa_index_mapping
self.swa_page_size = page_size
self.swa_kv_pool = _Pool(page_size)
def get_swa_key_buffer_radix(self, layer_id: int) -> torch.Tensor:
_ = layer_id
@@ -86,7 +86,7 @@ class _TokenToKVPool:
def get_extra_key_page_size(self, layer_id: int) -> int:
_ = layer_id
return self.swa_page_size
return self.swa_kv_pool.page_size
def get_extra_key_buffer(self, layer_id: int) -> torch.Tensor:
_ = layer_id