# Copyright 2023 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.

"""Module containing decode attention."""
from __future__ import annotations

import functools
from typing import Any

import jax
from jax import lax
from jax.experimental import pallas as pl
from jax.experimental.pallas import gpu as plgpu
import jax.numpy as jnp


def attn_forward_kernel(
    q_ref,  # [num_heads, head_dim]
    k_ref,  # [k_seq_len, head_dim]
    v_ref,  # [k_seq_len, head_dim]
    o_ref: Any,  # [num_heads, head_dim]
    *residual_refs: Any,  # Residual outputs: [num_heads,], [num_heads,]
    sm_scale: float,
    block_k: int,
):
  block_h, head_dim = q_ref.shape
  k_seq_len, _ = k_ref.shape
  start_q = pl.program_id(0)

  # o is the buffer where we accumulate the output on sram.
  # m_i and l_i (see FlashAttention2 paper) are updated during the k,v loop.
  m_i = jnp.zeros(block_h, dtype=jnp.float32) - float("inf")
  l_i = jnp.zeros(block_h, dtype=jnp.float32)
  o = jnp.zeros((block_h, head_dim), dtype=jnp.float32)

  # Load q: it will stay in L1 throughout. Indices form a matrix because we
  # read, compute, and write all in 2d chunks. 1 element ~= 1 CUDA thread index.
  # q tile has shape [block_h, head_dim].
  curr_q_slice = pl.dslice(start_q * block_h, block_h)
  q = pl.load(q_ref, (curr_q_slice, pl.dslice(None)))

  def _dot(a, b):
    # if a.shape[0] == 1:
    #   # Use matrix vector product
    #   return (a.T * b).sum(axis=0, keepdims=True)
    return pl.dot(a, b)

  # Loop over blocks of kv to process entire kv seq_len.
  # Grid loops over q blocks over num_heads.
  def body(start_k, carry):
    o_prev, m_prev, l_prev = carry
    curr_k_slice = pl.dslice(start_k * block_k, block_k)

    k = pl.load(k_ref, (curr_k_slice, slice(None)))
    qk = _dot(q, k.T)  # [block_h, block_k]
    if sm_scale != 1.0:
      qk *= sm_scale  # [block_h, block_k]

    m_curr = qk.max(axis=-1)
    m_next = jnp.maximum(m_prev, m_curr)
    correction = jnp.exp(m_prev - m_next)
    l_prev_corr = correction * l_prev
    s_curr = jnp.exp(
        qk - m_next[:, None]
    )  # Use m_next instead of m_curr to avoid a correction on l_curr
    l_curr = s_curr.sum(axis=-1)
    l_next = l_prev_corr + l_curr
    v = pl.load(v_ref, (curr_k_slice, slice(None)))
    o_curr = _dot(s_curr.astype(v.dtype), v)

    # flash2 unscaled_o
    o_next = correction[:, None] * o_prev + o_curr
    return o_next, m_next, l_next

  upper_bound = pl.cdiv(k_seq_len, block_k)
  # o is left unscaled; it will be scaled in the final reduction step
  o, m_i, l_i = lax.fori_loop(0, upper_bound, body, (o, m_i, l_i))

  if residual_refs:
    l_ref, m_ref = residual_refs
    pl.store(l_ref, (curr_q_slice,), l_i)
    pl.store(m_ref, (curr_q_slice,), m_i)
  # Write output to dram.
  o = o.astype(o_ref.dtype)
  pl.store(o_ref, (curr_q_slice, pl.dslice(None)), o)


def attn_unbatched(
    q,  # [num_heads, head_dim]
    k,  # [k_seq_len, head_dim]
    v,  # [k_seq_len, head_dim]
    sm_scale: float,
    block_h: int,
    block_k: int,
    k_splits: int,
    num_warps: int | None,
    num_stages: int,
    grid: tuple[int, ...] | None,
    interpret: bool,
    debug: bool,
):
  num_heads, head_dim = q.shape
  k_seq_len, _ = k.shape
  # Pad num query heads to 16 if needed, and slice output at the end.
  original_num_heads = None
  if num_heads < 16:
    q = jnp.pad(q, ((0, 16 - num_heads), (0, 0)))
    original_num_heads = num_heads
    num_heads = q.shape[0]
  block_h = min(block_h, num_heads)
  head_splits = pl.cdiv(num_heads, block_h)
  grid_ = grid
  if grid_ is None:
    grid_ = (head_splits, k_splits)

  assert (
      k_seq_len % k_splits == 0
  ), f"{k_seq_len=} must be divisible by {k_splits=}"
  k = k.reshape(k_splits, k_seq_len // k_splits, head_dim)
  v = v.reshape(k_splits, k_seq_len // k_splits, head_dim)
  k_seq_len = k_seq_len // k_splits
  assert min(num_heads, head_dim, k_seq_len) >= 16, "Minimum pl.dot size is 16"
  block_k = min(block_k, k_seq_len)
  num_warps_ = num_warps
  if num_warps_ is None:
    num_warps_ = 4
  kernel = functools.partial(
      attn_forward_kernel,
      sm_scale=sm_scale,
      block_k=block_k,
  )

  o, l, m = pl.pallas_call(
      kernel,
      grid=grid_,
      in_specs=[
          pl.BlockSpec((block_h, head_dim), lambda i, j: (i, 0)),
          pl.BlockSpec((None, k_seq_len, head_dim), lambda i, j: (j, 0, 0)),
          pl.BlockSpec((None, k_seq_len, head_dim), lambda i, j: (j, 0, 0)),
      ],
      out_specs=[
          pl.BlockSpec((None, block_h, head_dim), lambda i, j: (j, i, 0)),  # o
          pl.BlockSpec((None, block_h), lambda i, j: (j, i)),  # l
          pl.BlockSpec((None, block_h), lambda i, j: (j, i)),  # m
      ],
      compiler_params=plgpu.TritonCompilerParams(
          num_warps=num_warps_, num_stages=num_stages
      ),
      out_shape=[
          jax.ShapeDtypeStruct(shape=(k_splits, *q.shape), dtype=q.dtype),  # o
          jax.ShapeDtypeStruct(
              shape=(k_splits, num_heads), dtype=jnp.float32
          ),  # l
          jax.ShapeDtypeStruct(
              shape=(k_splits, num_heads), dtype=jnp.float32
          ),  # m
      ],
      debug=debug,
      interpret=interpret,
      name="mha_forward",
  )(q, k, v)

  # final round of flash
  m_next = m.max(axis=0)
  correction = jnp.exp(m - m_next[None])
  o = o * correction[:, :, None]
  l_next = (l * correction).sum(axis=0)
  o = o.sum(axis=0) / l_next[:, None]

  if original_num_heads is not None:
    o = o[:original_num_heads, :]
  return o


@functools.partial(
    jax.jit,
    static_argnames=[
        "sm_scale",
        "block_h",
        "block_k",
        "k_splits",
        "num_warps",
        "num_stages",
        "grid",
        "interpret",
        "debug",
    ],
)
def mqa(
    q,  # [batch_size, num_heads, head_dim]
    k,  # [batch_size, k_seq_len, head_dim]
    v,  # [batch_size, k_seq_len, head_dim]
    sm_scale: float = 1.0,
    block_h: int = 16,
    block_k: int = 256,
    k_splits: int = 16,
    num_warps: int | None = None,
    num_stages: int = 2,
    grid: tuple[int, ...] | None = None,
    interpret: bool = False,
    debug: bool = False,
):
  inner = functools.partial(
      attn_unbatched,
      sm_scale=sm_scale,
      block_h=block_h,
      block_k=block_k,
      k_splits=k_splits,
      num_warps=num_warps,
      num_stages=num_stages,
      grid=grid,
      interpret=interpret,
      debug=debug,
  )
  return jax.vmap(inner)(q, k, v)


@functools.partial(
    jax.jit,
    static_argnames=[
        "sm_scale",
        "block_h",
        "block_k",
        "k_splits",
        "num_warps",
        "num_stages",
        "grid",
        "interpret",
        "debug",
    ],
)
def gqa(
    q,  # [batch_size, num_q_heads, head_dim]
    k,  # [batch_size, k_seq_len, num_kv_heads, head_dim]
    v,  # [batch_size, k_seq_len, num_kv_heads, head_dim]
    sm_scale: float = 1.0,
    block_h: int = 16,
    block_k: int = 256,
    k_splits: int = 16,
    num_warps: int | None = None,
    num_stages: int = 2,
    grid: tuple[int, ...] | None = None,
    interpret: bool = False,
    debug: bool = False,
):
  batch_size, q_heads, head_dim = q.shape
  kv_heads = k.shape[2]
  assert kv_heads == v.shape[2]
  assert q_heads % kv_heads == 0
  q_heads_per_kv_head = q_heads // kv_heads
  q_reshaped = q.reshape(batch_size, kv_heads, q_heads_per_kv_head, head_dim)
  k_transposed = jnp.swapaxes(
      k, 1, 2
  )  # [batch_size, num_kv_heads, k_seq_len, head_dim]
  v_transposed = jnp.swapaxes(
      v, 1, 2
  )  # [batch_size, num_kv_heads, k_seq_len, head_dim]
  inner = functools.partial(
      attn_unbatched,
      sm_scale=sm_scale,
      block_h=block_h,
      block_k=block_k,
      k_splits=k_splits,
      num_warps=num_warps,
      num_stages=num_stages,
      grid=grid,
      interpret=interpret,
      debug=debug,
  )
  with_kv_heads = jax.vmap(inner)
  o = jax.vmap(with_kv_heads)(q_reshaped, k_transposed, v_transposed)
  return o.reshape(batch_size, q_heads, head_dim)


@functools.partial(jax.jit, static_argnames=["sm_scale"])
def mqa_reference(
    q,  # [bs, num_q_heads, head_dim]
    k,  # [bs, k_seq_len, head_dim]
    v,  # [bs, k_seq_len, head_dim]
    sm_scale=1.0,
):
  logits = jnp.einsum("bnd,bsd->bns", q, k).astype(jnp.float32)
  weights = jax.nn.softmax(logits * sm_scale).astype(q.dtype)
  return jnp.einsum("bns,bsd->bnd", weights, v)


@functools.partial(jax.jit, static_argnames=["sm_scale"])
def mha_reference(
    q,  # [bs, num_q_heads, head_dim]
    k,  # [bs, k_seq_len, num_k_heads, head_dim]
    v,  # [bs, k_seq_len, num_v_heads, head_dim]
    sm_scale=1.0,
):
  assert q.shape[1] == k.shape[2]
  logits = jnp.einsum("bnd,bsnd->bns", q, k).astype(jnp.float32)
  weights = jax.nn.softmax(logits * sm_scale).astype(q.dtype)
  return jnp.einsum("bns,bsnd->bnd", weights, v)


@functools.partial(jax.jit, static_argnames=["sm_scale"])
def gqa_reference(
    q,  # [bs, num_q_heads, head_dim]
    k,  # [bs, k_seq_len, num_k_heads, head_dim]
    v,  # [bs, k_seq_len, num_v_heads, head_dim]
    sm_scale=1.0,
):
  bs, num_q_heads, head_dim = q.shape
  num_kv_heads = k.shape[2]
  assert num_q_heads % num_kv_heads == 0
  q_reshaped = q.reshape(
      bs, num_kv_heads, num_q_heads // num_kv_heads, head_dim
  )
  k_transposed = jnp.swapaxes(
      k, 1, 2
  )  # [batch_size, num_kv_heads, k_seq_len, head_dim]
  v_transposed = jnp.swapaxes(
      v, 1, 2
  )  # [batch_size, num_kv_heads, k_seq_len, head_dim]
  logits = jnp.einsum("bkgd,bksd->bkgs", q_reshaped, k_transposed).astype(
      jnp.float32
  )
  weights = jax.nn.softmax(logits * sm_scale).astype(q.dtype)
  o = jnp.einsum("bkgs,bksd->bkgd", weights, v_transposed)
  return o.reshape(bs, num_q_heads, head_dim)
