ejkernel.kernels._pallas.tpu.ragged_page_attention_v2 package

• ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._interface
  • ragged_page_attention_v2()
• ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._pallas_impl_fwd
  • MultiPageAsyncCopyDescriptor
  • cdiv()
  • dynamic_validate_inputs()
  • get_dtype_packing()
  • get_min_heads_per_blk()
  • ragged_page_attention_kernel()
  • ref_ragged_page_attention()
  • static_validate_inputs()
• ejkernel.kernels._pallas.tpu.ragged_page_attention_v2._utils
  • get_device_name()
  • get_min_page_size()
  • get_tpu_version()
  • get_tuned_block_sizes()
  • next_power_of_2()
  • simplify_key()