code: def _flash_attention_kernel(q_tile_ref, *args, **kwargs):
      block_b = q_tile_ref.shape[0]
      # If we're not going to tile the softmax, then we can avoid a bunch of VPU ops.
      if kwargs["block_k"] == kwargs["kv_seq_len"]:
          kernel = _flash_attention_kernel_single_batch_single_step
      else:
          kernel = _flash_attention_kernel_single_batch
      for batch_idx in range(block_b):
          kernel((batch_idx, 0), q_tile_ref, *args, **kwargs)
  
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 6 for calling a global function with example usage: _flash_attention_kernel(q_tile_ref) without return types