tpu-inference 0.11.1.dev202511150811__py3-none-any.whl → 0.11.1.dev202511270815__py3-none-any.whl

tpu_inference/kernels/ragged_paged_attention/v3/kernel.py

@@ -440,42 +440,54 @@ 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)
 
-        # 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,
+        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,
             )
-            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)
+            # 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)]
             _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,
+                src=dst,
+                dst=dst,
+                sem=sem,
+                wait=True,
             )
-
-        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,
@@ -511,30 +523,41 @@ def _ragged_paged_attention_kernel(
         debug_print("[RPA debug] p_ignore={}", p_ignore)
         debug_print("[RPA debug] page_indices_offset={}", page_indices_offset)
 
-        def loop_body(i, states):
-            update_sz, ignore = states
-            sz = jnp.minimum(page_size - ignore, update_sz)
-
+        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)]
             _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,
+                src=dst,
+                dst=dst,
+                sem=sem,
+                wait=True,
             )
-            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]

tpu_inference/kernels/ragged_paged_attention/v3/kernel_hd64.py

@@ -267,6 +267,7 @@ def _ragged_paged_attention_kernel(
     *,
     sm_scale: float,
     sliding_window: int | None = None,
+    strict_sliding_window: bool = True,
     soft_cap: float | None = None,
     mask_value: float = DEFAULT_MASK_VALUE,
     q_scale: float | None = None,
@@ -317,19 +318,20 @@ def _ragged_paged_attention_kernel(
     q_len = q_end - q_start
     kv_len = kv_lens_ref[seq_idx]
 
-    bkv_idx_start = 0 if sliding_window is None else jnp.maximum(
-        kv_len - sliding_window, 0) // bkv_sz
-
     if sliding_window is None:
-        next_bkv_idx_start = 0
+        bkv_idx_start = next_seq_bkv_idx_start = 0
     else:
+        bkv_idx_start = jnp.maximum(kv_len - q_len - sliding_window,
+                                    0) // bkv_sz
 
         def get_next_bkv_idx_start():
             next_kv_len = kv_lens_ref[seq_idx + 1]
-            return jnp.maximum(next_kv_len - sliding_window, 0) // bkv_sz
+            next_q_len = cu_q_lens_ref[seq_idx + 2] - q_end
+            return jnp.maximum(next_kv_len - next_q_len - sliding_window,
+                               0) // bkv_sz
 
-        next_bkv_idx_start = lax.cond(seq_idx + 1 < num_seqs,
-                                      get_next_bkv_idx_start, lambda: 0)
+        next_seq_bkv_idx_start = lax.cond(seq_idx + 1 < num_seqs,
+                                          get_next_bkv_idx_start, lambda: 0)
 
     def debug_print(msg, *args):
         if debug_mode:
@@ -386,16 +388,19 @@ def _ragged_paged_attention_kernel(
             s *= k_scale
         if q_scale is not None:
             s *= q_scale
+        if soft_cap is not None:
+            s = soft_cap * jnp.tanh(s / soft_cap)
 
         q_span = (kv_len - q_len + bq_idx * bq_sz +
                   lax.broadcasted_iota(jnp.int32, s.shape, 0) //
                   num_q_heads_per_kv_head)
         k_span = bkv_idx * bkv_sz + lax.broadcasted_iota(jnp.int32, s.shape, 1)
-        mask = q_span < k_span
+        mask = k_span <= q_span
 
-        if soft_cap is not None:
-            s = soft_cap * jnp.tanh(s / soft_cap)
-        s += jnp.where(mask, mask_value, 0.0)
+        if sliding_window is not None and strict_sliding_window:
+            mask = jnp.logical_and(mask, q_span - sliding_window < k_span)
+
+        s = jnp.where(mask, s, mask_value)
         s_rowmax = jnp.max(s, axis=1, keepdims=True)
 
         if attention_sink_ref is not None:
@@ -475,42 +480,54 @@ 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)
 
-        # 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,
+        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,
             )
-            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)
+            # 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)]
             _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,
+                src=dst,
+                dst=dst,
+                sem=sem,
+                wait=True,
             )
-
-        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,
@@ -546,30 +563,41 @@ def _ragged_paged_attention_kernel(
         debug_print("[RPA debug] p_ignore={}", p_ignore)
         debug_print("[RPA debug] page_indices_offset={}", page_indices_offset)
 
-        def loop_body(i, states):
-            update_sz, ignore = states
-            sz = jnp.minimum(page_size - ignore, update_sz)
-
+        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)]
             _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,
+                src=dst,
+                dst=dst,
+                sem=sem,
+                wait=True,
             )
-            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]
@@ -737,13 +765,17 @@ def _ragged_paged_attention_kernel(
             next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx)
             next_bkv_sem_idx = lax.select(bkv_sem_idx == 0, 1, 0)
 
-            next_bkv_idx = lax.select(
-                is_last_bkv,
-                lax.select(
+            if sliding_window is None:
+                next_bkv_start_idx = 0
+            else:
+                next_bkv_start_idx = lax.select(
                     is_last_bq,
-                    next_bkv_idx_start,
+                    next_seq_bkv_idx_start,
                     bkv_idx_start,
-                ), next_bkv_idx)
+                )
+            next_bkv_idx = lax.select(is_last_bkv, next_bkv_start_idx,
+                                      next_bkv_idx)
+
             return next_seq_idx, next_bq_idx, next_bkv_idx, next_bkv_sem_idx
 
         def compute_with_bq(bq_idx, _):
@@ -1226,6 +1258,7 @@ def static_validate_inputs(
     static_argnames=(
         "sm_scale",
         "sliding_window",
+        "strict_sliding_window",
         "soft_cap",
         "mask_value",
         "q_scale",
@@ -1255,6 +1288,7 @@ def ragged_paged_attention_hd64(
     *,
     sm_scale: float = 1.0,
     sliding_window: int | None = None,
+    strict_sliding_window: bool = True,
     soft_cap: float | None = None,
     mask_value: float | None = DEFAULT_MASK_VALUE,
     q_scale: float | None = None,
@@ -1269,42 +1303,41 @@ def ragged_paged_attention_hd64(
     # Debug params.
     debug_mode: bool = False,
 ):
-    """A special Ragged paged attention version for head_dim=64 that supports mixed
-
-    prefill and decode.
-
-    Args:
-        queries: concatenated all sequences' queries.
-        keys: concatenated all sequences' keys (quantized).
-        values: concatenated all sequences' values (quantized).
-        kv_cache: paged KV cache with TPU-friendly shape.
-        kv_lens: padded kv lengths. Only the first num_seqs values are valid.
-        page_indices: flattened page indices look-up table by (seq_id, page_id).
-        cu_q_lens: the cumulative sum of the effective query lengths. Similar to
-        kv_lens, only the first num_seqs+1 values are valid.
-        distribution: (i, j, k) represents that sequences[0:i] are decode-only,
-        sequences[i:j] are chunked-prefill-only, and sequences[j:k] are mixed. The
-        k is also the total number of sequences.
-        attention_sink: optional attention sink for each q head.
-        actual_head_dim: the actual head size of the attention. Here we assume k and
-        v have the same actual head size.
-        sm_scale: the softmax scale which will be applied to the Q@K^T.
-        sliding_window: the sliding window size for the attention.
-        soft_cap: the logit soft cap for the attention.
-        mask_value: mask value for causal mask.
-        k_scale: the scale for the key cache.
-        v_scale: the scale for the value cache.
-        num_kv_pages_per_block: number of kv pages to be processed in one flash
-        attention block in the pallas kernel.
-        num_queries_per_block: number of kv pages to be processed in one flash
-        attention block in the pallas kernel.
-        vmem_limit_bytes: the vmem limit for the pallas kernel.
-        debug_mode: if true, RPA does not issue any DMAs or run flash attention but
-        print debug info. Need to compile with `--xla_tpu_enable_log_recorder`.
-
-    Returns:
-        The output of the attention.
-    """
+    """A variant of ragged paged attention for head_dim=64.
+
+  Args:
+    queries: concatenated all sequences' queries.
+    keys: concatenated all sequences' keys (quantized).
+    values: concatenated all sequences' values (quantized).
+    kv_cache: paged KV cache with TPU-friendly shape.
+    kv_lens: padded kv lengths. Only the first num_seqs values are valid.
+    page_indices: flattened page indices look-up table by (seq_id, page_id).
+    cu_q_lens: the cumulative sum of the effective query lengths. Similar to
+      kv_lens, only the first num_seqs+1 values are valid.
+    distribution: (i, j, k) represents that sequences[0:i] are decode-only,
+      sequences[i:j] are chunked-prefill-only, and sequences[j:k] are mixed. The
+      k is also the total number of sequences.
+    attention_sink: optional attention sink for each q head.
+    sm_scale: the softmax scale which will be applied to the Q@K^T.
+    sliding_window: the sliding window size for the attention.
+    strict_sliding_window: compute tokens that are strictly within the window.
+    soft_cap: the logit soft cap for the attention.
+    mask_value: mask value for causal mask.
+    q_scale: the scale for the query.
+    k_scale: the scale for the key cache.
+    v_scale: the scale for the value cache.
+    chunk_prefill_size: the chunk prefill size for the attention.
+    num_kv_pages_per_block: number of kv pages to be processed in one flash
+      attention block in the pallas kernel.
+    num_queries_per_block: number of kv pages to be processed in one flash
+      attention block in the pallas kernel.
+    vmem_limit_bytes: the vmem limit for the pallas kernel.
+    debug_mode: if true, RPA does not issue any DMAs or run flash attention but
+      print debug info. Need to compile with `--xla_tpu_enable_log_recorder`.
+
+  Returns:
+    The output of the attention.
+  """
     q, k, v = queries, keys, values
     static_validate_inputs(
         q,
@@ -1374,7 +1407,7 @@ def ragged_paged_attention_hd64(
         pl.BlockSpec(memory_space=pltpu.HBM),
         pl.BlockSpec(memory_space=pltpu.HBM),
         None if attention_sink is None else pl.BlockSpec(
-            memory_space=pltpu.VMEM)
+            memory_space=pltpu.VMEM),
     ]
 
     out_specs = [
@@ -1438,6 +1471,7 @@ def ragged_paged_attention_hd64(
                 _ragged_paged_attention_kernel,
                 sm_scale=sm_scale,
                 sliding_window=sliding_window,
+                strict_sliding_window=strict_sliding_window,
                 soft_cap=soft_cap,
                 mask_value=mask_value,
                 q_scale=q_scale,

tpu_inference/layers/common/attention_interface.py

@@ -308,7 +308,13 @@ def sharded_ragged_paged_attention(
     args = (q, k, v, kv_cache, kv_lens, page_indices, cu_q_lens, distribution)
 
     use_hd64 = q.shape[-1] == 64
-    func = ragged_paged_attention_hd64 if use_hd64 else ragged_paged_attention
+
+    func = ragged_paged_attention
+    if use_hd64:
+        func = functools.partial(ragged_paged_attention_hd64,
+                                 strict_sliding_window=False)
+    else:
+        func = ragged_paged_attention
 
     if attention_sink is not None:
         if not use_hd64:

tpu_inference/layers/common/sharding.py

@@ -166,9 +166,10 @@ class ShardingConfigManager:
                     f"LoRA is not supported with data parallelism "
                     f"(DP size: {total_dp_size}). Please disable LoRA or "
                     f"set data parallelism to 1.")
+        if sharding_strategy.attention_data_parallelism > 1:
             if not os.environ.get("NEW_MODEL_DESIGN", False):
                 raise ValueError(
-                    "Must run DP with NEW_MODEL_DESIGN enabled. Please set the "
+                    "Must run Attention DP with NEW_MODEL_DESIGN enabled. Please set the "
                     "NEW_MODEL_DESIGN=True.")
 
     @property