tpu-inference 0.0.1rc1__py3-none-any.whl → 0.11.1.dev202511180814__py3-none-any.whl

tpu_inference/kernels/ragged_paged_attention/v3/kernel.py

@@ -440,54 +440,42 @@ def _ragged_paged_attention_kernel(
         debug_print("[RPA debug] bkv_sz_frm_new={}", bkv_sz_frm_new)
         debug_print("[RPA debug] page_indices_offset={}", page_indices_offset)
 
-        if not wait:
-            # Fetch effective kv from kv cache.
-            def loop_body(i, offset):
-                sz = jnp.minimum(page_size, kv_left_frm_cache - i * page_size)
-                _async_copy(
-                    cache_hbm_ref.at[pl.ds(
-                        page_indices_ref[page_indices_offset + i] * page_size,
-                        sz)],
-                    vmem_ref.at[pl.ds(i * page_size, sz)],
-                    sem,
-                    wait=False,
-                )
-                debug_print("[RPA debug] loop_body i={}, sz={}", i, sz)
-                return offset + sz
-
-            offset = lax.fori_loop(
-                0,
-                bkv_p_frm_cache,
-                loop_body,
-                0,  # offset
-                unroll=False,
+        # Fetch effective kv from kv cache.
+        def loop_body(i, offset):
+            sz = jnp.minimum(page_size, kv_left_frm_cache - i * page_size)
+            _async_copy(
+                cache_hbm_ref.at[pl.ds(
+                    page_indices_ref[page_indices_offset + i] * page_size,
+                    sz)],
+                vmem_ref.at[pl.ds(i * page_size, sz)],
+                sem,
+                wait,
             )
+            debug_print("[RPA debug] loop_body i={}, sz={}", i, sz)
+            return offset + sz
+
+        offset = lax.fori_loop(
+            0,
+            bkv_p_frm_cache,
+            loop_body,
+            0,  # offset
+            unroll=False,
+        )
 
-            # Fetch kv directly from new kv.
-            @pl.when(bkv_sz_frm_new > 0)
-            def _fetch_bkv_from_new_kv():
-                new_kv_len_start = q_end - kv_left_frm_new
-                debug_print("[RPA debug] new_kv_len_start={}",
-                            new_kv_len_start)
-                debug_print("[RPA debug] offset_in_bkv={}", offset)
-                _async_copy(
-                    kv_hbm_ref.at[pl.ds(new_kv_len_start, bkv_sz_frm_new)],
-                    vmem_ref.at[pl.ds(offset, bkv_sz_frm_new)],
-                    sem,
-                    wait,
-                )
-
-            return kv_len_start + offset, bkv_sz_frm_new
-        else:
-            offset = jnp.minimum(kv_left_frm_cache, page_size * bkv_p)
-            dst = vmem_ref.at[pl.ds(0, offset + bkv_sz_frm_new)]
+        # Fetch kv directly from new kv.
+        @pl.when(bkv_sz_frm_new > 0)
+        def _fetch_bkv_from_new_kv():
+            new_kv_len_start = q_end - kv_left_frm_new
+            debug_print("[RPA debug] new_kv_len_start={}", new_kv_len_start)
+            debug_print("[RPA debug] offset_in_bkv={}", offset)
             _async_copy(
-                src=dst,
-                dst=dst,
-                sem=sem,
-                wait=True,
+                kv_hbm_ref.at[pl.ds(new_kv_len_start, bkv_sz_frm_new)],
+                vmem_ref.at[pl.ds(offset, bkv_sz_frm_new)],
+                sem,
+                wait,
             )
-            return kv_len_start + offset, bkv_sz_frm_new
+
+        return kv_len_start + offset, bkv_sz_frm_new
 
     def _update_kv_cache(seq_idx,
                          bkv_sem_idx,
@@ -523,41 +511,30 @@ def _ragged_paged_attention_kernel(
         debug_print("[RPA debug] p_ignore={}", p_ignore)
         debug_print("[RPA debug] page_indices_offset={}", page_indices_offset)
 
-        if not wait:
-
-            def loop_body(i, states):
-                update_sz, ignore = states
-                sz = jnp.minimum(page_size - ignore, update_sz)
-
-                _async_copy(
-                    vmem_ref.at[pl.ds((p_ignore + i) * page_size + ignore,
-                                      sz)],
-                    cache_hbm_ref.at[pl.ds(
-                        page_indices_ref[page_indices_offset + i] * page_size +
-                        ignore,
-                        sz,
-                    )],
-                    sem,
-                    wait=False,
-                )
-                debug_print("[RPA debug] loop_body i={}, sz={}", i, sz)
-                return update_sz - sz, 0
-
-            lax.fori_loop(
-                0,
-                kv_p_end - kv_p_start,
-                loop_body,
-                (update_sz, ignore),  # total transfer size
-                unroll=False,
-            )
-        else:
-            dst = cache_hbm_ref.at[pl.ds(0, update_sz)]
+        def loop_body(i, states):
+            update_sz, ignore = states
+            sz = jnp.minimum(page_size - ignore, update_sz)
+
             _async_copy(
-                src=dst,
-                dst=dst,
-                sem=sem,
-                wait=True,
+                vmem_ref.at[pl.ds((p_ignore + i) * page_size + ignore, sz)],
+                cache_hbm_ref.at[pl.ds(
+                    page_indices_ref[page_indices_offset + i] * page_size +
+                    ignore,
+                    sz,
+                )],
+                sem,
+                wait,
             )
+            debug_print("[RPA debug] loop_body i={}, sz={}", i, sz)
+            return update_sz - sz, 0
+
+        lax.fori_loop(
+            0,
+            kv_p_end - kv_p_start,
+            loop_body,
+            (update_sz, ignore),  # total transfer size
+            unroll=False,
+        )
 
     def _fetch_bq(seq_idx, bq_idx, bq_sem_idx, *, wait=False):
         sem = sems.at[1, bq_sem_idx]