ejkernel.kernels._pallas.tpu.ragged_page_attention_v3 package

• ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._interface
  • ragged_page_attention_v3()
• ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd
  • dynamic_validate_inputs()
  • get_kv_cache_shape()
  • get_smem_estimate_bytes()
  • get_vmem_estimate_bytes()
  • merge_kv()
  • prepare_inputs()
  • prepare_outputs()
  • ragged_paged_attention()
  • ref_ragged_paged_attention()
  • static_validate_inputs()
• ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._pallas_impl_fwd_h64
  • dynamic_validate_inputs()
  • get_kv_cache_shape()
  • get_smem_estimate_bytes()
  • get_vmem_estimate_bytes()
  • merge_kv()
  • prepare_inputs()
  • prepare_outputs()
  • ragged_paged_attention()
  • ref_ragged_paged_attention_hd64()
  • static_validate_inputs()
• ejkernel.kernels._pallas.tpu.ragged_page_attention_v3._utils
  • align_to()
  • cdiv()
  • get_device_name()
  • get_dtype_bitwidth()
  • get_dtype_packing()
  • get_lookup_keys()
  • get_lookup_keys_h64()
  • get_simplified_raw_key()
  • get_simplified_raw_key_h64()
  • get_tpu_version()
  • get_tuned_block_sizes()
  • get_tuned_block_sizes_h64()
  • next_power_of_2()