code: def _flash_attention_dq_kernel(
      q_tile_ref,
      k_tile_ref,
      v_tile_ref,
      ab_tile_ref,
      q_segment_ids_tile_ref,
      kv_segment_ids_tile_ref,
      l_tile_ref,
      m_tile_ref,
      do_tile_ref,
      di_tile_ref,
      dq_tile_ref,
      dq_scratch_ref,
      ds_tile_ref,
      *,
      sm_scale: float,
      causal: bool,
      mask_value: float,
      kv_seq_len: int,
      block_k: int,
  ):
      _, _, block_k_major, _ = k_tile_ref.shape
      _, _, block_q_major, _ = q_tile_ref.shape
  
      kv_seq_index = pl.program_id(axis=3)
      q_seq_index = pl.program_id(axis=2)
  
      @pl.when(kv_seq_index == 0)
      def start_new_sequence():
          dq_scratch_ref[:, :] = jnp.zeros(
              dq_scratch_ref.shape, dq_scratch_ref.dtype
          )
  
      def body(i, _):
          k_slice = pl.ds(i * block_k, block_k)
          q = q_tile_ref[0, 0, :, :]
          k = pl.load(
              k_tile_ref,
              (0, 0, k_slice, slice(None)),
          )  # [block_k, head_dim]
          v = pl.load(
              v_tile_ref,
              (0, 0, k_slice, slice(None)),
          )  # [block_k, head_dim]
          l = l_tile_ref[0, 0, :, :]  # [block_q_major, 128]
          m = m_tile_ref[0, 0, :, :]  # [block_q_major, 128]
          do = do_tile_ref[0, 0, :, :]  # [block_q_major, head_dim]
          di = di_tile_ref[0, 0, :].astype(
              jnp.float32
          )  # [block_q_major, 128]
  
          capped_logits = jax.lax.dot_general(
              q,
              k,
              TRANS_B_DIM_NUMBERS,
              preferred_element_type=jnp.float32,
          )
  
          if ab_tile_ref is not None:
              ab = pl.load(
                  ab_tile_ref,
                  (
                      0,
                      0,
                      pl.dslice(None),
                      pl.dslice(i * block_k, block_k),
                  ),
              ).astype(jnp.float32)
              capped_logits += ab
  
          if sm_scale != 1.0:
              capped_logits *= sm_scale
  
          mask = None
          if q_segment_ids_tile_ref is not None:
              repeats, rem = divmod(block_k, NUM_LANES)
              if rem:
                  raise NotImplementedError(
                      f"kv block size must be a multiple of {NUM_LANES}"
                  )
              q_segment_ids = pltpu.repeat(
                  q_segment_ids_tile_ref[0], repeats, axis=1
              )  # [block_q, block_k].
              kv_segment_ids = pl.load(
                  kv_segment_ids_tile_ref, (slice(None), 0, k_slice)
              )  # [1, block_k].
              mask = jnp.equal(q_segment_ids, kv_segment_ids).astype(
                  jnp.bool_
              )
  
          if causal:
              mask_shape = (block_q_major, block_k)
              row_ids = jax.lax.broadcasted_iota(
                  jnp.int32, mask_shape, 0
              )
              row_ids += q_seq_index * block_q_major
              col_ids = jax.lax.broadcasted_iota(
                  jnp.int32, mask_shape, 1
              )
              col_ids += kv_seq_index * block_k_major + i * block_k
              causal_mask = col_ids <= row_ids
              mask = (
                  causal_mask
                  if mask is None
                  else jnp.logical_and(mask, causal_mask)
              )
          capped_logits = (
              capped_logits
              if mask is None
              else capped_logits + jnp.where(mask, 0.0, mask_value)
          )
  
          p = jnp.exp(
              capped_logits
              - pltpu.repeat(m, block_k // MIN_BLOCK_SIZE, axis=1)
          )
          p = p * pltpu.repeat(
              1 / l, block_k // MIN_BLOCK_SIZE, axis=1
          )  # [block_q_major, block_k]
  
          # di: [block_q_major, 128]
          # do: [block_q_major, head_dim]
          # v: [block_k_major, head_dim]
          dp = jax.lax.dot_general(
              do,
              v,
              TRANS_B_DIM_NUMBERS,
              preferred_element_type=jnp.float32,
          )
          ds = (
              dp - pltpu.repeat(di, block_k // MIN_BLOCK_SIZE, axis=1)
          ) * p
          # dp = jnp.dot(do, v.T)
          # ds = (dp - (dp * p).sum(axis=1)[:, None]) * p
  
          if sm_scale != 1.0:
              ds = ds * sm_scale
  
          if ds_tile_ref is not None:
              pl.store(
                  ds_tile_ref,
                  (
                      0,
                      0,
                      pl.dslice(None),
                      pl.dslice(i * block_k, block_k),
                  ),
                  ds.astype(ds_tile_ref.dtype),
              )
  
          # dp: [block_q_major, block_k]
          # k: [block_k, head_dim]
          dq_scratch_ref[:, :] += lax.dot(
              ds.astype(k.dtype),
              k,
              preferred_element_type=jnp.float32,
          ).astype(dq_scratch_ref.dtype)
  
      if causal:
          should_run = below_or_on_diag(
              q_seq_index, block_q_major, kv_seq_index, block_k_major
          )
          should_not_run = lax.select(should_run, False, True)
      else:
          should_run = True
          should_not_run = False  # type: ignore
  
      @pl.when(should_run)
      def run():
          lax.fori_loop(
              0, block_k_major // block_k, body, None, unroll=True
          )
  
      @pl.when(should_not_run)
      def zero_out_ds():
          if ds_tile_ref is not None:
              ds_tile_ref[...] = jnp.zeros_like(ds_tile_ref)
  
      @pl.when(kv_seq_index == kv_seq_len // block_k_major - 1)
      def end_of_kv_sequence():
          dq_tile_ref[0, 0, :, :] = dq_scratch_ref[...].astype(
              dq_tile_ref
          )
          dq_scratch_ref[...] = jnp.zeros_like(dq_scratch_ref)
  
imports: from __future__ import annotations
  import dataclasses
  import functools
  from typing import Any, NamedTuple
  import jax
  from jax import lax
  from jax.experimental import pallas
  from jax.experimental.pallas import tpu
  import jax.numpy
label_desc: Use image node_id 12 for calling a global function with example usage: _flash_attention_dq_kernel(q_tile_ref, k_tile_ref, v_tile_ref, ab_tile_ref, q_segment_ids_tile_ref, kv_segment_ids_tile_ref, l_tile_ref, m_tile_ref, do_tile_ref, di_tile_ref, dq_tile_ref, dq_scratch_ref, ds_tile_ref) without return types