Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix PagedPrefill python api and some typos #441

Merged
merged 3 commits into from
Aug 13, 2024
Merged

Fix PagedPrefill python api and some typos #441

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
298b176
Fix: typos confused between "NHD" and "HND"
jianfei-wangg Aug 13, 2024
4152d10
Fix: paged_prefill error when cuda_graph is disabled
jianfei-wangg Aug 13, 2024
eaf2c66
Mod: remove not-needed modification
jianfei-wangg Aug 13, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions python/csrc/batch_prefill.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -289,11 +289,11 @@ std::vector<torch::Tensor> BatchPrefillWithPagedKVCachePyTorchWrapper::ForwardCu

if (paged_kv_defined) {
// [max_num_pages, 2, num_kv_heads, page_size, head_dim] for HND
// [max_num_pages, 2, page_size, num_kv_heads, head_dim] for HND
// [max_num_pages, 2, page_size, num_kv_heads, head_dim] for NHD
CHECK_DIM(5, paged_kv_cache.value());
} else {
// [max_num_pages, num_kv_heads, page_size, head_dim] for HND
// [max_num_pages, page_size, num_kv_heads, head_dim] for HND
// [max_num_pages, page_size, num_kv_heads, head_dim] for NHD
CHECK_DIM(4, paged_k_cache.value());
CHECK_DIM(4, paged_v_cache.value());
}
Expand Down
4 changes: 2 additions & 2 deletions python/flashinfer/cascade.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -374,7 +374,7 @@ def forward(
``[max_num_pages, 2, page_size, num_kv_heads, head_dim]`` if
:attr:`kv_layout` is ``NHD``, and
``[max_num_pages, 2, num_kv_heads, page_size, head_dim]`` if
:attr:`kv_layout` is ``NHD``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
:attr:`kv_layout` is ``HND``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
``paged_kv_cache[:, 1]`` is the value-cache.

allow_fp16_qk_reduction : bool
Expand Down Expand Up @@ -631,7 +631,7 @@ def forward(
``[max_num_pages, 2, page_size, num_kv_heads, head_dim]`` if
:attr:`kv_layout` is ``NHD``, and
``[max_num_pages, 2, num_kv_heads, page_size, head_dim]`` if
:attr:`kv_layout` is ``NHD``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
:attr:`kv_layout` is ``HND``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
``paged_kv_cache[:, 1]`` is the value-cache.

causal : bool
Expand Down
4 changes: 2 additions & 2 deletions python/flashinfer/decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -577,7 +577,7 @@ def forward(
``[max_num_pages, 2, page_size, num_kv_heads, head_dim]`` if
:attr:`kv_layout` is ``NHD``, and
``[max_num_pages, 2, num_kv_heads, page_size, head_dim]`` if
:attr:`kv_layout` is ``NHD``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
:attr:`kv_layout` is ``HND``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
``paged_kv_cache[:, 1]`` is the value-cache.

pos_encoding_mode : str
Expand Down Expand Up @@ -696,7 +696,7 @@ def forward_return_lse(
``[max_num_pages, 2, page_size, num_kv_heads, head_dim]`` if
:attr:`kv_layout` is ``NHD``, and
``[max_num_pages, 2, num_kv_heads, page_size, head_dim]`` if
:attr:`kv_layout` is ``NHD``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
:attr:`kv_layout` is ``HND``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
``paged_kv_cache[:, 1]`` is the value-cache.

pos_encoding_mode : str
Expand Down
2 changes: 1 addition & 1 deletion python/flashinfer/page.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ def append_paged_kv_cache(
``[max_num_pages, 2, page_size, num_kv_heads, head_dim]`` if
:attr:`kv_layout` is ``NHD``, and
``[max_num_pages, 2, num_kv_heads, page_size, head_dim]`` if
:attr:`kv_layout` is ``NHD``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
:attr:`kv_layout` is ``HND``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
``paged_kv_cache[:, 1]`` is the value-cache.

kv_indices : torch.Tensor
Expand Down
8 changes: 4 additions & 4 deletions python/flashinfer/prefill.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -778,8 +778,8 @@ def begin_forward(
self._paged_kv_indices_buf = paged_kv_indices
self._paged_kv_last_page_len_buf = paged_kv_last_page_len
if packed_custom_mask is not None:
self._custom_mask = packed_custom_mask
self._qk_indptr = qk_indptr
self._custom_mask_buf = packed_custom_mask
self._qk_indptr_buf = qk_indptr
empty_q_data = torch.empty(
0,
dtype=(
Expand Down Expand Up @@ -843,7 +843,7 @@ def forward(
``[max_num_pages, 2, page_size, num_kv_heads, head_dim]`` if
:attr:`kv_layout` is ``NHD``, and
``[max_num_pages, 2, num_kv_heads, page_size, head_dim]`` if
:attr:`kv_layout` is ``NHD``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
:attr:`kv_layout` is ``HND``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
``paged_kv_cache[:, 1]`` is the value-cache.

causal : bool
Expand Down Expand Up @@ -969,7 +969,7 @@ def forward_return_lse(
``[max_num_pages, 2, page_size, num_kv_heads, head_dim]`` if
:attr:`kv_layout` is ``NHD``, and
``[max_num_pages, 2, num_kv_heads, page_size, head_dim]`` if
:attr:`kv_layout` is ``NHD``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
:attr:`kv_layout` is ``HND``. Where ``paged_kv_cache[:, 0]`` is the key-cache and
``paged_kv_cache[:, 1]`` is the value-cache.

causal : bool
Expand Down