From 5d2faa3c005cab59b258b1819b59f2dd80302510 Mon Sep 17 00:00:00 2001
From: root
 <root@cc-58e31738-7d16-4b93-aa09-745898b4ddee-0.cc-58e31738-7d16-4b93-aa09-745898b4ddee.ws-19872517457.svc.cluster.local>
Date: Wed, 19 Aug 2026 01:07:16 +0000
Subject: [PATCH] [fix] baseline4 ixf to xllm

---
 qwen3_6_scripts/paged_attn.py | 36 ++++++++++++++++++++++++++++-------
 1 file changed, 29 insertions(+), 7 deletions(-)

diff --git a/qwen3_6_scripts/paged_attn.py b/qwen3_6_scripts/paged_attn.py
index 23118933..984ebf4a 100644
--- a/qwen3_6_scripts/paged_attn.py
+++ b/qwen3_6_scripts/paged_attn.py
@@ -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)
\ No newline at end of file
+        # 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)
\ No newline at end of file