# Copyright 2025 The JAX Authors.
#
# 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
#
#     https://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._src import dtypes
from jax.experimental import pallas as pl
from jax.experimental.pallas import tpu as pltpu
from jax.experimental.pallas.ops.tpu.ragged_paged_attention.tuned_block_sizes import \
    get_tuned_block_sizes

DEFAULT_MASK_VALUE = -0.7 * float(jnp.finfo(jnp.dtype("float32")).max)


class MultiPageAsyncCopyDescriptor:
    """Descriptor for async copy of multiple K/V pages from HBM."""

    def __init__(
            self,
            pages_hbm_ref,  # [total_num_pages, page_size, num_combined_kv_heads_per_blk, head_dim]
            vmem_buf,  # [num_kv_pages_per_blk, page_size, num_combined_kv_heads_per_blk, head_dim]
            sem,
            page_indices_ref,  # i32[max_num_seqs, pages_per_seq]
            metadata,  # [seq_idx, start_page_idx, end_page_idx]
    ):
        self._vmem_buf = vmem_buf
        seq_id, start_page_idx, end_page_idx = metadata
        self._async_copies = []
        # TODO(jevinjiang): Only fetch dynamic shape in need! This will insert
        # a bunch of if-ops. Check the performance when we have benchmarking setup.
        for i in range(vmem_buf.shape[0]):
            page_idx = start_page_idx + i
            page_idx = jax.lax.select(page_idx < end_page_idx, page_idx, 0)
            self._async_copies.append(
                pltpu.make_async_copy(
                    pages_hbm_ref.at[page_indices_ref[seq_id, page_idx]],
                    vmem_buf.at[i],
                    sem,
                ))

    def start(self):
        """Starts the async copies."""
        for async_copy in self._async_copies:
            async_copy.start()

    def wait(self):
        for async_copy in self._async_copies:
            async_copy.wait()
        return self._vmem_buf


def ref_ragged_paged_attention(
    queries: jax.Array,  # [max_num_batched_tokens, num_q_heads, head_dim]
    kv_pages: jax.
    Array,  # [total_num_pages, page_size, num_combined_kv_heads, 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]
    num_seqs: jax.Array,  # i32[1],
    *,
    sm_scale: float = 1.0,
    sliding_window: int | None = None,
    soft_cap: float | None = None,
    mask_value: float | None = DEFAULT_MASK_VALUE,
    k_scale: float | None = None,
    v_scale: float | None = None,
):
    static_validate_inputs(
        queries,
        kv_pages,
        kv_lens,
        page_indices,
        cu_q_lens,
        num_seqs,
        sm_scale=sm_scale,
        k_scale=k_scale,
        v_scale=v_scale,
        sliding_window=sliding_window,
        soft_cap=soft_cap,
        mask_value=mask_value,
    )
    if mask_value is None:
        mask_value = DEFAULT_MASK_VALUE
    _, _, num_combined_kv_heads, head_dim = kv_pages.shape
    assert num_combined_kv_heads % 2 == 0
    num_kv_heads = num_combined_kv_heads // 2
    num_q_heads = queries.shape[1]
    assert num_q_heads % num_kv_heads == 0
    num_query_per_kv = num_q_heads // num_kv_heads
    outputs = []
    for i in range(num_seqs[0]):
        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 = page_indices[i]
        q = queries[q_start:q_end]
        k = kv_pages[indices, :, 0::2, :].reshape(-1, num_kv_heads,
                                                  head_dim)[:kv_len]
        v = kv_pages[indices, :, 1::2, :].reshape(-1, num_kv_heads,
                                                  head_dim)[:kv_len]
        if k_scale is not None:
            k = k.astype(jnp.float32) * k_scale
            k = k.astype(q.dtype)
        if v_scale is not None:
            v = v.astype(jnp.float32) * v_scale
            v = v.astype(q.dtype)
        k = jnp.repeat(k, num_query_per_kv, axis=1)
        v = jnp.repeat(v, num_query_per_kv, axis=1)
        attn = jnp.einsum("qhd,khd->hqk",
                          q,
                          k,
                          preferred_element_type=jnp.float32)
        attn *= sm_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)
        outputs.append(out)

    return jnp.concatenate(outputs, axis=0)


# Expect to run these checks during runtime.
def dynamic_validate_inputs(
    q: jax.Array,  # [max_num_batched_tokens, num_q_heads, head_dim]
    kv_pages: jax.
    Array,  # [total_num_pages, page_size, num_combined_kv_heads, 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]
    num_seqs: jax.Array,  # i32[1]
    *,
    # These inputs are optional. If not specified, we will not validate them.
    sm_scale: float | None = None,
    sliding_window: int | None = None,
    soft_cap: float | None = None,
    mask_value: float | None = None,
    k_scale: float | None = None,
    v_scale: float | 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,
):
    static_validate_inputs(
        q,
        kv_pages,
        kv_lens,
        page_indices,
        cu_q_lens,
        num_seqs,
        sm_scale=sm_scale,
        sliding_window=sliding_window,
        soft_cap=soft_cap,
        mask_value=mask_value,
        k_scale=k_scale,
        v_scale=v_scale,
        num_kv_pages_per_block=num_kv_pages_per_block,
        num_queries_per_block=num_queries_per_block,
        vmem_limit_bytes=vmem_limit_bytes,
    )
    max_num_batched_tokens = q.shape[0]
    page_size = kv_pages.shape[1]
    max_num_seqs, pages_per_seq = page_indices.shape
    if num_seqs[0] > max_num_seqs:
        raise ValueError(
            f"{num_seqs[0]=} must be less or equal to {max_num_seqs=}")
    max_kv_len = jnp.max(kv_lens)
    min_pages_per_seq = cdiv(max_kv_len, page_size)
    if pages_per_seq < min_pages_per_seq:
        raise ValueError(
            f"{pages_per_seq=} must be greater or equal to"
            f" {min_pages_per_seq=} given {max_kv_len=} and {page_size=}.")
    if cu_q_lens[num_seqs[0]] > max_num_batched_tokens:
        raise ValueError(
            f"Total q tokens {cu_q_lens[num_seqs[0]]} must be less or equal to"
            f" {max_num_batched_tokens=}.")
    for i in range(num_seqs[0]):
        q_len = cu_q_lens[i + 1] - cu_q_lens[i]
        kv_len = kv_lens[i]
        if q_len > kv_len:
            raise ValueError(
                f"{q_len=} must be less or equal to {kv_len=} at sequence {i}."
            )


# Expect to run these checks during compile time.
def static_validate_inputs(
    q: jax.Array,  # [max_num_batched_tokens, num_q_heads, head_dim]
    kv_pages: jax.
    Array,  # [total_num_pages, page_size, num_combined_kv_heads, 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]
    num_seqs: jax.Array,  # i32[1]
    *,
    # These inputs are optional. If not specified, we will not validate them.
    sm_scale: float | None = None,
    sliding_window: int | None = None,
    soft_cap: float | None = None,
    mask_value: float | None = None,
    k_scale: float | None = None,
    v_scale: float | 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,
):
    _, num_q_heads, head_dim = q.shape
    _, _, num_combined_kv_heads, head_dim_k = kv_pages.shape
    assert num_combined_kv_heads % 2 == 0
    assert isinstance(k_scale, float) or k_scale is None
    assert isinstance(v_scale, float) or v_scale is None
    num_kv_heads = num_combined_kv_heads // 2
    max_num_seqs, pages_per_seq = page_indices.shape
    if num_seqs.shape != (1, ):
        raise ValueError(f"{num_seqs.shape=} must be (1,)")
    if head_dim_k != head_dim:
        raise ValueError(
            f"Q head_dim {head_dim} must be the same as that of K/V {head_dim_k}."
        )
    if kv_lens.shape != (max_num_seqs, ):
        raise ValueError(
            f"Expected {kv_lens.shape=} to be ({max_num_seqs},) where"
            " `max_num_seqs` is `page_indices.shape[0]`.")
    if cu_q_lens.shape != (max_num_seqs + 1, ):
        raise ValueError(
            f"Expected {cu_q_lens.shape=} to be ({max_num_seqs + 1},)  where"
            " `max_num_seqs` is `page_indices.shape[0]`.")
    if (kv_lens.dtype != jnp.int32 or page_indices.dtype != jnp.int32
            or cu_q_lens.dtype != jnp.int32):
        raise ValueError(
            "The dtype of `kv_lens`, `page_indices`, and `cu_q_lens` must be"
            f" int32. Got {kv_lens.dtype=}, {page_indices.dtype=},"
            f" {cu_q_lens.dtype=}.")
    if num_q_heads % num_kv_heads != 0:
        raise ValueError(
            f"{num_q_heads=} must be divisible by {num_kv_heads=}")
    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 (num_kv_pages_per_block is not None
            and not 0 < num_kv_pages_per_block <= pages_per_seq):
        raise ValueError(
            f"{num_kv_pages_per_block=} must be in range (0, {pages_per_seq}]."
        )
    if num_queries_per_block is not None and 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.")
    del sm_scale  # No constraints on sm_scale.
    del mask_value  # No consstraints on mask_value.


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]
    seq_buf_idx_ref,
    # TODO(jevinjiang): if OOM in SMEM, consider pack to other scalar refs.
    num_seqs_ref,
    # Input
    q_ref,  # [num_q_per_blk, num_q_heads_per_blk, head_dim]
    kv_pages_hbm_ref,  # [total_num_pages, page_size, num_combined_kv_heads, head_dim]
    # Output
    o_ref,  # [num_q_per_blk, num_q_heads_per_blk, head_dim]
    # Scratch
    kv_bufs,  # [2, num_kv_pages_per_blk, page_size, num_combined_kv_heads_per_blk, head_dim]
    sems,  # [2, 2]
    l_ref,  # [num_kv_heads_per_blk, num_q_per_blk * num_q_heads_per_kv_head, 128]
    m_ref,  # [num_kv_heads_per_blk, num_q_per_blk * num_q_heads_per_kv_head, 128]
    acc_ref,  # [num_q_per_blk, num_q_heads_per_blk, head_dim]
    *,
    sm_scale: float,
    sliding_window: int | None = None,
    soft_cap: float | None = None,
    mask_value: float | None = DEFAULT_MASK_VALUE,
    k_scale: float | None = None,
    v_scale: float | None = None,
):
    if mask_value is None:
        mask_value = DEFAULT_MASK_VALUE
    num_q_per_blk, num_q_heads_per_blk, head_dim = q_ref.shape
    pages_per_seq = page_indices_ref.shape[-1]
    num_seqs = num_seqs_ref[0]
    _, num_kv_pages_per_blk, page_size, num_combined_kv_heads_per_blk, _ = (
        kv_bufs.shape)
    num_kv_heads_per_blk = num_combined_kv_heads_per_blk // 2
    num_kv_per_blk = num_kv_pages_per_blk * page_size
    num_q_heads_per_kv_head = num_q_heads_per_blk // num_kv_heads_per_blk
    heads_blk_idx, q_blk_idx = (
        pl.program_id(0),
        pl.program_id(1),
    )
    num_heads_blks = pl.num_programs(0)
    init_seq_idx = seq_buf_idx_ref[0]
    init_buf_idx = seq_buf_idx_ref[1]
    q_len_start = q_blk_idx * num_q_per_blk
    q_len_end = q_len_start + num_q_per_blk

    def create_kv_async_copy_descriptors(heads_blk_idx, seq_idx, kv_blk_idx,
                                         buf_idx):
        start_kv_page_idx = kv_blk_idx * num_kv_pages_per_blk
        end_kv_page_idx = jnp.minimum(pages_per_seq,
                                      cdiv(kv_lens_ref[seq_idx], page_size))
        metadata = (seq_idx, start_kv_page_idx, end_kv_page_idx)
        heads_start = heads_blk_idx * num_combined_kv_heads_per_blk
        async_copy_kv = MultiPageAsyncCopyDescriptor(
            kv_pages_hbm_ref.
            at[:, :,
               pl.ds(heads_start, num_combined_kv_heads_per_blk), :],
            kv_bufs.at[buf_idx],
            sems.at[buf_idx],
            page_indices_ref,
            metadata,
        )
        return async_copy_kv

    # TODO(jevinjiang): Add these to Mosaic:
    # 1. Support arbitrary strided load/store for int4 and int8 dtype.
    # 2. Support arbitrary strided load/store for any last dimension.
    def strided_load_kv(ref, start, step):
        packing = get_dtype_packing(ref.dtype)
        if packing == 1:
            return [ref[start::step, :]], [ref[start + 1::step, :]]
        assert packing in (2, 4, 8)
        assert step % packing == 0
        k_list, v_list = [], []
        b_start = start // packing
        b_step = step // packing
        b_ref = ref.bitcast(jnp.uint32)
        b = b_ref[b_start::b_step, :]

        # TODO(chengjiyao): use the general strided loading logic for bf16 after
        # fixing the issue in mosaic's infer vector layout pass
        if ref.dtype == jnp.bfloat16:
            bk = b << 16
            bv = b & jnp.uint32(0xFFFF0000)
            k = pltpu.bitcast(bk, jnp.float32).astype(jnp.bfloat16)
            v = pltpu.bitcast(bv, jnp.float32).astype(jnp.bfloat16)
            k_list.append(k)
            v_list.append(v)
        else:
            bitwidth = 32 // packing
            bitcast_dst_dtype = jnp.dtype(f"uint{bitwidth}")
            for i in range(0, packing, 2):
                bk = b >> (i * bitwidth)
                k = pltpu.bitcast(bk.astype(bitcast_dst_dtype), ref.dtype)
                k_list.append(k)
                bv = b >> ((i + 1) * bitwidth)
                v = pltpu.bitcast(bv.astype(bitcast_dst_dtype), ref.dtype)
                v_list.append(v)

        return k_list, v_list

    def fold_on_2nd_minor(vec):
        assert vec.dtype == jnp.bfloat16 or vec.dtype == jnp.float32
        assert len(vec.shape) >= 2
        last_dim = vec.shape[-1]
        packing = get_dtype_packing(vec.dtype)
        if vec.shape[-2] % packing != 0:
            vec = vec.astype(jnp.float32)
        return vec.reshape(-1, last_dim)

    @pl.when(heads_blk_idx + q_blk_idx == 0)
    def prefetch_first_kv_blk():
        async_copy_kv = create_kv_async_copy_descriptors(
            heads_blk_idx, init_seq_idx, 0, init_buf_idx)
        async_copy_kv.start()

    def is_cur_q_blk_needed(q_states):
        done, cur_seq_idx, _ = q_states
        should_run = jnp.logical_and(q_len_start < cu_q_lens_ref[num_seqs],
                                     cur_seq_idx < num_seqs)
        return jnp.logical_and(done == 0, should_run)

    def compute_with_cur_q_blk(q_states):
        done, cur_seq_idx, cur_buf_idx = q_states
        q_start = cu_q_lens_ref[cur_seq_idx]
        q_end = cu_q_lens_ref[cur_seq_idx + 1]
        q_len = q_end - q_start
        kv_len = kv_lens_ref[cur_seq_idx]

        def get_next_prefetch_ids(heads_blk_idx, cur_seq_idx, kv_blk_idx,
                                  cur_buf_idx):
            next_kv_blk_idx = kv_blk_idx + 1
            is_last_kv_blk = next_kv_blk_idx * num_kv_per_blk >= kv_len
            next_kv_blk_idx = lax.select(
                is_last_kv_blk,
                0,
                next_kv_blk_idx,
            )
            is_cur_seq_end_in_cur_q_blk = q_end <= q_len_end
            next_seq_idx = lax.select(
                is_last_kv_blk,
                lax.select(is_cur_seq_end_in_cur_q_blk, cur_seq_idx + 1,
                           cur_seq_idx),
                cur_seq_idx,
            )
            is_last_seq = next_seq_idx == num_seqs
            next_seq_idx = lax.select(
                is_last_seq,
                0,
                next_seq_idx,
            )
            next_heads_blk_idx = lax.select(
                is_last_seq,
                heads_blk_idx + 1,
                heads_blk_idx,
            )
            next_buf_idx = lax.select(cur_buf_idx == 0, 1, 0)
            return next_heads_blk_idx, next_seq_idx, next_kv_blk_idx, next_buf_idx

        def flash_attention(
            q,  # [num_q_per_blk * num_q_heads_per_kv_head, head_dim]
            k,  # [num_kv_per_blk, head_dim]
            v,  # [num_kv_per_blk, head_dim]
            head_l_ref,  # [num_q_per_blk * num_q_heads_per_kv_head, 128]
            head_m_ref,  # [num_q_per_blk * num_q_heads_per_kv_head, 128]
            head_acc_ref,  # [num_q_per_blk, num_q_heads_per_kv_head, head_dim]
            *,
            kv_blk_idx,
        ):
            assert q.shape == (
                num_q_per_blk * num_q_heads_per_kv_head,
                head_dim,
            )
            assert (k.shape == v.shape == (
                num_kv_per_blk,
                head_dim,
            ))
            assert k.dtype == v.dtype
            assert (head_m_ref.shape == head_l_ref.shape == (
                num_q_per_blk * num_q_heads_per_kv_head,
                128,
            ))
            assert head_acc_ref.shape == (
                num_q_per_blk,
                num_q_heads_per_kv_head,
                head_dim,
            )
            kv_len_start = kv_blk_idx * num_kv_per_blk

            def masked_store(ref, val, start, end, group=1):
                iota = lax.broadcasted_iota(jnp.int32, ref.shape, 0) // group
                mask = jnp.logical_and(iota >= start, iota < end)
                pl.store(ref,
                         idx=tuple(slice(None) for _ in ref.shape),
                         val=val,
                         mask=mask)

            def load_with_init(ref, init_val):
                return jnp.where(kv_blk_idx == 0, jnp.full_like(ref, init_val),
                                 ref[...])

            # kv lens will be contracting dim, we should mask out the NaNs.
            kv_mask = (lax.broadcasted_iota(jnp.int32, k.shape, 0)
                       < kv_len - kv_len_start)
            k = jnp.where(kv_mask, k.astype(jnp.float32), 0).astype(k.dtype)
            v = jnp.where(kv_mask, v.astype(jnp.float32), 0).astype(v.dtype)

            qk = (jnp.einsum(
                "nd,md->nm", q, k, preferred_element_type=jnp.float32) *
                  sm_scale)
            store_start = jnp.maximum(q_start - q_len_start, 0)
            store_end = jnp.minimum(q_end - q_len_start, num_q_per_blk)

            row_ids = (
                (kv_len - q_len) + q_len_start - q_start +
                jax.lax.broadcasted_iota(
                    jnp.int32,
                    (num_q_per_blk * num_q_heads_per_kv_head, num_kv_per_blk),
                    0,
                ) // num_q_heads_per_kv_head)
            col_ids = kv_len_start + jax.lax.broadcasted_iota(
                jnp.int32,
                (num_q_per_blk * num_q_heads_per_kv_head, num_kv_per_blk),
                1,
            )
            causal_mask = row_ids < col_ids
            if sliding_window is not None:
                causal_mask = jnp.logical_or(
                    causal_mask, row_ids - sliding_window >= col_ids)
            if soft_cap is not None:
                qk = soft_cap * jnp.tanh(qk / soft_cap)
            qk += jnp.where(causal_mask, mask_value, 0.0)
            m_curr = jnp.max(qk, axis=1, keepdims=True)
            s_curr = jnp.exp(qk - m_curr)
            qkv = jnp.dot(s_curr, v, preferred_element_type=jnp.float32)
            lm_store_shape = head_m_ref.shape
            m_curr = jnp.broadcast_to(m_curr, lm_store_shape)
            l_curr = jnp.broadcast_to(s_curr.sum(axis=1, keepdims=True),
                                      lm_store_shape)
            m_prev = load_with_init(head_m_ref, -jnp.inf)
            l_prev = load_with_init(head_l_ref, 0.0)
            m_next = jnp.maximum(m_prev, m_curr)
            masked_store(head_m_ref, m_next, store_start, store_end,
                         num_q_heads_per_kv_head)
            alpha = jnp.exp(m_prev - m_next)
            beta = jnp.exp(m_curr - m_next)
            l_alpha = alpha * l_prev
            l_next = l_alpha + beta * l_curr
            l_next_safe = jnp.where(l_next == 0.0, 1.0, l_next)
            masked_store(
                head_l_ref,
                l_next_safe,
                store_start,
                store_end,
                num_q_heads_per_kv_head,
            )

            def broadcast_to_shape(arr, shape):
                if arr.shape == shape:
                    return arr
                assert len(arr.shape) == len(shape)
                assert arr.shape[0] == shape[0]
                assert shape[1] % arr.shape[1] == 0
                # no-op concatenation.
                return jnp.concatenate(
                    [arr for _ in range(shape[1] // arr.shape[1])], axis=1)

            o_curr = load_with_init(head_acc_ref, 0.0).reshape(-1, head_dim)
            l_alpha = broadcast_to_shape(l_alpha, qkv.shape)
            beta = broadcast_to_shape(beta, qkv.shape)
            l_next_safe = broadcast_to_shape(l_next_safe, qkv.shape)
            out = lax.div(
                l_alpha * o_curr + beta * qkv,
                l_next_safe,
            )
            masked_store(
                head_acc_ref,
                out.reshape(head_acc_ref.shape),
                store_start,
                store_end,
            )

        def is_valid_kv_blk_in_cur_seq(kv_states):
            kv_blk_idx, _ = kv_states
            return kv_blk_idx * num_kv_per_blk < kv_len

        def compute_with_kv_blk_in_cur_seq(kv_states):
            kv_blk_idx, cur_buf_idx = kv_states
            next_heads_blk_idx, next_seq_idx, next_kv_blk_idx, next_buf_idx = (
                get_next_prefetch_ids(heads_blk_idx, cur_seq_idx, kv_blk_idx,
                                      cur_buf_idx))

            @pl.when(next_heads_blk_idx < num_heads_blks)
            def prefetch_next_kv_blk():
                # TODO(jevinjiang): reuse the same buffer if it is already prefetched!
                # TODO(jevinjiang): only fetch effective dynamic size to hold kv_len and
                # DMA to fixed size buffer!
                next_async_copy_kv = create_kv_async_copy_descriptors(
                    next_heads_blk_idx, next_seq_idx, next_kv_blk_idx,
                    next_buf_idx)
                next_async_copy_kv.start()

            cur_async_copy_kv = create_kv_async_copy_descriptors(
                heads_blk_idx, cur_seq_idx, kv_blk_idx, cur_buf_idx)
            kv_ref = cur_async_copy_kv.wait().reshape(
                num_kv_pages_per_blk * page_size *
                num_combined_kv_heads_per_blk,
                head_dim,
            )
            kv_packing = get_dtype_packing(kv_ref.dtype)
            # NOTE: kv_packing is divided by 2 because k and v are packed together.
            kv_load_step = max(1, kv_packing // 2)
            for kv_head_chunk_idx in range(0, num_kv_heads_per_blk,
                                           kv_load_step):
                k_list, v_list = strided_load_kv(
                    kv_ref, kv_head_chunk_idx * 2,
                    num_combined_kv_heads_per_blk)
                for step_idx in range(kv_load_step):
                    k = k_list[step_idx]
                    v = v_list[step_idx]
                    if k_scale is not None:
                        # NOTE: Conversion between arbitrary data types is not supported.
                        # That's why it is converted to float32 first.
                        k = k.astype(jnp.float32) * k_scale
                        k = k.astype(q_ref.dtype)
                    if v_scale is not None:
                        v = v.astype(jnp.float32) * v_scale
                        v = v.astype(q_ref.dtype)
                    kv_head_idx = kv_head_chunk_idx + step_idx
                    q_head_idx = kv_head_idx * num_q_heads_per_kv_head
                    # TODO(jevinjiang): extra handling for packed type that can start at
                    # unaligned position!
                    q = fold_on_2nd_minor(q_ref[:, q_head_idx:q_head_idx +
                                                num_q_heads_per_kv_head, :])
                    flash_attention(
                        q,
                        k,
                        v,
                        l_ref.at[kv_head_idx],
                        m_ref.at[kv_head_idx],
                        acc_ref.at[:, q_head_idx:q_head_idx +
                                   num_q_heads_per_kv_head, :],
                        kv_blk_idx=kv_blk_idx,
                    )
            return kv_blk_idx + 1, next_buf_idx

        _, next_buf_idx = lax.while_loop(
            is_valid_kv_blk_in_cur_seq,
            compute_with_kv_blk_in_cur_seq,
            (0, cur_buf_idx),  # (kv_blk_idx, buf_idx)
        )
        next_seq_idx = lax.select(q_end <= q_len_end, cur_seq_idx + 1,
                                  cur_seq_idx)
        done = lax.select(q_end < q_len_end, done, 1)
        return done, next_seq_idx, next_buf_idx

    _, seq_idx, buf_idx = lax.while_loop(
        is_cur_q_blk_needed,
        compute_with_cur_q_blk,
        (0, init_seq_idx, init_buf_idx),  # (done, seq_idx, buf_idx)
    )
    # Reset seq_idx for next kv_heads_blk if run out of seqs!
    seq_buf_idx_ref[0] = lax.select(seq_idx < num_seqs, seq_idx, 0)
    seq_buf_idx_ref[1] = buf_idx
    o_ref[...] = acc_ref[...].astype(q_ref.dtype)


def cdiv(a, b):
    assert b != 0
    return (a + b - 1) // b


def get_dtype_packing(dtype):
    bits = (dtypes.bit_width(dtype)
            if hasattr(dtypes, "bit_width") else dtypes.itemsize_bits(dtype))
    return 32 // bits


def get_min_heads_per_blk(num_q_heads, num_combined_kv_heads, q_dtype,
                          kv_dtype):
    q_packing = get_dtype_packing(q_dtype)
    kv_packing = get_dtype_packing(kv_dtype)

    def can_be_xla_fully_tiled(x, packing):
        if x % packing != 0:
            return False
        x //= packing
        return x in (1, 2, 4, 8) or x % 8 == 0

    # TODO(jevinjiang): support unaligned number of heads!
    if not can_be_xla_fully_tiled(num_combined_kv_heads, kv_packing):
        raise ValueError(
            f"Not implemented: {num_combined_kv_heads=} can not be XLA fully tiled."
        )
    assert num_combined_kv_heads % 2 == 0
    num_kv_heads = num_combined_kv_heads // 2
    assert num_q_heads % num_kv_heads == 0
    ratio = num_q_heads // num_kv_heads
    # TODO(jevinjiang): we can choose smaller tiling for packed type if large
    # second minor tiling is not on.
    max_combined_kv_tiling = 8 * kv_packing
    min_combined_kv_heads = (max_combined_kv_tiling if num_combined_kv_heads %
                             max_combined_kv_tiling == 0 else
                             num_combined_kv_heads)
    min_q_heads = min_combined_kv_heads // 2 * ratio
    if can_be_xla_fully_tiled(min_q_heads, q_packing):
        return min_q_heads, min_combined_kv_heads
    return num_q_heads, num_combined_kv_heads


@functools.partial(
    jax.jit,
    static_argnames=[
        "sm_scale",
        "mask_value",
        "num_kv_pages_per_block",
        "num_queries_per_block",
        "vmem_limit_bytes",
        "sliding_window",
        "soft_cap",
        "k_scale",
        "v_scale",
    ],
)
def ragged_paged_attention(
    q: jax.Array,  # [max_num_batched_tokens, num_q_heads, head_dim]
    # TODO(jevinjiang): create a write_to_kv_cache kernel!
    kv_pages: jax.
    Array,  # [total_num_pages, page_size, num_combined_kv_heads, 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]
    num_seqs: jax.Array,  # i32[1]
    *,
    sm_scale: float = 1.0,
    sliding_window: int | None = None,
    soft_cap: float | None = None,
    mask_value: float | None = DEFAULT_MASK_VALUE,
    k_scale: float | None = None,
    v_scale: float | None = None,
    num_kv_pages_per_block: int | None = None,
    num_queries_per_block: int | None = None,
    vmem_limit_bytes: int | None = None,
):
    """Ragged paged attention that supports mixed prefill and decode.

  Args:
    q: concatenated all sequences' queries.
    kv_pages: paged KV cache. Normally in HBM.
    kv_lens: padded kv lengths. Only the first num_seqs values are valid.
    page_indices: the first index indicates which page to use in the kv cache
      for each sequence. Only the first num_seqs values are valid.
    cu_q_lens: the cumulative sum of the effective query lengths. Similar to
      kv_lens, only the first num_seqs+1 values are valid.
    num_seqs: the dynamic 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.
    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.

  Returns:
    The output of the attention.
  """
    static_validate_inputs(
        q,
        kv_pages,
        kv_lens,
        page_indices,
        cu_q_lens,
        num_seqs,
        sm_scale=sm_scale,
        sliding_window=sliding_window,
        soft_cap=soft_cap,
        mask_value=mask_value,
        k_scale=k_scale,
        v_scale=v_scale,
        num_kv_pages_per_block=num_kv_pages_per_block,
        num_queries_per_block=num_queries_per_block,
        vmem_limit_bytes=vmem_limit_bytes,
    )
    if mask_value is None:
        mask_value = DEFAULT_MASK_VALUE
    num_q_tokens, num_q_heads, head_dim = q.shape
    _, page_size, num_combined_kv_heads, _ = kv_pages.shape
    assert num_combined_kv_heads % 2 == 0
    num_kv_heads = num_combined_kv_heads // 2
    _, pages_per_seq = page_indices.shape
    num_q_heads_per_blk, num_combined_kv_heads_per_blk = get_min_heads_per_blk(
        num_q_heads, num_combined_kv_heads, q.dtype, kv_pages.dtype)
    num_q_per_blk = num_queries_per_block
    num_kv_pages_per_blk = num_kv_pages_per_block
    if num_q_per_blk is None or num_kv_pages_per_blk is None:
        num_kv_pages_per_blk, num_q_per_blk = get_tuned_block_sizes(
            q.dtype,
            kv_pages.dtype,
            num_q_heads_per_blk,
            num_combined_kv_heads_per_blk // 2,
            head_dim,
            page_size,
            num_q_tokens,
            pages_per_seq,
        )
    num_q_heads_per_kv_head = num_q_heads // num_kv_heads
    num_q_blks = cdiv(num_q_tokens, num_q_per_blk)
    assert num_combined_kv_heads_per_blk % 2 == 0
    num_kv_heads_per_blk = num_combined_kv_heads_per_blk // 2
    assert num_q_heads_per_blk % num_q_heads_per_kv_head == 0
    num_heads_blks = num_q_heads // num_q_heads_per_blk
    grid = (num_heads_blks, num_q_blks)

    def q_index_map(heads_blk_idx, q_blk_idx, *_):
        return (q_blk_idx, heads_blk_idx, 0)

    q_block_spec = pl.BlockSpec(
        (num_q_per_blk, num_q_heads_per_blk, head_dim),
        q_index_map,
    )
    in_specs = [
        q_block_spec,
        pl.BlockSpec(memory_space=pltpu.ANY),
    ]
    out_specs = q_block_spec
    lm_scratch = pltpu.VMEM(
        # TODO(jevinjiang): use 128 instead of 1 is due to Mosaic does not support
        # unaligned slicing!
        (num_kv_heads_per_blk, num_q_per_blk * num_q_heads_per_kv_head, 128),
        jnp.float32,
    )
    acc_scratch = pltpu.VMEM(
        (num_q_per_blk, num_q_heads_per_blk, head_dim),
        jnp.float32,
    )
    double_buf_scratch = pltpu.VMEM(
        (
            2,  # For double buffering during DMA copies.
            num_kv_pages_per_blk,
            page_size,
            num_combined_kv_heads_per_blk,
            head_dim,
        ),
        kv_pages.dtype,
    )
    scratch_shapes = [
        double_buf_scratch,  # kv_bufs
        pltpu.SemaphoreType.DMA((2, )),  # Semaphores for double buffers.
        lm_scratch,  # l_ref
        lm_scratch,  # m_ref
        acc_scratch,
    ]
    scalar_prefetches = (
        kv_lens,
        page_indices,
        cu_q_lens,
        jnp.array((0, 0), jnp.int32),  # seq_idx, buf_idx
        num_seqs,
    )
    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,
            k_scale=k_scale,
            v_scale=v_scale,
        ),
        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(
            dimension_semantics=(
                "arbitrary",
                "arbitrary",
            ),
            vmem_limit_bytes=vmem_limit_bytes,
        ),
        out_shape=jax.ShapeDtypeStruct(shape=q.shape, dtype=q.dtype),
        name="ragged_paged_attention_kernel",
    )

    return kernel(*scalar_prefetches, q, kv_pages)
