Home
last modified time | relevance | path

Searched defs:kHeadDim (Results 1 – 5 of 5) sorted by relevance

/aosp_15_r20/external/pytorch/aten/src/ATen/native/transformers/cuda/flash_attn/
H A Dflash_bwd_preprocess_kernel.h71 constexpr int kHeadDim = Kernel_traits::kHeadDim; in compute_dot_do_o() local
154 constexpr int kHeadDim = Kernel_traits::kHeadDim; in clear_dKVaccum() local
198 constexpr int kHeadDim = Kernel_traits::kHeadDim; in convert_dQ() local
286 constexpr int kHeadDim = Kernel_traits::kHeadDim; in convert_dKV() local
H A Dflash_fwd_kernel.h43 constexpr int kHeadDim = Kernel_traits::kHeadDim; in compute_attn_1rowblock() local
493 constexpr int kHeadDim = Kernel_traits::kHeadDim; in compute_attn_1rowblock_splitkv() local
1092 constexpr int kHeadDim = Kernel_traits::kHeadDim; in combine_attn_seqk_parallel() local
H A Dkernel_traits.h70 static constexpr int kHeadDim = kHeadDim_; member
186 static constexpr int kHeadDim = kHeadDim_; member
H A Dflash_bwd_kernel.h94 constexpr int kHeadDim = Kernel_traits::kHeadDim; in compute_dq_dk_dv_1colblock() local
/aosp_15_r20/external/executorch/examples/mediatek/executor_runner/llama_runner/llm_helper/include/
H A Drotary_embedding.h74 const size_t kHeadDim; variable