Searched defs:qkv (Results 1 – 5 of 5) sorted by relevance
| /external/pytorch/test/ |
| D | test_native_mha.py | 55 def simple_transform_bias_rescale_qkv(qkv, bias): argument
160 def __init__(self, embed_dim, num_heads, qkv, proj): argument
|
| /external/pytorch/aten/src/ATen/native/transformers/ |
| D | attention.cpp | 192 Tensor qkv; in qkv_projection() local
231 const Tensor& qkv, in transform_bias_rescale_qkv_cpu()
337 auto qkv = qkv_projection(query, key, value, embed_dim, qkv_weight); in native_multi_head_attention_cpu() local
913 auto qkv = qkv_projection(query, key, value, embed_dim, qkv_weight); in triton_multi_head_attention() local
|
| /external/pytorch/aten/src/ATen/native/cpu/ |
| D | NativeMultiheadAttnKernel.cpp | 94 const scalar_t* qkv = static_cast<const scalar_t*>(_qkv); in transform_bias_rescale_qkv_kernel_impl() local
|
| /external/pytorch/aten/src/ATen/native/transformers/cuda/ |
| D | attention.cu | 100 const PackedTensorAccessor64<scalar_t, 3, RestrictPtrTraits> qkv, in transform_bias_rescale_qkv_kernel()
201 const PackedTensorAccessor64<scalar_t, 1, RestrictPtrTraits> qkv, in transform_bias_rescale_qkv_add_padding_kernel()
377 const Tensor& qkv, in transform_bias_rescale_qkv_cuda()
596 auto qkv = qkv_projection(query, key, value, embed_dim, qkv_weight); in native_multi_head_attention_cuda() local
|
| /external/pytorch/aten/src/ATen/native/nested/cuda/ |
| D | NestedTensorTransformerUtils.cpp | 40 std::tuple<Tensor, int64_t, int64_t> cumulative_and_max_seq_len_nnz(const Tensor& qkv) { in cumulative_and_max_seq_len_nnz()
|