upd

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()]
 

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<https://arxiv.org/abs/2305.13245>`_.
+        `grouped query attention <https://arxiv.org/abs/2305.13245>`_.
         """
         self._batch_decode_wrapper.begin_forward(
             unique_kv_indptr,

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