From b36b5bf0e0a73d549776a1e0cc5fd41f29889bf8 Mon Sep 17 00:00:00 2001 From: Zihao Ye Date: Wed, 31 Jan 2024 22:17:45 +0800 Subject: [PATCH] upd --- docs/conf.py | 4 ++-- python/flashinfer/cascade.py | 2 +- python/flashinfer/decode.py | 8 ++++---- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index aff57f29..ea23f246 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -94,6 +94,6 @@ } # add additional overrides -# templates_path += [tlcpack_sphinx_addon.get_templates_path()] -# html_static_path += [tlcpack_sphinx_addon.get_static_path()] +templates_path += [tlcpack_sphinx_addon.get_templates_path()] +html_static_path += [tlcpack_sphinx_addon.get_static_path()] diff --git a/python/flashinfer/cascade.py b/python/flashinfer/cascade.py index 93ba7583..1d5b8890 100644 --- a/python/flashinfer/cascade.py +++ b/python/flashinfer/cascade.py @@ -289,7 +289,7 @@ def begin_forward( The ``num_qo_heads`` must be a multiple of ``num_kv_heads``. If ``num_qo_heads`` is not equal to ``num_kv_heads``, the function will use - `grouped query attention`_. + `grouped query attention `_. """ self._batch_decode_wrapper.begin_forward( unique_kv_indptr, diff --git a/python/flashinfer/decode.py b/python/flashinfer/decode.py index 215e3673..d4f42cc7 100644 --- a/python/flashinfer/decode.py +++ b/python/flashinfer/decode.py @@ -157,12 +157,12 @@ def batch_decode_with_padded_kv_cache( The padded key tensor, shape: ``[batch_size, padded_seq_len, num_kv_heads, head_dim]`` if :attr:`kv_layout` is ``NHD`` or ``[batch_size, num_kv_heads, padded_seq_len, head_dim]`` if - :attr:`kv_layout` is ``HND`. + :attr:`kv_layout` is ``HND``. v_padded : torch.Tensor The padded value tensor, shape: ``[batch_size, padded_seq_len, num_kv_heads, head_dim]`` if :attr:`kv_layout` is ``NHD`` or ``[batch_size, num_kv_heads, padded_seq_len, head_dim]`` if - :attr:`kv_layout` is ``HND`. + :attr:`kv_layout` is ``HND``. kv_layout : str The layout of the input k/v tensors, could be either ``NHD`` or ``HND``. rotary_mode : str @@ -228,12 +228,12 @@ def batch_decode_with_padded_kv_cache_return_lse( The padded key tensor, shape: ``[batch_size, padded_seq_len, num_kv_heads, head_dim]`` if :attr:`kv_layout` is ``NHD`` or ``[batch_size, num_kv_heads, padded_seq_len, head_dim]`` if - :attr:`kv_layout` is ``HND`. + :attr:`kv_layout` is ``HND``. v_padded : torch.Tensor The padded value tensor, shape: ``[batch_size, padded_seq_len, num_kv_heads, head_dim]`` if :attr:`kv_layout` is ``NHD`` or ``[batch_size, num_kv_heads, padded_seq_len, head_dim]`` if - :attr:`kv_layout` is ``HND`. + :attr:`kv_layout` is ``HND``. kv_layout : str The layout of the input k/v tensors, could be either ``NHD`` or ``HND``. rotary_mode : str