doc: fix logits cap docstring (#300)

follow up of #299, pre-attention -> pre-softmax.

Author: yzh119

python/flashinfer/decode.py

@@ -86,7 +86,7 @@ def single_decode_with_kv_cache(
         ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
         Defaults to ``NONE``.
     logits_cap : bool
-        Whether to apply logits cap to pre-attention logits.
+        Whether to apply logits cap to pre-softmax logits.
         If ``True``, the logits will be capped according to formula (proposed in
         Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
         Defaults to ``False``.
@@ -199,7 +199,7 @@ def batch_decode_with_padded_kv_cache(
         ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
         Defaults to ``NONE``.
     logits_cap : bool
-        Whether to apply logits cap to pre-attention logits.
+        Whether to apply logits cap to pre-softmax logits.
         If ``True``, the logits will be capped according to formula (proposed in
         Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
         Defaults to ``False``.
@@ -312,7 +312,7 @@ def batch_decode_with_padded_kv_cache_return_lse(
         ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
         Defaults to ``NONE``.
     logits_cap : bool
-        Whether to apply logits cap to pre-attention logits.
+        Whether to apply logits cap to pre-softmax logits.
         If ``True``, the logits will be capped according to formula (proposed in
         Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
         Defaults to ``False``.
@@ -592,7 +592,7 @@ def begin_forward(
             ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
             Defaults to ``NONE``.
         logits_cap: bool
-            Whether to apply logits cap to pre-attention logits.
+            Whether to apply logits cap to pre-softmax logits.
             If ``True``, the logits will be capped according to formula (proposed in
             Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
             Defaults to ``False``.
@@ -704,7 +704,7 @@ def forward(
             ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
             Defaults to ``NONE``.
         logits_cap: bool
-            Whether to apply logits cap to pre-attention logits.
+            Whether to apply logits cap to pre-softmax logits.
             If ``True``, the logits will be capped according to formula (proposed in
             Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
             Defaults to ``False``.
@@ -789,7 +789,7 @@ def forward_return_lse(
             ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
             Defaults to ``NONE``.
         logits_cap: bool
-            Whether to apply logits cap to pre-attention logits.
+            Whether to apply logits cap to pre-softmax logits.
             If ``True``, the logits will be capped according to formula (proposed in
             Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
             Defaults to ``False``.

python/flashinfer/prefill.py

@@ -96,7 +96,7 @@ def single_prefill_with_kv_cache(
         ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
         Default is ``NONE``.
     logits_cap : bool
-        Whether to apply logits cap to pre-attention logits.
+        Whether to apply logits cap to pre-softmax logits.
         If ``True``, the logits will be capped according to formula (proposed in
         Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
         Defaults to ``False``.
@@ -240,7 +240,7 @@ def single_prefill_with_kv_cache_return_lse(
         ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
         Default is ``NONE``.
     logits_cap : bool
-        Whether to apply logits cap to pre-attention logits.
+        Whether to apply logits cap to pre-softmax logits.
         If ``True``, the logits will be capped according to formula (proposed in
         Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
         Defaults to ``False``.
@@ -770,7 +770,7 @@ def forward(
             ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
             Default is ``NONE``.
         logits_cap : bool
-            Whether to apply logits cap to pre-attention logits, 
+            Whether to apply logits cap to pre-softmax logits, 
             If ``True``, the logits will be capped according to formula (proposed in
             Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
             Defaults to ``False``.
@@ -874,7 +874,7 @@ def forward_return_lse(
             ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
             Default is ``NONE``.
         logits_cap : bool
-            Whether to apply logits cap to pre-attention logits.
+            Whether to apply logits cap to pre-softmax logits.
             If ``True``, the logits will be capped according to formula (proposed in
             Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
             Defaults to ``False``.
@@ -1276,7 +1276,7 @@ def forward(
             ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
             Default is ``NONE``.
         logits_cap : bool
-            Whether to apply logits cap to pre-attention logits.
+            Whether to apply logits cap to pre-softmax logits.
             If ``True``, the logits will be capped according to formula (proposed in
             Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
             Defaults to ``False``.
@@ -1378,7 +1378,7 @@ def forward_return_lse(
             ``NONE``/``ROPE_LLAMA`` (LLAMA style rotary embedding) /``ALIBI``.
             Default is ``NONE``.
         logits_cap : bool
-            Whether to apply logits cap to pre-attention logits.
+            Whether to apply logits cap to pre-softmax logits.
             If ``True``, the logits will be capped according to formula (proposed in
             Grok-1): :math:`30 \times \mathrm{tanh}(x / 30)`, where :math:`x` is the input logits.
             Defaults to ``False``.