tpu-inference 0.12.0.dev20251222__py3-none-any.whl

tpu_inference/kernels/ragged_paged_attention/v3/kernel.py

@@ -0,0 +1,1594 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TPU-Friendly Ragged Paged Attention kernel.
+
+This kernel offers a highly optimized implementation of ragged paged attention,
+specifically designed for TPU and compatible with a wide range of model
+specifications. It supports mixed prefill and decoding, enhancing throughput
+during inference.
+"""
+import functools
+
+import jax
+import jax.numpy as jnp
+from jax import lax
+from jax.experimental import pallas as pl
+from jax.experimental.pallas import tpu as pltpu
+
+from tpu_inference.kernels.ragged_paged_attention.v3.tuned_block_sizes import \
+    get_tuned_block_sizes
+from tpu_inference.kernels.ragged_paged_attention.v3.util import (
+    align_to, cdiv, get_dtype_bitwidth, get_dtype_packing)
+
+DEFAULT_MASK_VALUE = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max)
+
+DEFAULT_VMEM_LIMIT_BYTES = 100 * 1024 * 1024
+
+
+def ref_ragged_paged_attention(
+    queries: jax.
+    Array,  # [max_num_tokens, actual_num_q_heads, actual_head_dim]
+    keys: jax.Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim]
+    values: jax.
+    Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim]
+    kv_cache: jax.
+    Array,  # [total_num_pages, page_size, num_kv_heads_x2 // kv_packing, kv_packing, head_dim]
+    kv_lens: jax.Array,  # i32[max_num_seqs]
+    page_indices: jax.Array,  # i32[max_num_seqs * pages_per_seq]
+    cu_q_lens: jax.Array,  # i32[max_num_seqs + 1]
+    distribution: jax.Array,  # i32[3]
+    *,
+    sm_scale: float = 1.0,
+    sliding_window: int | None = None,
+    soft_cap: float | None = None,
+    mask_value: float | None = DEFAULT_MASK_VALUE,
+    q_scale: float | None = None,
+    k_scale: float | None = None,
+    v_scale: float | None = None,
+):
+    if mask_value is None:
+        mask_value = DEFAULT_MASK_VALUE
+
+    dynamic_validate_inputs(
+        queries,
+        keys,
+        values,
+        kv_cache,
+        kv_lens,
+        page_indices,
+        cu_q_lens,
+        distribution,
+        sm_scale=sm_scale,
+        sliding_window=sliding_window,
+        soft_cap=soft_cap,
+        mask_value=mask_value,
+        q_scale=q_scale,
+        k_scale=k_scale,
+        v_scale=v_scale,
+    )
+    actual_head_dim = queries.shape[2]
+    actual_num_q_heads = queries.shape[1]
+    actual_num_kv_heads = keys.shape[1]
+    merged_kv = merge_kv(keys, values)
+    assert merged_kv.shape[-3:] == kv_cache.shape[-3:]
+
+    _, page_size, num_kv_heads_x2_per_kv_packing, kv_packing, head_dim = (
+        kv_cache.shape)
+    num_kv_heads_x2 = num_kv_heads_x2_per_kv_packing * kv_packing
+    assert num_kv_heads_x2 % 2 == 0
+    assert actual_num_q_heads % actual_num_kv_heads == 0
+    assert head_dim % 128 == 0
+    assert get_dtype_packing(kv_cache.dtype) == kv_packing
+    assert num_kv_heads_x2 == align_to(actual_num_kv_heads * 2, kv_packing)
+    actual_num_q_heads_per_kv_head = actual_num_q_heads // actual_num_kv_heads
+    max_num_seqs = kv_lens.shape[0]
+    num_page_indices = page_indices.shape[0]
+    assert num_page_indices % max_num_seqs == 0
+    pages_per_seq = num_page_indices // max_num_seqs
+    outputs = []
+
+    for i in range(distribution[-1]):
+        q_start = cu_q_lens[i]
+        q_end = cu_q_lens[i + 1]
+        q_len = q_end - q_start
+
+        kv_len = kv_lens[i]
+        indices_start = i * pages_per_seq
+        indices_end = indices_start + cdiv(kv_len, page_size)
+        indices = page_indices[indices_start:indices_end]
+        q = queries[q_start:q_end, :, :actual_head_dim]
+
+        # Update the kv cache.
+        assert kv_len - q_len >= 0
+        gathered_kv = kv_cache[indices]
+        gathered_shape = gathered_kv.shape
+        gathered_kv = gathered_kv.reshape(-1, *gathered_shape[-3:])
+        gathered_kv = gathered_kv.at[kv_len - q_len:kv_len].set(
+            merged_kv[q_start:q_end])
+        kv_cache = kv_cache.at[indices].set(
+            gathered_kv.reshape(gathered_shape))
+
+        kv = gathered_kv.reshape(
+            -1, num_kv_heads_x2,
+            head_dim)[:, :actual_num_kv_heads * 2, :].reshape(
+                -1, actual_num_kv_heads, head_dim * 2)
+        k = kv[:kv_len, :, :head_dim][:, :, :actual_head_dim]
+        v = kv[:kv_len, :, head_dim:][:, :, :actual_head_dim]
+        k = jnp.repeat(k, actual_num_q_heads_per_kv_head, axis=1)
+        v = jnp.repeat(v, actual_num_q_heads_per_kv_head, axis=1)
+
+        if q_scale is not None:
+            q = q / q_scale
+            if jnp.issubdtype(k.dtype, jnp.floating):
+                dtype_info = jnp.finfo(k.dtype)
+                minval = float(dtype_info.min)
+                maxval = float(dtype_info.max)
+                q = jnp.clip(q, min=minval, max=maxval)
+            q = q.astype(k.dtype)
+
+        attn = jnp.einsum("qhd,khd->hqk",
+                          q,
+                          k,
+                          preferred_element_type=jnp.float32)
+        attn *= sm_scale
+        if k_scale is not None:
+            attn *= k_scale
+        if q_scale is not None:
+            attn *= q_scale
+
+        q_span = (kv_len - q_len) + jax.lax.broadcasted_iota(
+            jnp.int32, attn.shape, 1)
+        kv_span = jax.lax.broadcasted_iota(jnp.int32, attn.shape, 2)
+        mask = q_span < kv_span
+        if sliding_window is not None:
+            mask = jnp.logical_or(mask, q_span - sliding_window >= kv_span)
+        if soft_cap is not None:
+            attn = soft_cap * jnp.tanh(attn / soft_cap)
+        attn += jnp.where(mask, mask_value, 0.0)
+        attn = jax.nn.softmax(attn, axis=-1).astype(v.dtype)
+
+        out = jnp.einsum("hqk,khd->qhd", attn, v).astype(queries.dtype)
+        if v_scale is not None:
+            out *= v_scale
+
+        outputs.append(out)
+
+    result = jnp.concatenate(outputs, axis=0)
+    return result, kv_cache
+
+
+def get_smem_estimate_bytes(max_num_seqs, pages_per_seq):
+    total_bits = (
+        # kv_lens_ref: i32[max_num_seqs]
+        align_to(max_num_seqs, 128) * 32 +
+        # page_indices_ref: i32[max_num_seqs * pages_per_seq]
+        align_to(max_num_seqs * pages_per_seq, 128) * 32 +
+        # cu_q_lens_ref: i32[max_num_seqs + 1]
+        align_to(max_num_seqs + 1, 128) * 32 +
+        # distribution_ref: i32[3]
+        128 * 32 +
+        # sem_ids_ref: i32[3]
+        128 * 32 +
+        # bo_ids_ref: i32[4]
+        128 * 32 +
+        # bkv_update_ids_ref: i32[6]
+        128 * 32)
+    return cdiv(total_bits, 8)
+
+
+def get_vmem_estimate_bytes(
+    actual_num_kv_heads,
+    actual_num_q_heads_per_kv_head,
+    actual_head_dim,
+    bq_sz,
+    bkv_sz,
+    q_dtype,
+    kv_dtype,
+):
+    q_packing = get_dtype_packing(q_dtype)
+    kv_packing = get_dtype_packing(kv_dtype)
+    num_q_heads_per_kv_head = align_to(actual_num_q_heads_per_kv_head,
+                                       q_packing)
+    num_kv_heads_x2 = align_to(actual_num_kv_heads * 2, kv_packing)
+    head_dim = align_to(actual_head_dim, 128)
+
+    total_bits = (
+        # bkv_x2_ref
+        (2 * bkv_sz * num_kv_heads_x2 * head_dim) * (32 // kv_packing) +
+        # bq_x2_ref + bo_x2_ref
+        2 * (2 * actual_num_kv_heads * bq_sz * num_q_heads_per_kv_head *
+             head_dim) * (32 // q_packing) +
+        # l_ref + m_ref
+        2 *
+        (actual_num_kv_heads * bq_sz * num_q_heads_per_kv_head * 128) * 32 +
+        # acc_ref
+        (actual_num_kv_heads * bq_sz * num_q_heads_per_kv_head * head_dim) *
+        32)
+    return cdiv(total_bits, 8)
+
+
+def get_kv_cache_shape(
+    total_num_pages,
+    page_size,
+    actual_num_kv_heads,
+    actual_head_dim,
+    kv_dtype,
+):
+    kv_packing = get_dtype_packing(kv_dtype)
+    return (
+        total_num_pages,
+        page_size,
+        align_to(actual_num_kv_heads * 2, kv_packing) // kv_packing,
+        kv_packing,
+        align_to(actual_head_dim, 128),
+    )
+
+
+def _ragged_paged_attention_kernel(
+    # Prefetch
+    kv_lens_ref,  # [max_num_seqs]
+    page_indices_ref,  # [max_num_seqs * pages_per_seq]
+    cu_q_lens_ref,  # [max_num_seqs + 1]
+    # TODO(jevinjiang): merge these into one so we can save SMEM.
+    distribution_ref,  # [3] (decode_end, prefill_end, mixed_end)
+    sem_ids_ref,  # [3] (bq_sem_idx, bkv_sem_idx, bo_sem_idx)
+    bo_ids_ref,  # [4] (bo_sem_0_seq_idx, bo_sem_1_seq_idx, bo_sem_0_bo_idx, bo_sem_1_bo_idx)
+    bkv_update_ids_ref,  # [6] (bkv_sem_0_seq_idx, bkv_sem_1_seq_idx, bkv_sem_0_offset, bkv_sem_1_offset, bkv_sem_0_sz, bkv_sem_1_sz)
+    # Input
+    q_hbm_ref,  # [actual_num_kv_heads, max_num_tokens, num_q_heads_per_kv_head // q_packing, q_packing, head_dim]
+    kv_hbm_ref,  # [max_num_tokens, num_kv_heads_x2 // kv_packing, kv_packing, head_dim]
+    kv_cache_hbm_ref,  # [total_num_pages, page_size, num_kv_heads_x2 // kv_packing, kv_packing, head_dim]
+    # Output
+    o_hbm_ref,  # [actual_num_kv_heads, max_num_tokens, num_q_heads_per_kv_head // q_packing, q_packing, head_dim]
+    updated_kv_cache_hbm_ref,  # [total_num_pages, page_size, num_kv_heads_x2 // kv_packing, kv_packing, head_dim]
+    # Scratch
+    bkv_x2_ref,  # [2, bkv_sz, num_kv_heads_x2 // kv_packing, kv_packing, head_dim]
+    bq_x2_ref,  # [2, actual_num_kv_heads, bq_sz, num_q_heads_per_kv_head // q_packing, q_packing, head_dim]
+    bo_x2_ref,  # [2, actual_num_kv_heads, bq_sz, num_q_heads_per_kv_head // q_packing, q_packing, head_dim]
+    sems,  # [4, 2]
+    l_ref,  # [actual_num_kv_heads, bq_sz * num_q_heads_per_kv_head, 128],
+    m_ref,  # [actual_num_kv_heads, bq_sz * num_q_heads_per_kv_head, 128],
+    acc_ref,  # [actual_num_kv_heads, bq_sz * num_q_heads_per_kv_head, head_dim],
+    *,
+    sm_scale: float,
+    sliding_window: int | None = None,
+    soft_cap: float | None = None,
+    mask_value: float = DEFAULT_MASK_VALUE,
+    q_scale: float | None = None,
+    k_scale: float | None = None,
+    v_scale: float | None = None,
+    chunk_prefill_size: int | None = None,
+    bkv_p,
+    bq_sz,
+    debug_mode: bool = False,
+):
+    assert q_hbm_ref.shape == o_hbm_ref.shape
+    assert q_hbm_ref.shape[-1] == kv_cache_hbm_ref.shape[-1]
+    (
+        actual_num_kv_heads,
+        max_num_tokens,
+        num_q_heads_per_kv_head_per_packing,
+        q_packing,
+        head_dim,
+    ) = q_hbm_ref.shape
+    (
+        total_num_pages,
+        page_size,
+        num_kv_heads_x2_per_kv_packing,
+        kv_packing,
+        _,
+    ) = kv_cache_hbm_ref.shape
+    max_num_seqs = kv_lens_ref.shape[0]
+    num_page_indices = page_indices_ref.shape[0]
+    assert num_page_indices % max_num_seqs == 0
+    pages_per_seq = num_page_indices // max_num_seqs
+    num_kv_heads_x2 = num_kv_heads_x2_per_kv_packing * kv_packing
+    num_q_heads_per_kv_head = num_q_heads_per_kv_head_per_packing * q_packing
+    q_dtype = q_hbm_ref.dtype
+    kv_dtype = kv_cache_hbm_ref.dtype
+    assert o_hbm_ref.dtype == q_dtype
+    assert get_dtype_packing(q_dtype) == q_packing
+    assert get_dtype_packing(kv_dtype) == kv_packing
+    assert head_dim % 128 == 0
+    bkv_sz = bkv_p * page_size
+    seq_idx = pl.program_id(0)
+    num_seqs = pl.num_programs(0)
+    decode_end = distribution_ref[0]
+    prefill_end = distribution_ref[1]
+    mixed_end = distribution_ref[2]
+
+    q_start = cu_q_lens_ref[seq_idx]
+    q_end = cu_q_lens_ref[seq_idx + 1]
+    q_len = q_end - q_start
+    kv_len = kv_lens_ref[seq_idx]
+
+    if sliding_window is None:
+        bkv_idx_start = next_seq_bkv_idx_start = 0
+    else:
+        bkv_idx_start = jnp.maximum(kv_len - q_len - sliding_window,
+                                    0) // bkv_sz
+
+        # If seq_idx + 1 == num_seqs, kv_lens_ref[seq_idx + 1] will trigger a
+        # out-of-bound error. To avoid this, we set upperbound of next_seq_idx
+        # to be num_seqs - 1.
+        next_seq_idx = jnp.minimum(seq_idx + 1, num_seqs - 1)
+        next_kv_len = kv_lens_ref[next_seq_idx]
+        next_q_len = cu_q_lens_ref[next_seq_idx + 1] - q_end
+        next_seq_bkv_idx_start = (
+            jnp.maximum(next_kv_len - next_q_len - sliding_window, 0) //
+            bkv_sz)
+
+    def debug_print(msg, *args):
+        if debug_mode:
+            pl.debug_print(msg, *args)
+
+    debug_print("[RPA debug] ======= In loop seq_idx={}", seq_idx)
+    debug_print("[RPA debug] num_seqs={}", num_seqs)
+    debug_print("[RPA debug] decode_end={}", decode_end)
+    debug_print("[RPA debug] prefill_end={}", prefill_end)
+    debug_print("[RPA debug] mixed_end={}", mixed_end)
+    debug_print("[RPA debug] bkv_p={}", bkv_p)
+    debug_print("[RPA debug] page_size={}", page_size)
+    debug_print("[RPA debug] pages_per_seq={}", pages_per_seq)
+    debug_print("[RPA debug] bkv_sz={}", bkv_sz)
+    debug_print("[RPA debug] bq_sz={}", bq_sz)
+    debug_print("[RPA debug] q_start={}", q_start)
+    debug_print("[RPA debug] q_end={}", q_end)
+    debug_print("[RPA debug] q_len={}", q_len)
+    debug_print("[RPA debug] kv_len={}", kv_len)
+
+    def flash_attention_step1_qk_softmax(
+        q,  # [actual_bq_sz * num_q_heads_per_kv_head, head_dim]
+        k,  # [bkv_sz, head_dim]
+        v,  # [bkv_sz, head_dim]
+        *,
+        bq_idx,
+        bkv_idx,
+        kv_head_idx,
+    ):
+        assert len(q.shape) == 2
+        assert q.shape[0] % num_q_heads_per_kv_head == 0
+        assert q.shape[1] == head_dim
+        assert k.shape == v.shape == (bkv_sz, head_dim)
+        assert k.dtype == v.dtype
+        head_l_ref = l_ref.at[kv_head_idx, :q.shape[0]]
+        head_m_ref = m_ref.at[kv_head_idx, :q.shape[0]]
+
+        def load_with_init(ref, init_val):
+            return jnp.where(bkv_idx == bkv_idx_start,
+                             jnp.full_like(ref, init_val), ref[...])
+
+        # Follow FlashAttention-2 forward pass.
+        if q_scale is not None:
+            q = q / q_scale
+            if jnp.issubdtype(k.dtype, jnp.floating):
+                dtype_info = jnp.finfo(k.dtype)
+                minval = float(dtype_info.min)
+                maxval = float(dtype_info.max)
+                q = jnp.clip(q, min=minval, max=maxval)
+            q = q.astype(k.dtype)
+
+        s = jnp.einsum("nd,md->nm", q, k, preferred_element_type=jnp.float32)
+        s *= sm_scale
+        if k_scale is not None:
+            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 = k_span <= q_span
+
+        if sliding_window is not None:
+            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)
+
+        m_prev = load_with_init(head_m_ref, -jnp.inf)
+        m_curr = jnp.maximum(m_prev, s_rowmax)
+        head_m_ref[...] = m_curr
+        p = jnp.exp(s - broadcast_minor(m_curr, s.shape))
+
+        p_rowsum = jnp.sum(p, axis=1, keepdims=True)
+        exp_m_diff = jnp.exp(m_prev - m_curr)
+        l_prev = load_with_init(head_l_ref, 0.0)
+        l_curr = exp_m_diff * l_prev + p_rowsum
+        head_l_ref[...] = l_curr
+
+        return p, exp_m_diff
+
+    def flash_attention_step2_pv(
+        q_shape_0,
+        v,  # [bkv_sz, head_dim]
+        p,  # from step1
+        exp_m_diff,  # from step1
+        *,
+        bkv_idx,
+        kv_head_idx,
+    ):
+        head_acc_ref = acc_ref.at[kv_head_idx, :q_shape_0]
+
+        def load_with_init(ref, init_val):
+            return jnp.where(bkv_idx == 0, jnp.full_like(ref, init_val),
+                             ref[...])
+
+        pv = jnp.einsum("nm,md->nd", p, v, preferred_element_type=jnp.float32)
+        if v_scale is not None:
+            pv *= v_scale
+        o_prev = load_with_init(head_acc_ref, 0.0)
+        o_curr = broadcast_minor(exp_m_diff, o_prev.shape) * o_prev + pv
+        head_acc_ref[...] = o_curr
+
+    def _async_copy(src, dst, sem, wait):
+        if debug_mode:
+            # Skip DMA if debug mode is enabled.
+            return
+        cp = pltpu.make_async_copy(src, dst, sem)
+        if wait:
+            cp.wait()
+        else:
+            cp.start()
+
+    def _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, *, wait=False):
+        sem = sems.at[0, bkv_sem_idx]
+        vmem_ref = bkv_x2_ref.at[bkv_sem_idx]
+
+        cache_hbm_shape = kv_cache_hbm_ref.shape
+        cache_hbm_ref = kv_cache_hbm_ref.reshape(
+            cache_hbm_shape[0] * cache_hbm_shape[1], *cache_hbm_shape[2:])
+        kv_len = kv_lens_ref[seq_idx]
+        kv_len_start = bkv_idx * bkv_sz
+        kv_p_start = bkv_idx * bkv_p
+        q_start = cu_q_lens_ref[seq_idx]
+        q_end = cu_q_lens_ref[seq_idx + 1]
+        q_len = q_end - q_start
+
+        kv_left = kv_len - kv_len_start
+        kv_left_frm_cache = jnp.maximum(kv_left - q_len, 0)
+        kv_left_frm_new = kv_left - kv_left_frm_cache
+        bkv_p_frm_cache = jnp.minimum(cdiv(kv_left_frm_cache, page_size),
+                                      bkv_p)
+        bkv_sz_frm_new = jnp.minimum(
+            jnp.maximum(bkv_sz - kv_left_frm_cache, 0), kv_left_frm_new)
+        page_indices_offset = seq_idx * pages_per_seq + kv_p_start
+
+        # Make sure the current bkv buffer is safe to overwrite.
+        wait_update_kv_cache(bkv_sem_idx)
+
+        debug_print(
+            "[RPA debug]"
+            f" -----------{'wait' if wait else 'start'}_fetch_bkv-----------")
+        debug_print("[RPA debug] seq_idx={}", seq_idx)
+        debug_print("[RPA debug] bkv_idx={}", bkv_idx)
+        debug_print("[RPA debug] bkv_sem_idx={}", bkv_sem_idx)
+        debug_print("[RPA debug] kv_len_start={}", kv_len_start)
+        debug_print("[RPA debug] kv_p_start={}", kv_p_start)
+        debug_print("[RPA debug] kv_left={}", kv_left)
+        debug_print("[RPA debug] kv_left_frm_cache={}", kv_left_frm_cache)
+        debug_print("[RPA debug] kv_left_frm_new={}", kv_left_frm_new)
+        debug_print("[RPA debug] bkv_p_frm_cache={}", bkv_p_frm_cache)
+        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,
+            )
+
+            size = lax.select(bkv_sz_frm_new > 0, bkv_sz_frm_new, 0)
+            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, size)],
+                vmem_ref.at[pl.ds(offset, size)],
+                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(
+                src=dst,
+                dst=dst,
+                sem=sem,
+                wait=True,
+            )
+            return kv_len_start + offset, bkv_sz_frm_new
+
+    def _update_kv_cache(seq_idx,
+                         bkv_sem_idx,
+                         offset,
+                         update_sz,
+                         *,
+                         wait=False):
+        sem = sems.at[3, bkv_sem_idx]
+        vmem_ref = bkv_x2_ref.at[bkv_sem_idx]
+        bkv_id = offset // bkv_sz
+        kv_p_start = offset // page_size
+        kv_p_end = cdiv(offset + update_sz, page_size)
+        ignore = offset % page_size
+        p_ignore = kv_p_start - bkv_id * bkv_p
+        page_indices_offset = seq_idx * pages_per_seq + kv_p_start
+
+        cache_hbm_shape = updated_kv_cache_hbm_ref.shape
+        cache_hbm_ref = updated_kv_cache_hbm_ref.reshape(
+            cache_hbm_shape[0] * cache_hbm_shape[1], *cache_hbm_shape[2:])
+
+        debug_print(
+            "[RPA debug]"
+            f" -----------{'wait' if wait else 'start'}_update_kv_cache-----------"
+        )
+        debug_print("[RPA debug] seq_idx={}", seq_idx)
+        debug_print("[RPA debug] bkv_sem_idx={}", bkv_sem_idx)
+        debug_print("[RPA debug] offset={}", offset)
+        debug_print("[RPA debug] update_sz={}", update_sz)
+        debug_print("[RPA debug] bkv_id={}", bkv_id)
+        debug_print("[RPA debug] kv_p_start={}", kv_p_start)
+        debug_print("[RPA debug] kv_p_end={}", kv_p_end)
+        debug_print("[RPA debug] ignore={}", ignore)
+        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)]
+            _async_copy(
+                src=dst,
+                dst=dst,
+                sem=sem,
+                wait=True,
+            )
+
+    def _fetch_bq(seq_idx, bq_idx, bq_sem_idx, *, wait=False):
+        sem = sems.at[1, bq_sem_idx]
+        vmem_ref = bq_x2_ref.at[bq_sem_idx]
+        q_len_start = cu_q_lens_ref[seq_idx] + bq_idx * bq_sz
+        q_end = cu_q_lens_ref[seq_idx + 1]
+        sz = jnp.minimum(bq_sz, q_end - q_len_start)
+
+        debug_print(
+            "[RPA debug]"
+            f" -----------{'wait' if wait else 'start'}_fetch_bq-----------")
+        debug_print("[RPA debug] seq_idx={}", seq_idx)
+        debug_print("[RPA debug] bq_idx={}", bq_idx)
+        debug_print("[RPA debug] bq_sem_idx={}", bq_sem_idx)
+        debug_print("[RPA debug] q_len_start={}", q_len_start)
+        debug_print("[RPA debug] q_end={}", q_end)
+        debug_print("[RPA debug] sz={}", sz)
+
+        _async_copy(
+            q_hbm_ref.at[:, pl.ds(q_len_start, sz)],
+            vmem_ref.at[:, pl.ds(0, sz)],
+            sem,
+            wait,
+        )
+
+    def _send_bo(seq_idx, bo_idx, bo_sem_idx, *, wait=False):
+        sem = sems.at[2, bo_sem_idx]
+        vmem_ref = bo_x2_ref.at[bo_sem_idx]
+        q_len_start = cu_q_lens_ref[seq_idx] + bo_idx * bq_sz
+        q_end = cu_q_lens_ref[seq_idx + 1]
+        sz = jnp.minimum(bq_sz, q_end - q_len_start)
+
+        debug_print(
+            "[RPA debug]"
+            f" -----------{'wait' if wait else 'start'}_send_bo-----------")
+        debug_print("[RPA debug] seq_idx={}", seq_idx)
+        debug_print("[RPA debug] bo_idx={}", bo_idx)
+        debug_print("[RPA debug] bo_sem_idx={}", bo_sem_idx)
+        debug_print("[RPA debug] q_len_start={}", q_len_start)
+        debug_print("[RPA debug] q_end={}", q_end)
+        debug_print("[RPA debug] sz={}", sz)
+
+        _async_copy(
+            vmem_ref.at[:, pl.ds(0, sz)],
+            o_hbm_ref.at[:, pl.ds(q_len_start, sz)],
+            sem,
+            wait,
+        )
+
+    def start_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx):
+        return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx)
+
+    def wait_fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx):
+        return _fetch_bkv(seq_idx, bkv_idx, bkv_sem_idx, wait=True)
+
+    def start_fetch_bq(seq_idx, bq_idx, bq_sem_idx):
+        return _fetch_bq(seq_idx, bq_idx, bq_sem_idx)
+
+    def wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx):
+        return _fetch_bq(seq_idx, bq_idx, bq_sem_idx, wait=True)
+
+    def start_send_bo(seq_idx, bo_idx, bo_sem_idx):
+        bo_ids_ref[bo_sem_idx] = seq_idx
+        bo_ids_ref[bo_sem_idx + 2] = bo_idx
+        _send_bo(seq_idx, bo_idx, bo_sem_idx)
+
+    def wait_send_bo(bo_sem_idx):
+        old_seq_idx = bo_ids_ref[bo_sem_idx]
+        old_bo_idx = bo_ids_ref[bo_sem_idx + 2]
+
+        @pl.when(jnp.logical_and(0 <= old_seq_idx, old_seq_idx <= seq_idx))
+        def _():
+            _send_bo(old_seq_idx, old_bo_idx, bo_sem_idx, wait=True)
+
+    def start_update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz):
+        bkv_update_ids_ref[bkv_sem_idx] = seq_idx
+        bkv_update_ids_ref[bkv_sem_idx + 2] = offset
+        bkv_update_ids_ref[bkv_sem_idx + 4] = update_sz
+        _update_kv_cache(seq_idx, bkv_sem_idx, offset, update_sz)
+
+    def wait_update_kv_cache(bkv_sem_idx):
+        update_sz = bkv_update_ids_ref[bkv_sem_idx + 4]
+
+        @pl.when(update_sz > 0)
+        def _():
+            seq_idx = bkv_update_ids_ref[bkv_sem_idx]
+            offset = bkv_update_ids_ref[bkv_sem_idx + 2]
+            bkv_update_ids_ref[bkv_sem_idx + 4] = 0
+            _update_kv_cache(seq_idx,
+                             bkv_sem_idx,
+                             offset,
+                             update_sz,
+                             wait=True)
+
+    def load_bq(bq_sem_idx, kv_head_idx, *, actual_bq_sz=bq_sz):
+        q_ref = (bq_x2_ref.bitcast(
+            jnp.uint32).at[bq_sem_idx, kv_head_idx].reshape(
+                bq_sz * num_q_heads_per_kv_head_per_packing, head_dim))
+        return pltpu.bitcast(
+            q_ref[:actual_bq_sz * num_q_heads_per_kv_head_per_packing],
+            q_dtype)
+
+    def strided_load(ref, start, step):
+        assert get_dtype_packing(ref.dtype) == 1
+        assert len(ref.shape) == 2
+        r, l = ref.shape  # noqa
+        assert l % 128 == 0
+        folds = l // 128
+        ref = ref.reshape(r * folds, 128)
+        start *= folds
+        step *= folds
+        vec = jnp.concat([ref[start + i::step] for i in range(folds)], axis=1)
+        return vec
+
+    def strided_load_bkv(bkv_sem_idx, start, step):
+        assert start % kv_packing == 0
+        assert step % kv_packing == 0
+        start //= kv_packing
+        step //= kv_packing
+        kv_ref = (bkv_x2_ref.bitcast(jnp.uint32).at[bkv_sem_idx].reshape(
+            bkv_sz * step, head_dim))
+
+        if kv_packing == 1:
+            k = strided_load(kv_ref, start, step)
+            v = strided_load(kv_ref, start + 1, step)
+
+            k = pltpu.bitcast(k, kv_dtype)
+            v = pltpu.bitcast(v, kv_dtype)
+            return [(k, v)]
+
+        kv = strided_load(kv_ref, start, step)
+        bitwidth = 32 // kv_packing
+
+        # If we want to convert 32-bits into 32//N number of N-bits value, naive
+        # approach would be to perform 32//N number of 32-bits to N-bits conversion.
+        # However, we can reduce number of instructions by utilizing binary tree.
+        # 0: [32]
+        # 1: [16, 16]
+        # ...
+        # log2(32//N): [N, N, ... N]
+
+        def _convert_to_target_bitwidth(val, target_bitwidth: int):
+            curr_dtype = val.dtype
+            curr_bitwidth = get_dtype_bitwidth(curr_dtype)
+            assert target_bitwidth != curr_bitwidth, "No conversion is needed."
+
+            # We split val into two vals (left and right) where each have half of the
+            # original bitwidth.
+            next_bitwidth = curr_bitwidth // 2
+            next_dtype = jnp.dtype(f"uint{next_bitwidth}")
+
+            left = val.astype(next_dtype)
+
+            # Bitwise shift is only supported in uint32.
+            val_u32 = pltpu.bitcast(val, jnp.uint32)
+            val_u32_shifted = val_u32 >> next_bitwidth
+            # Convert back to original dtype.
+            val_shifted = pltpu.bitcast(val_u32_shifted, curr_dtype)
+            right = val_shifted.astype(next_dtype)
+
+            if next_bitwidth == target_bitwidth:
+                k = pltpu.bitcast(left, kv_dtype)
+                v = pltpu.bitcast(right, kv_dtype)
+                return [(k, v)]
+            else:
+                left_out = _convert_to_target_bitwidth(
+                    left,
+                    target_bitwidth=target_bitwidth,
+                )
+                right_out = _convert_to_target_bitwidth(
+                    right,
+                    target_bitwidth=target_bitwidth,
+                )
+                return left_out + right_out
+
+        return _convert_to_target_bitwidth(kv, target_bitwidth=bitwidth)
+
+    def broadcast_minor(src, shape):
+        if src.shape == shape:
+            return src
+        assert src.shape[:-1] == shape[:-1]
+        assert src.shape[-1] % 128 == 0
+        target_minor = align_to(shape[-1], src.shape[-1])
+        # no-op concatenation.
+        return jnp.concatenate(
+            [src for _ in range(target_minor // src.shape[-1])],
+            axis=-1)[..., :shape[-1]]
+
+    def process(static_q_len=None):
+        num_bkv = cdiv(kv_len, bkv_sz)
+        if static_q_len is None:
+            actual_bq_sz = bq_sz
+            num_bq = cdiv(q_len, actual_bq_sz)
+        else:
+            actual_bq_sz = min(bq_sz, static_q_len)
+            num_bq = cdiv(static_q_len, actual_bq_sz)
+
+        def get_next_bq_ids(seq_idx, bq_idx, bq_sem_idx):
+            next_bq_idx = bq_idx + 1
+            is_last_bq = next_bq_idx == num_bq
+            next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx)
+            next_seq_idx = lax.select(is_last_bq, seq_idx + 1, seq_idx)
+            next_bq_sem_idx = lax.select(bq_sem_idx == 0, 1, 0)
+            return next_seq_idx, next_bq_idx, next_bq_sem_idx
+
+        def get_next_bkv_ids(seq_idx, bq_idx, bkv_idx, bkv_sem_idx):
+            next_bkv_idx = bkv_idx + 1
+            is_last_bkv = next_bkv_idx == num_bkv
+            next_bq_idx = lax.select(is_last_bkv, bq_idx + 1, bq_idx)
+            is_last_bq = next_bq_idx == num_bq
+            next_bq_idx = lax.select(is_last_bq, 0, next_bq_idx)
+            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)
+
+            if sliding_window is None:
+                # When sliding window is disabled, starting bkv_idx of next request is
+                # always 0 regardless of seq_idx of next request.
+                next_bkv_idx_start = 0
+            else:
+                # Determine starting bkv_idx of next request based on whether next
+                # request is from the same sequence or next sequence.
+                next_bkv_idx_start = lax.select(
+                    is_last_bq,
+                    next_seq_bkv_idx_start,
+                    bkv_idx_start,
+                )
+            next_bkv_idx = lax.select(is_last_bkv, next_bkv_idx_start,
+                                      next_bkv_idx)
+
+            return next_seq_idx, next_bq_idx, next_bkv_idx, next_bkv_sem_idx
+
+        def compute_with_bq(bq_idx, _):
+            bq_sem_idx = sem_ids_ref[0]
+            next_seq_idx, next_bq_idx, next_bq_sem_idx = get_next_bq_ids(
+                seq_idx, bq_idx, bq_sem_idx)
+
+            # Prefetch next bq
+            @pl.when(next_seq_idx < num_seqs)
+            def prefetch_next_bq():
+                sem_ids_ref[0] = next_bq_sem_idx
+                start_fetch_bq(next_seq_idx, next_bq_idx, next_bq_sem_idx)
+
+            def compute_with_bkv(bkv_idx, _):
+                # Create bitmask for KV.
+                assert bkv_sz % kv_packing == 0
+
+                # Get next bkv ids.
+                bkv_sem_idx = sem_ids_ref[1]
+                next_seq_idx, _, next_bkv_idx, next_bkv_sem_idx = get_next_bkv_ids(
+                    seq_idx, bq_idx, bkv_idx, bkv_sem_idx)
+
+                # Prefetch next bkv
+                @pl.when(next_seq_idx < num_seqs)
+                def prefetch_next_bkv():
+                    sem_ids_ref[1] = next_bkv_sem_idx
+                    start_fetch_bkv(next_seq_idx, next_bkv_idx,
+                                    next_bkv_sem_idx)
+
+                # Wait for cur bq if not ready yet
+                @pl.when(bkv_idx == 0)
+                def wait_cur_bq():
+                    wait_fetch_bq(seq_idx, bq_idx, bq_sem_idx)
+
+                # Wait for cur bkv
+                offset, update_sz = wait_fetch_bkv(seq_idx, bkv_idx,
+                                                   bkv_sem_idx)
+
+                # Start updating bkv to kv cache if applicable.
+                # Only needed in first bq loop.
+                @pl.when(jnp.logical_and(update_sz > 0, bq_idx == 0))
+                def update_cur_bkv_to_cache():
+                    start_update_kv_cache(seq_idx, bkv_sem_idx, offset,
+                                          update_sz)
+
+                debug_print(
+                    "[RPA debug] -----------flash attention-----------")
+                debug_print("[RPA debug] seq_idx={}", seq_idx)
+                debug_print("[RPA debug] bq_idx={}", bq_idx)
+                debug_print("[RPA debug] bkv_idx={}", bkv_idx)
+                if debug_mode:
+                    # Skip flash attention if debug mode is enabled.
+                    return
+
+                # Flash attention with cur bkv and bq
+                # NOTE: kv_packing is divided by 2 because k and v are packed together.
+                prev_bq_shape_0 = None
+                prev_kv_head_bv = None
+                prev_kv_head_idx = None
+                prev_kv_head_p = None
+                prev_kv_head_exp_m_diff = None
+                heads_per_load = max(1, kv_packing // 2)
+                for kv_head_start in range(0, actual_num_kv_heads,
+                                           heads_per_load):
+                    bkv_lst = strided_load_bkv(
+                        bkv_sem_idx,
+                        kv_head_start * 2,
+                        num_kv_heads_x2,
+                    )
+                    assert len(bkv_lst) == heads_per_load
+                    for i in range(heads_per_load):
+                        cur_kv_head_idx = kv_head_start + i
+                        if cur_kv_head_idx >= actual_num_kv_heads:
+                            break
+
+                        cur_kv_head_bq = load_bq(bq_sem_idx,
+                                                 cur_kv_head_idx,
+                                                 actual_bq_sz=actual_bq_sz)
+                        bk, bv = bkv_lst[i]
+                        # FlashAttention is divided into `flash_attention_step1_qk_softmax`
+                        # and `flash_attention_step2_pv` to pipeline the computation.
+                        # `step2_pv` for the previous KV head, which depends on the softmax
+                        # output, is overlapped with `step1_qk_softmax` for the current KV
+                        # head, reducing overall wait times.
+                        cur_kv_head_p, cur_kv_head_exp_m_diff = (
+                            flash_attention_step1_qk_softmax(
+                                cur_kv_head_bq,
+                                bk,
+                                bv,
+                                bq_idx=bq_idx,
+                                bkv_idx=bkv_idx,
+                                kv_head_idx=cur_kv_head_idx,
+                            ))
+                        if prev_bq_shape_0 is not None:
+                            flash_attention_step2_pv(
+                                prev_bq_shape_0,
+                                prev_kv_head_bv,
+                                prev_kv_head_p,
+                                prev_kv_head_exp_m_diff,
+                                bkv_idx=bkv_idx,
+                                kv_head_idx=prev_kv_head_idx,
+                            )
+                        prev_bq_shape_0 = cur_kv_head_bq.shape[0]
+                        prev_kv_head_bv = bv
+                        prev_kv_head_p = cur_kv_head_p
+                        prev_kv_head_exp_m_diff = cur_kv_head_exp_m_diff
+                        prev_kv_head_idx = cur_kv_head_idx
+
+                # Execute pv of last attention head.
+                assert prev_bq_shape_0 is not None
+                flash_attention_step2_pv(
+                    prev_bq_shape_0,
+                    prev_kv_head_bv,
+                    prev_kv_head_p,
+                    prev_kv_head_exp_m_diff,
+                    bkv_idx=bkv_idx,
+                    kv_head_idx=prev_kv_head_idx,
+                )
+
+            lax.fori_loop(0, num_bkv, compute_with_bkv, None, unroll=False)
+
+            # Load acc and calculate final output.
+            acc = acc_ref[...]
+            l = broadcast_minor(l_ref[...], acc.shape)  # noqa
+            out = (lax.div(acc, l) if q_dtype == jnp.float32 else
+                   (acc * pl.reciprocal(l, approx=True)).astype(q_dtype))
+
+            # Wait for previous bo to be fully sent before storing new bo.
+            bo_sem_idx = sem_ids_ref[2]
+            sem_ids_ref[2] = lax.select(bo_sem_idx == 0, 1, 0)
+            wait_send_bo(bo_sem_idx)
+
+            # Store output from acc to bo.
+            bo_x2_ref.at[bo_sem_idx].bitcast(jnp.int32).reshape(
+                actual_num_kv_heads,
+                bq_sz * num_q_heads_per_kv_head_per_packing,
+                head_dim,
+            )[...] = pltpu.bitcast(out, jnp.int32)
+
+            # Send cur bo
+            start_send_bo(seq_idx, bq_idx, bo_sem_idx)
+
+        lax.fori_loop(0, num_bq, compute_with_bq, None, unroll=False)
+
+    ### ------- Kernel start ------- ###
+
+    @pl.when(seq_idx == 0)
+    def prologue():
+        start_fetch_bq(0, 0, 0)
+
+        # Initialize bkv_x2_ref to zeros to avoid NaN issues from accessing
+        # uninitialized memory. Bitcast into int32 to avoid tiling issues.
+        bkv_x2_int32_ref = bkv_x2_ref.bitcast(jnp.int32).reshape(
+            (2, -1, 8, 128))
+        zeros = jnp.zeros(bkv_x2_int32_ref.shape[1:], jnp.int32)
+
+        # To pipeline VST and DMA, we divide the initialization into two steps.
+        bkv_x2_int32_ref[0] = zeros
+        start_fetch_bkv(0, bkv_idx_start, 0)
+        bkv_x2_int32_ref[1] = zeros
+
+    @pl.when(seq_idx < decode_end)
+    def process_decode():
+        process(static_q_len=1)
+
+    @pl.when(jnp.logical_and(decode_end <= seq_idx, seq_idx < prefill_end))
+    def process_prefill():
+        process(static_q_len=chunk_prefill_size)
+
+    @pl.when(jnp.logical_and(prefill_end <= seq_idx, seq_idx < mixed_end))
+    def process_mixed():
+        process()
+
+    @pl.when(seq_idx == num_seqs - 1)
+    def epilogue():
+        for i in range(2):
+            wait_send_bo(i)
+            wait_update_kv_cache(i)
+
+    ### ------- Kernel end ------- ###
+
+
+def merge_kv(
+        k: jax.
+    Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim],
+        v: jax.
+    Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim],
+):
+    assert k.shape == v.shape
+    assert k.dtype == v.dtype
+    max_num_tokens, actual_num_kv_heads, actual_head_dim = k.shape
+    kv_packing = get_dtype_packing(k.dtype)
+    actual_num_kv_heads_x2 = actual_num_kv_heads * 2
+    num_kv_heads_x2 = align_to(actual_num_kv_heads_x2, kv_packing)
+    head_dim = align_to(actual_head_dim, 128)
+    kv = jnp.pad(
+        jnp.concat([k, v],
+                   axis=-1).reshape(max_num_tokens, actual_num_kv_heads_x2,
+                                    actual_head_dim),
+        (
+            (0, 0),
+            (0, num_kv_heads_x2 - actual_num_kv_heads_x2),
+            (0, head_dim - actual_head_dim),
+        ),
+        constant_values=0,
+    ).reshape(
+        max_num_tokens,
+        num_kv_heads_x2 // kv_packing,
+        kv_packing,
+        head_dim,
+    )
+    return kv
+
+
+def prepare_inputs(
+        q: jax.Array,  # [max_num_tokens, actual_num_q_heads, actual_head_dim],
+        k: jax.
+    Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim],
+        v: jax.
+    Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim],
+):
+    max_num_tokens, actual_num_q_heads, actual_head_dim = q.shape
+    actual_num_kv_heads = k.shape[1]
+    assert actual_num_q_heads % actual_num_kv_heads == 0
+    actual_num_q_heads_per_kv_head = actual_num_q_heads // actual_num_kv_heads
+    q_packing = get_dtype_packing(q.dtype)
+    num_q_heads_per_kv_head = align_to(actual_num_q_heads_per_kv_head,
+                                       q_packing)
+    head_dim = align_to(actual_head_dim, 128)
+    q = (
+        jnp.pad(
+            q.reshape(
+                max_num_tokens,
+                actual_num_kv_heads,
+                actual_num_q_heads_per_kv_head,
+                actual_head_dim,
+            ),
+            (
+                (0, 0),
+                (0, 0),
+                (0, num_q_heads_per_kv_head - actual_num_q_heads_per_kv_head),
+                (0, head_dim - actual_head_dim),
+            ),
+            constant_values=0,
+        ).reshape(
+            max_num_tokens,
+            actual_num_kv_heads,
+            num_q_heads_per_kv_head // q_packing,
+            q_packing,
+            head_dim,
+        )
+        # TODO(jevinjiang): Explore fusing swapping non-tiling axis to DMA.
+        .swapaxes(0, 1))
+    # TODO(kyuyeunk, chengjiyao): Add kv quantization here.
+    kv = merge_kv(k, v)
+    return q, kv
+
+
+def prepare_outputs(
+    out,  # [actual_num_kv_heads, max_num_tokens, num_q_heads_per_kv_head // q_packing, q_packing, head_dim]
+    actual_num_q_heads_per_kv_head: int,
+    actual_head_dim: int,
+):
+    (
+        actual_num_kv_heads,
+        max_num_tokens,
+        num_q_heads_per_kv_head_per_q_packing,
+        q_packing,
+        head_dim,
+    ) = out.shape
+    actual_num_q_heads = actual_num_q_heads_per_kv_head * actual_num_kv_heads
+    return (out.swapaxes(0, 1).reshape(
+        max_num_tokens,
+        actual_num_kv_heads,
+        num_q_heads_per_kv_head_per_q_packing * q_packing,
+        head_dim,
+    )[:, :, :actual_num_q_heads_per_kv_head, :actual_head_dim].reshape(
+        max_num_tokens, actual_num_q_heads, actual_head_dim))
+
+
+# Expect to run this validation during runtime.
+def dynamic_validate_inputs(
+    queries: jax.
+    Array,  # [max_num_tokens, actual_num_q_heads, actual_head_dim]
+    keys: jax.Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim]
+    values: jax.
+    Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim]
+    kv_cache: jax.
+    Array,  # [total_num_pages, page_size, num_kv_heads_x2 // kv_packing, kv_packing, head_dim]
+    kv_lens: jax.Array,  # i32[max_num_seqs]
+    page_indices: jax.Array,  # i32[max_num_seqs * pages_per_seq]
+    cu_q_lens: jax.Array,  # i32[max_num_seqs + 1]
+    distribution: jax.Array,  # i32[3]
+    *,
+    sm_scale: float = 1.0,
+    sliding_window: int | None = None,
+    soft_cap: float | None = None,
+    mask_value: float | None = DEFAULT_MASK_VALUE,
+    q_scale: float | None = None,
+    k_scale: float | None = None,
+    v_scale: float | None = None,
+    # Kernel optimization params.
+    chunk_prefill_size: int | None = None,
+    # Kernel tuning params.
+    num_kv_pages_per_block: int | None = None,
+    num_queries_per_block: int | None = None,
+    vmem_limit_bytes: int | None = None,
+    # Debug params.
+    debug_mode: bool = False,
+):
+    q, k, v = queries, keys, values
+    static_validate_inputs(
+        q,
+        k,
+        v,
+        kv_cache,
+        kv_lens,
+        page_indices,
+        cu_q_lens,
+        distribution,
+        sm_scale=sm_scale,
+        sliding_window=sliding_window,
+        soft_cap=soft_cap,
+        mask_value=mask_value,
+        q_scale=q_scale,
+        k_scale=k_scale,
+        v_scale=v_scale,
+        chunk_prefill_size=chunk_prefill_size,
+        num_kv_pages_per_block=num_kv_pages_per_block,
+        num_queries_per_block=num_queries_per_block,
+        vmem_limit_bytes=vmem_limit_bytes,
+        debug_mode=debug_mode,
+    )
+    max_num_tokens = q.shape[0]
+    total_num_pages = kv_cache.shape[0]
+    page_size = kv_cache.shape[1]
+    max_num_seqs = kv_lens.shape[0]
+    num_page_indices = page_indices.shape[0]
+    assert num_page_indices % max_num_seqs == 0
+    pages_per_seq = num_page_indices // max_num_seqs
+
+    i, j, k = distribution
+    if not (i <= j <= k):
+        raise ValueError(f"Invalid distribution: {distribution=}")
+
+    if k > max_num_seqs:
+        raise ValueError(f"num_seqs={k} must be <= {max_num_seqs=}")
+
+    if cu_q_lens[k] > max_num_tokens:
+        raise ValueError(
+            f"Total q tokens {cu_q_lens[k]} must be <= {max_num_tokens=}.")
+    for i in range(k):
+        q_len = cu_q_lens[i + 1] - cu_q_lens[i]
+        kv_len = kv_lens[i]
+        if not (0 < q_len <= kv_len):
+            raise ValueError(
+                f"Require 0 < {q_len=} <= {kv_len=} at sequence {i}.")
+        page_cnt = cdiv(kv_len, page_size)
+        if page_cnt > pages_per_seq:
+            raise ValueError(
+                f"Require {page_cnt=} <= {pages_per_seq=} at sequence {i} where"
+                f" {kv_len=} and {page_size=}.")
+        for p in range(page_cnt):
+            page_idx = page_indices[i * pages_per_seq + p]
+            if not (0 <= page_idx < total_num_pages):
+                raise ValueError(
+                    f"Require 0 <= {page_idx=} < {total_num_pages=} at sequence"
+                    f" {i} where {kv_len=} and {page_size=}.")
+
+
+# Expect to run this validation during compile time.
+def static_validate_inputs(
+    queries: jax.
+    Array,  # [max_num_tokens, actual_num_q_heads, actual_head_dim]
+    keys: jax.Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim]
+    values: jax.
+    Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim]
+    kv_cache: jax.
+    Array,  # [total_num_pages, page_size, num_kv_heads_x2 // kv_packing, kv_packing, head_dim]
+    kv_lens: jax.Array,  # i32[max_num_seqs]
+    page_indices: jax.Array,  # i32[max_num_seqs * pages_per_seq]
+    cu_q_lens: jax.Array,  # i32[max_num_seqs + 1]
+    distribution: jax.Array,  # i32[3]
+    *,
+    sm_scale: float = 1.0,
+    sliding_window: int | None = None,
+    soft_cap: float | None = None,
+    mask_value: float | None = DEFAULT_MASK_VALUE,
+    q_scale: float | None = None,
+    k_scale: float | None = None,
+    v_scale: float | None = None,
+    # Kernel optimization params.
+    chunk_prefill_size: int | None = None,
+    # Kernel tuning params.
+    num_kv_pages_per_block: int | None = None,
+    num_queries_per_block: int | None = None,
+    vmem_limit_bytes: int | None = None,
+    # Debug params.
+    debug_mode: bool = False,
+):
+    """Validate inputs to the RPA kernel statically."""
+    q, k, v = queries, keys, values
+    if not (len(q.shape) == len(k.shape) == len(v.shape) == 3):
+        raise ValueError(
+            f"Expected 3D array for {q.shape=}, {k.shape=}, {v.shape=}")
+    if k.shape != v.shape:
+        raise ValueError(f"Expected {k.shape=} to be equal to {v.shape=}")
+    if not (q.shape[0] == k.shape[0] == v.shape[0]):
+        raise ValueError(
+            f"Expected {q.shape[0]=} to be equal to {k.shape[0]=} and {v.shape[0]=}"
+        )
+    if not (q.shape[2] == k.shape[2] == v.shape[2]):
+        raise ValueError(
+            f"Expected {q.shape[2]=} to be equal to {k.shape[2]=} and {v.shape[2]=}"
+        )
+
+    actual_head_dim = q.shape[2]
+    actual_num_q_heads = q.shape[1]
+    actual_num_kv_heads = k.shape[1]
+
+    if actual_num_q_heads % actual_num_kv_heads != 0:
+        raise ValueError(f"Expected {actual_num_q_heads=} to be divisible by"
+                         f" {actual_num_kv_heads=}.")
+
+    (
+        _,
+        page_size,
+        num_kv_heads_x2_per_kv_packing,
+        kv_packing,
+        head_dim,
+    ) = kv_cache.shape
+
+    if head_dim != align_to(actual_head_dim, 128):
+        raise ValueError(
+            f"Expected {head_dim=} is equal to {align_to(actual_head_dim, 128)=}"
+        )
+    # Note: we expect the kv quantization happens outside of the RPA kernel.
+    if not (kv_cache.dtype == k.dtype == v.dtype):
+        raise ValueError(
+            f"Expected {kv_cache.dtype=} to be equal to {k.dtype=} and {v.dtype=}."
+        )
+    # Integer kv quantization is currently not supported.
+    if not jnp.issubdtype(kv_cache.dtype, jnp.floating):
+        raise ValueError(f"Expected {kv_cache.dtype=} to be a floating point.")
+    if kv_packing != get_dtype_packing(kv_cache.dtype):
+        raise ValueError(
+            f"{kv_packing=} does not match with {kv_cache.dtype=}")
+
+    num_kv_heads_x2 = num_kv_heads_x2_per_kv_packing * kv_packing
+    if num_kv_heads_x2 % 2 != 0:
+        raise ValueError(
+            f"Combined KV heads must be divisible by 2, but got {num_kv_heads_x2}"
+        )
+    if align_to(actual_num_kv_heads * 2, kv_packing) != num_kv_heads_x2:
+        raise ValueError(
+            f"Invalid {num_kv_heads_x2=}, {actual_num_kv_heads=}, {kv_packing=}"
+        )
+
+    if not (jnp.int32 == kv_lens.dtype == page_indices.dtype == cu_q_lens.dtype
+            == distribution.dtype):
+        raise ValueError(
+            f"Expected int32 dtype for {kv_lens.dtype=}, {page_indices.dtype=},"
+            f" {cu_q_lens.dtype=}, {distribution.dtype=}")
+
+    if not (len(kv_lens.shape) == len(page_indices.shape) == len(
+            cu_q_lens.shape) == 1):
+        raise ValueError(
+            f"Expected 1D array for {kv_lens.shape=}, {page_indices.shape=},"
+            f" {cu_q_lens.shape=}")
+
+    max_num_seqs = kv_lens.shape[0]
+    num_page_indices = page_indices.shape[0]
+    if num_page_indices % max_num_seqs != 0:
+        raise ValueError(
+            f"Expected {num_page_indices=} to be divisible by {max_num_seqs=}."
+        )
+    if cu_q_lens.shape != (max_num_seqs + 1, ):
+        raise ValueError(
+            f"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},).")
+    if distribution.shape != (3, ):
+        raise ValueError(f"Expected {distribution.shape=} to be (3,).")
+
+    if page_size % kv_packing != 0:
+        raise ValueError(f"{page_size=} must be divisible by {kv_packing=}.")
+    if sliding_window is not None and sliding_window <= 0:
+        raise ValueError(f"{sliding_window=} must be positive.")
+    if soft_cap is not None and soft_cap == 0.0:
+        raise ValueError(f"{soft_cap=} must not be 0.0.")
+    if chunk_prefill_size is not None and chunk_prefill_size <= 0:
+        raise ValueError(f"{chunk_prefill_size=} must be positive.")
+    if num_kv_pages_per_block is not None:
+        if num_kv_pages_per_block <= 0:
+            raise ValueError(f"{num_kv_pages_per_block=} must be positive.")
+    if num_queries_per_block is not None:
+        if num_queries_per_block <= 0:
+            raise ValueError(f"{num_queries_per_block=} must be positive.")
+    if vmem_limit_bytes is not None and vmem_limit_bytes <= 0:
+        raise ValueError(f"{vmem_limit_bytes=} must be positive.")
+
+    # No constraints for the following inputs.
+    del sm_scale
+    del mask_value
+    del q_scale
+    del k_scale
+    del v_scale
+    del debug_mode
+
+
+def get_kernel_scope_name(bq_size, bkv_p, page_size):
+    return f"RPA-bq_{bq_size}-bkvp_{bkv_p}-p_{page_size}-"
+
+
+@functools.partial(
+    jax.jit,
+    static_argnames=(
+        "sm_scale",
+        "sliding_window",
+        "soft_cap",
+        "mask_value",
+        "q_scale",
+        "k_scale",
+        "v_scale",
+        "chunk_prefill_size",
+        "num_kv_pages_per_block",
+        "num_queries_per_block",
+        "vmem_limit_bytes",
+        "debug_mode",
+    ),
+    donate_argnames=("kv_cache", ),
+)
+def ragged_paged_attention(
+    queries: jax.
+    Array,  # [max_num_tokens, actual_num_q_heads, actual_head_dim]
+    keys: jax.Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim]
+    values: jax.
+    Array,  # [max_num_tokens, actual_num_kv_heads, actual_head_dim]
+    kv_cache: jax.
+    Array,  # [total_num_pages, page_size, num_kv_heads_x2 // kv_packing, kv_packing, head_dim]
+    kv_lens: jax.Array,  # i32[max_num_seqs]
+    page_indices: jax.Array,  # i32[max_num_seqs * pages_per_seq]
+    cu_q_lens: jax.Array,  # i32[max_num_seqs + 1]
+    distribution: jax.Array,  # i32[3]
+    *,
+    sm_scale: float = 1.0,
+    sliding_window: int | None = None,
+    soft_cap: float | None = None,
+    mask_value: float | None = DEFAULT_MASK_VALUE,
+    q_scale: float | None = None,
+    k_scale: float | None = None,
+    v_scale: float | None = None,
+    # Kernel optimization params.
+    chunk_prefill_size: int | None = None,
+    # Kernel tuning params.
+    num_kv_pages_per_block: int | None = None,
+    num_queries_per_block: int | None = None,
+    vmem_limit_bytes: int | None = None,
+    # Debug params.
+    debug_mode: bool = False,
+):
+    """Ragged paged attention 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.
+    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.
+    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,
+        k,
+        v,
+        kv_cache,
+        kv_lens,
+        page_indices,
+        cu_q_lens,
+        distribution,
+        sm_scale=sm_scale,
+        sliding_window=sliding_window,
+        soft_cap=soft_cap,
+        mask_value=mask_value,
+        q_scale=q_scale,
+        k_scale=k_scale,
+        v_scale=v_scale,
+        chunk_prefill_size=chunk_prefill_size,
+        num_kv_pages_per_block=num_kv_pages_per_block,
+        num_queries_per_block=num_queries_per_block,
+        vmem_limit_bytes=vmem_limit_bytes,
+    )
+
+    actual_num_q_heads = q.shape[1]
+    actual_head_dim = q.shape[2]
+    actual_num_kv_heads = k.shape[1]
+
+    actual_num_q_heads_per_kv_head = actual_num_q_heads // actual_num_kv_heads
+    q, kv = prepare_inputs(q, k, v)
+    (
+        _,
+        max_num_tokens,
+        num_q_heads_per_kv_head_per_q_packing,
+        q_packing,
+        head_dim,
+    ) = q.shape
+    page_size = kv_cache.shape[1]
+    max_num_seqs = kv_lens.shape[0]
+    num_page_indices = page_indices.shape[0]
+    assert num_page_indices % max_num_seqs == 0
+    pages_per_seq = num_page_indices // max_num_seqs
+    num_q_heads_per_kv_head = num_q_heads_per_kv_head_per_q_packing * q_packing
+
+    bkv_p = num_kv_pages_per_block
+    bq_sz = num_queries_per_block
+    if bq_sz is None or bkv_p is None:
+        bkv_p, bq_sz = get_tuned_block_sizes(
+            q.dtype,
+            kv_cache.dtype,
+            actual_num_q_heads,
+            actual_num_kv_heads,
+            head_dim,
+            page_size,
+            max_num_tokens,
+            pages_per_seq,
+            sliding_window,
+        )
+    bkv_sz = bkv_p * page_size
+    if vmem_limit_bytes is None:
+        # TODO (jevinjiang/jacobplatin): change this to use
+        # `get_vmem_estimate_bytes` when VREG spilling is fixed.
+        vmem_limit_bytes = DEFAULT_VMEM_LIMIT_BYTES
+    grid = (distribution[2], )
+
+    in_specs = [
+        pl.BlockSpec(memory_space=pltpu.HBM),
+        pl.BlockSpec(memory_space=pltpu.HBM),
+        pl.BlockSpec(memory_space=pltpu.HBM),
+    ]
+
+    out_specs = [
+        pl.BlockSpec(memory_space=pltpu.HBM),
+        pl.BlockSpec(memory_space=pltpu.HBM),
+    ]
+
+    bkv_double_buf = pltpu.VMEM(
+        (2, bkv_sz, *kv_cache.shape[2:]),
+        kv_cache.dtype,
+    )
+
+    bq_double_buf = pltpu.VMEM(
+        (2, actual_num_kv_heads, bq_sz, *q.shape[2:]),
+        q.dtype,
+    )
+
+    bo_double_buf = bq_double_buf
+
+    l_scratch = pltpu.VMEM(
+        (actual_num_kv_heads, bq_sz * num_q_heads_per_kv_head, 128),
+        jnp.float32,
+    )
+    m_scratch = l_scratch
+
+    acc_scratch = pltpu.VMEM(
+        (actual_num_kv_heads, bq_sz * num_q_heads_per_kv_head, head_dim),
+        jnp.float32,
+    )
+
+    scratch_shapes = [
+        bkv_double_buf,  # Double buffering for kv block.
+        bq_double_buf,  # Double buffering for q block.
+        bo_double_buf,  # Double buffering for output block.
+        # Semaphores for double buffering of bkv, bq, bo and bkv_update.
+        pltpu.SemaphoreType.DMA((4, 2)),
+        # Intermediate buffers per kv head for flash attention.
+        l_scratch,
+        m_scratch,
+        acc_scratch,
+    ]
+
+    scalar_prefetches = (
+        kv_lens,
+        # TODO(jevinjiang): can we use ragged page_indices to save some smem?
+        page_indices,
+        cu_q_lens,
+        distribution,
+        # (bq_sem_idx, bkv_sem_idx, bo_sem_idx)
+        jnp.zeros((3, ), jnp.int32),
+        # (bo_sem_0_seq_idx, bo_sem_1_seq_idx, bo_sem_0_bo_idx, bo_sem_1_bo_idx)
+        jnp.full((4, ), -1, jnp.int32),
+        # (bkv_sem_0_seq_idx, bkv_sem_1_seq_idx, bkv_sem_0_offset, bkv_sem_1_offset, bkv_sem_0_sz, bkv_sem_1_sz)
+        jnp.full((6, ), -1, jnp.int32),
+    )
+
+    scope_name = get_kernel_scope_name(bq_sz, bkv_p, page_size)
+    kernel = pl.pallas_call(
+        functools.partial(
+            _ragged_paged_attention_kernel,
+            sm_scale=sm_scale,
+            sliding_window=sliding_window,
+            soft_cap=soft_cap,
+            mask_value=mask_value,
+            q_scale=q_scale,
+            k_scale=k_scale,
+            v_scale=v_scale,
+            chunk_prefill_size=chunk_prefill_size,
+            bq_sz=bq_sz,
+            bkv_p=bkv_p,
+            debug_mode=debug_mode,
+        ),
+        grid_spec=pltpu.PrefetchScalarGridSpec(
+            num_scalar_prefetch=len(scalar_prefetches),
+            in_specs=in_specs,
+            out_specs=out_specs,
+            grid=grid,
+            scratch_shapes=scratch_shapes,
+        ),
+        compiler_params=pltpu.CompilerParams(
+            # TODO(jevinjiang): since each sequence depends on the previous
+            # one, we need some extra work to support Megacore mode.
+            dimension_semantics=("arbitrary", ),
+            vmem_limit_bytes=vmem_limit_bytes,
+        ),
+        out_shape=[
+            jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype),
+            jax.ShapeDtypeStruct(shape=kv_cache.shape, dtype=kv_cache.dtype),
+        ],
+        input_output_aliases={
+            7: 0,
+            9: 1
+        },
+        name=scope_name,
+    )
+
+    output, updated_kv_cache = kernel(*scalar_prefetches, q, kv, kv_cache)
+    return (
+        prepare_outputs(output, actual_num_q_heads_per_kv_head,
+                        actual_head_dim),
+        updated_kv_cache,
+    )