[fix] baseline4 ixf to xllm
This commit is contained in:
@@ -2290,10 +2290,32 @@ class PagedAttention:
|
||||
kv_caches: List[torch.Tensor],
|
||||
src_to_dists: torch.Tensor,
|
||||
) -> None:
|
||||
key_caches = [kv_cache[0] for kv_cache in kv_caches]
|
||||
value_caches = [kv_cache[1] for kv_cache in kv_caches]
|
||||
# BI100 CoreX 3.2.3: ixformer exposes vllm_copy_blocks, not
|
||||
# copy_blocks. Call the vendor symbol directly instead of going
|
||||
# through ops.copy_blocks (which hits the missing name).
|
||||
import ixformer.functions as _ixf
|
||||
_ixf.vllm_copy_blocks(key_caches, value_caches, src_to_dists)
|
||||
# BI100 CoreX 3.2.3: ixformer has no copy_blocks binding.
|
||||
# Use xllm_cache.block_copy kernel — single launch for all layers.
|
||||
if src_to_dists.numel() == 0:
|
||||
return
|
||||
from ex_engine.python.xllm_ops import _get as _xllm_get
|
||||
_xllm_block_copy = _xllm_get("xllm_cache").block_copy
|
||||
|
||||
device = src_to_dists.device
|
||||
n = src_to_dists.size(0)
|
||||
|
||||
# 1:1 mapping → each pair is its own group
|
||||
src_indices = src_to_dists[:, 0].to(torch.int32).contiguous()
|
||||
dst_indices = src_to_dists[:, 1].to(torch.int32).contiguous()
|
||||
cum_sum = torch.arange(1, n + 1, dtype=torch.int32, device=device)
|
||||
|
||||
# Build per-layer cache pointer tensors
|
||||
key_ptrs = torch.tensor(
|
||||
[kv[0].data_ptr() for kv in kv_caches],
|
||||
dtype=torch.int64, device=device)
|
||||
val_ptrs = torch.tensor(
|
||||
[kv[1].data_ptr() for kv in kv_caches],
|
||||
dtype=torch.int64, device=device)
|
||||
|
||||
numel_per_block = kv_caches[0][0][0].numel()
|
||||
cache_dtype = kv_caches[0][0].dtype
|
||||
|
||||
_xllm_block_copy(key_ptrs, val_ptrs,
|
||||
src_indices, dst_indices, cum_sum,
|
||||
numel_per_block, cache_dtype)
|
||||
Reference in New Issue
Block a user