Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
gaoqiong
composable_kernel_ROCM
Commits
86517ce4
"test/srt/vscode:/vscode.git/clone" did not exist on "24eaebeb4b43ca24c8bf9aaf8c9d0836487f07df"
Commit
86517ce4
authored
Dec 02, 2024
by
Po Yen Chen
Browse files
Use vllm paged-kcache layout to read blocks
parent
44828b7c
Changes
1
Show whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
32 additions
and
6 deletions
+32
-6
include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_kernel.hpp
include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_kernel.hpp
+32
-6
No files found.
include/ck_tile/ops/fmha/kernel/fmha_fwd_splitkv_kernel.hpp
View file @
86517ce4
...
...
@@ -635,12 +635,38 @@ struct FmhaFwdSplitKVKernel
}();
const
auto
make_k_dram
=
[
&
](
const
KDataType
*
data
,
index_t
height
)
{
const
auto
k_dram_naive
=
make_naive_tensor_view
<
address_space_enum
::
global
>
(
auto
k_dram_naive
=
[
&
]
{
if
constexpr
(
kIsPagedKV
)
{
constexpr
index_t
vector_size
=
FmhaPipeline
::
kAlignmentK
;
// (hdim_q/vector_size, seqlen_k, vector_size)
const
auto
view
=
make_naive_tensor_view
<
address_space_enum
::
global
>
(
data
,
// will update this pointer if using paged-kvcache
make_tuple
(
kargs
.
hdim_q
/
vector_size
,
height
,
number
<
vector_size
>
{}),
make_tuple
(
height
*
vector_size
,
number
<
vector_size
>
{},
number
<
1
>
{}),
number
<
vector_size
>
{},
number
<
1
>
{});
// (seqlen_k, hdim_q)
return
transform_tensor_view
(
view
,
make_tuple
(
make_pass_through_transform
(
height
),
make_merge_transform
(
make_tuple
(
kargs
.
hdim_q
/
vector_size
,
number
<
vector_size
>
{}))),
make_tuple
(
sequence
<
1
>
{},
sequence
<
0
,
2
>
{}),
make_tuple
(
sequence
<
0
>
{},
sequence
<
1
>
{}));
}
else
{
// (seqlen_k, hdim_q)
return
make_naive_tensor_view
<
address_space_enum
::
global
>
(
data
,
// will update this pointer if using paged-kvcache
make_tuple
(
height
,
kargs
.
hdim_q
),
make_tuple
(
kargs
.
stride_k
,
1
),
number
<
FmhaPipeline
::
kAlignmentK
>
{},
number
<
1
>
{});
}
}();
return
pad_tensor_view
(
k_dram_naive
,
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment