sglang 0.1.14__py3-none-any.whl → 0.1.21__py3-none-any.whl

sglang/srt/layers/logits_processor.py

@@ -1,11 +1,56 @@
+"""Logits processing."""
+
+import dataclasses
+from typing import List, Union
+
 import torch
-from sglang.srt.managers.router.model_runner import ForwardMode, InputMetadata
 from torch import nn
-from vllm.model_executor.parallel_utils.communication_op import (
+from vllm.distributed import (
     get_tensor_model_parallel_world_size,
     tensor_model_parallel_all_gather,
 )
 
+from sglang.srt.managers.controller.model_runner import ForwardMode, InputMetadata
+
+
+@dataclasses.dataclass
+class LogitProcessorOutput:
+    # The logits of the next tokens.       shape: [#seq, vocab_size]
+    next_token_logits: torch.Tensor
+    # The logprobs of the next tokens.     shape: [#seq, vocab_size]
+    next_token_logprobs: torch.Tensor
+
+    # The normlaized logprobs of prompts.  shape: [#seq]
+    normalized_prompt_logprobs: torch.Tensor
+    # The logprobs of prefill tokens.      shape: [#token, vocab_size]
+    prefill_token_logprobs: torch.Tensor
+
+    # The logprob and id of the top-k tokens in prefill positions.  shape [#seq, #token, k] of Tuple(logprob, token_id)
+    prefill_top_logprobs: List
+    # The logprob and id of the top-k tokens in decode positions.   shape [#seq, #token, k] of Tuple(logprob, token_id)
+    decode_top_logprobs: List
+
+
+@dataclasses.dataclass
+class LogitsMetadata:
+    forward_mode: ForwardMode
+    extend_seq_lens: torch.Tensor
+    extend_start_loc: torch.Tensor
+
+    # For logprobs
+    return_logprob: bool
+    top_logprobs_nums: List[int]
+
+    @classmethod
+    def from_input_metadata(cls, input_metadata: InputMetadata):
+        return cls(
+            forward_mode=input_metadata.forward_mode,
+            extend_seq_lens=input_metadata.extend_seq_lens,
+            extend_start_loc=input_metadata.extend_start_loc,
+            return_logprob=input_metadata.return_logprob,
+            top_logprobs_nums=input_metadata.top_logprobs_nums,
+        )
+
 
 class LogitsProcessor(nn.Module):
     def __init__(self, config):
@@ -13,78 +58,159 @@ class LogitsProcessor(nn.Module):
         self.config = config
         self.tp_size = get_tensor_model_parallel_world_size()
 
-    def forward(self, input_ids, hidden_states, weight, input_metadata):
-        last_index = None
+    def _get_normalized_prompt_logprobs(
+        self, prefill_token_logprobs, logits_metadata: LogitsMetadata
+    ):
+        logprobs_cumsum = torch.cumsum(
+            prefill_token_logprobs, dim=0, dtype=torch.float32
+        )
 
-        # Compute the last index (the first decode token) of each requeast
-        # if we are in prefill or extend mode.
-        if input_metadata.forward_mode != ForwardMode.DECODE:
-            last_index = (
-                torch.cumsum(
-                    input_metadata.seq_lens - input_metadata.prefix_lens,
-                    dim=0,
-                    dtype=torch.long,
+        start = logits_metadata.extend_start_loc.clone()
+        end = start + logits_metadata.extend_seq_lens - 2
+        start.clamp_(min=0, max=prefill_token_logprobs.shape[0] - 1)
+        end.clamp_(min=0, max=prefill_token_logprobs.shape[0] - 1)
+        sum_logp = (
+            logprobs_cumsum[end]
+            - logprobs_cumsum[start]
+            + prefill_token_logprobs[start]
+        )
+        normalized_prompt_logprobs = sum_logp / (
+            (logits_metadata.extend_seq_lens - 1).clamp(min=1)
+        )
+
+        return normalized_prompt_logprobs
+
+    def _get_top_logprobs(self, all_logprobs, logits_metadata: LogitsMetadata):
+        # TODO: vectorize the code below
+        if logits_metadata.forward_mode == ForwardMode.DECODE:
+            decode_top_logprobs = []
+            for i in range(all_logprobs.shape[0]):
+                k = logits_metadata.top_logprobs_nums[i]
+                t = all_logprobs[i].topk(k)
+                v_cpu = t.values.tolist()
+                p_cpu = t.indices.tolist()
+                decode_top_logprobs.append(list(zip(v_cpu, p_cpu)))
+            return None, decode_top_logprobs
+        else:
+            prefill_top_logprobs, decode_top_logprobs = [], []
+            pt = 0
+            extend_seq_lens_cpu = logits_metadata.extend_seq_lens.tolist()
+            for i, extend_seq_len in enumerate(extend_seq_lens_cpu):
+                if extend_seq_len == 0:
+                    prefill_top_logprobs.append([])
+                    decode_top_logprobs.append([])
+                    continue
+                k = logits_metadata.top_logprobs_nums[i]
+                t = all_logprobs[pt : pt + extend_seq_len].topk(k)
+                vs_cpu = t.values.tolist()
+                ps_cpu = t.indices.tolist()
+                prefill_top_logprobs.append(
+                    [list(zip(vs_cpu[j], ps_cpu[j])) for j in range(len(vs_cpu) - 1)]
                 )
+                decode_top_logprobs.append(list(zip(vs_cpu[-1], ps_cpu[-1])))
+                pt += extend_seq_len
+
+            return prefill_top_logprobs, decode_top_logprobs
+
+    def forward(
+        self,
+        input_ids,
+        hidden_states,
+        weight,
+        logits_metadata: Union[LogitsMetadata, InputMetadata],
+    ):
+        if isinstance(logits_metadata, InputMetadata):
+            logits_metadata = LogitsMetadata.from_input_metadata(logits_metadata)
+        assert isinstance(logits_metadata, LogitsMetadata)
+
+        # Get the last hidden states and last logits for the next token prediction
+        if logits_metadata.forward_mode == ForwardMode.DECODE:
+            last_index = None
+            last_hidden = hidden_states
+        else:
+            last_index = (
+                torch.cumsum(logits_metadata.extend_seq_lens, dim=0, dtype=torch.long)
                 - 1
             )
+            last_hidden = hidden_states[last_index]
 
-        if not input_metadata.return_logprob:
-            # When logprob is not requested, only compute the last logits.
-            if input_metadata.forward_mode == ForwardMode.DECODE:
-                last_hidden = hidden_states
-            else:
-                last_hidden = hidden_states[last_index]
-                hidden_states = None
-
-            last_logits = torch.matmul(last_hidden, weight.T)
-            if self.tp_size > 1:
-                last_logits = tensor_model_parallel_all_gather(last_logits)
-            last_logits = last_logits[:, : self.config.vocab_size]
-            return last_logits, (None, None, None)
+        last_logits = torch.matmul(last_hidden, weight.T)
+        if self.tp_size > 1:
+            last_logits = tensor_model_parallel_all_gather(last_logits)
+        last_logits = last_logits[:, : self.config.vocab_size]
+
+        if hasattr(self.config, "final_logit_softcapping"):
+            last_logits /= self.config.final_logit_softcapping
+            last_logits = torch.tanh(last_logits)
+            last_logits *= self.config.final_logit_softcapping
+
+        # Return only last_logits if logprob is not requested
+        if not logits_metadata.return_logprob:
+            return LogitProcessorOutput(
+                next_token_logits=last_logits,
+                next_token_logprobs=None,
+                normalized_prompt_logprobs=None,
+                prefill_token_logprobs=None,
+                prefill_top_logprobs=None,
+                decode_top_logprobs=None,
+            )
         else:
             # When logprob is requested, compute the logits for all tokens.
-            logits = torch.matmul(hidden_states, weight.T)
-            if self.tp_size > 1:
-                logits = tensor_model_parallel_all_gather(logits)
-            logits = logits[:, : self.config.vocab_size]
-            all_logprobs = torch.log(torch.softmax(logits.float(), dim=-1) + 1e-6)
-
-            if input_metadata.forward_mode == ForwardMode.DECODE:
-                last_logits = logits
-                last_logprobs = all_logprobs
-                prefill_logprobs = normalized_logprobs = None
+            if logits_metadata.forward_mode == ForwardMode.DECODE:
+                all_logits = last_logits
+            else:
+                all_logits = torch.matmul(hidden_states, weight.T)
+                if self.tp_size > 1:
+                    all_logits = tensor_model_parallel_all_gather(all_logits)
+                all_logits = all_logits[:, : self.config.vocab_size]
+
+            all_logprobs = all_logits.float()
+            del all_logits
+            all_logprobs[:] = torch.nn.functional.log_softmax(all_logprobs, dim=-1)
+
+            # Get the logprob of top-k tokens
+            return_top_logprob = any(x > 0 for x in logits_metadata.top_logprobs_nums)
+            if return_top_logprob:
+                prefill_top_logprobs, decode_top_logprobs = self._get_top_logprobs(
+                    all_logprobs, logits_metadata
+                )
+            else:
+                prefill_top_logprobs = decode_top_logprobs = None
+
+            if logits_metadata.forward_mode == ForwardMode.DECODE:
+                return LogitProcessorOutput(
+                    next_token_logits=last_logits,
+                    next_token_logprobs=all_logprobs,
+                    normalized_prompt_logprobs=None,
+                    prefill_token_logprobs=None,
+                    prefill_top_logprobs=None,
+                    decode_top_logprobs=decode_top_logprobs,
+                )
             else:
-                # Compute the logprobs for the last token of each request.
-                last_logits = logits[last_index]
                 last_logprobs = all_logprobs[last_index]
 
                 # Compute the logprobs and normalized logprobs for the prefill tokens.
                 # Note that we pad a zero at the end of each sequence for easy computation.
-                prefill_logprobs = all_logprobs[
+                prefill_token_logprobs = all_logprobs[
                     torch.arange(all_logprobs.shape[0], device="cuda"),
                     torch.cat([input_ids[1:], torch.tensor([0], device="cuda")]),
                 ]
-                logprobs_cumsum = torch.cumsum(
-                    prefill_logprobs, dim=0, dtype=torch.float32
-                )
 
-                start = input_metadata.extend_start_loc.clone()
-                end = start + input_metadata.extend_seq_lens - 2
-                start.clamp_(min=0, max=prefill_logprobs.shape[0] - 1)
-                end.clamp_(min=0, max=prefill_logprobs.shape[0] - 1)
-                sum_logp = (
-                    logprobs_cumsum[end]
-                    - logprobs_cumsum[start]
-                    + prefill_logprobs[start]
-                )
-                normalized_logprobs = sum_logp / (
-                    (input_metadata.extend_seq_lens - 1).clamp(min=1)
+                normalized_prompt_logprobs = self._get_normalized_prompt_logprobs(
+                    prefill_token_logprobs, logits_metadata
                 )
 
-            return last_logits, (prefill_logprobs, normalized_logprobs, last_logprobs)
+                return LogitProcessorOutput(
+                    next_token_logits=last_logits,
+                    next_token_logprobs=last_logprobs,
+                    normalized_prompt_logprobs=normalized_prompt_logprobs,
+                    prefill_token_logprobs=prefill_token_logprobs,
+                    prefill_top_logprobs=prefill_top_logprobs,
+                    decode_top_logprobs=decode_top_logprobs,
+                )
 
 
-if __name__ == "__main__":
+def test():
     all_logprobs = torch.tensor(
         #       s                     s                s
         [[0, 1, 2, 3], [1, 2, 3, 4], [2, 3, 4, 5], [3, 4, 5, 6], [4, 5, 6, 7]],
@@ -93,23 +219,26 @@ if __name__ == "__main__":
     )
     seq_lens = torch.tensor([2, 0, 3, 0], dtype=torch.int32, device="cuda")
     input_ids = torch.tensor([1, 2, 3, 0, 1], dtype=torch.int32, device="cuda")
-    logprobs = torch.zeros(5, dtype=torch.float32, device="cuda")
 
-    logprobs = all_logprobs[
+    token_logprobs = all_logprobs[
         torch.arange(all_logprobs.shape[0], device="cuda"),
         torch.cat([input_ids[1:], torch.tensor([0], device="cuda")]),
     ]
-    logprobs_cumsum = torch.cumsum(logprobs, dim=0, dtype=torch.float32)
+    logprobs_cumsum = torch.cumsum(token_logprobs, dim=0, dtype=torch.float32)
 
     len_cumsum = torch.cumsum(seq_lens, dim=0)
     start = torch.cat((torch.tensor([0], device="cuda"), len_cumsum[:-1]), 0)
     end = start + seq_lens - 2
-    start.clamp_(min=0, max=logprobs.shape[0] - 1)
-    end.clamp_(min=0, max=logprobs.shape[0] - 1)
-    sum_logp = logprobs_cumsum[end] - logprobs_cumsum[start] + logprobs[start]
+    start.clamp_(min=0, max=token_logprobs.shape[0] - 1)
+    end.clamp_(min=0, max=token_logprobs.shape[0] - 1)
+    sum_logp = logprobs_cumsum[end] - logprobs_cumsum[start] + token_logprobs[start]
 
     # assert logprobs == [2, _, 2, 4, _]
-    print("logprobs", logprobs)
+    print("token logprobs", token_logprobs)
     print("start", start)
     print("end", end)
     print("sum_logp", sum_logp)
+
+
+if __name__ == "__main__":
+    test()

sglang/srt/layers/radix_attention.py

@@ -1,46 +1,42 @@
+"""Radix attention."""
+
 import torch
-from sglang.srt.layers.context_flashattention_nopad import context_attention_fwd
+from flashinfer.cascade import merge_state
+from torch import nn
+
+from sglang.global_config import global_config
 from sglang.srt.layers.extend_attention import extend_attention_fwd
 from sglang.srt.layers.token_attention import token_attention_fwd
-from sglang.srt.managers.router.model_runner import ForwardMode, InputMetadata
-from torch import nn
+from sglang.srt.managers.controller.infer_batch import global_server_args_dict
+from sglang.srt.managers.controller.model_runner import ForwardMode, InputMetadata
 
 
 class RadixAttention(nn.Module):
-    def __init__(self, num_heads, head_dim, scaling, num_kv_heads, layer_id):
+    def __init__(
+        self,
+        num_heads: int,
+        head_dim: int,
+        scaling: float,
+        num_kv_heads: int,
+        layer_id: int,
+        logit_cap: int = -1,
+    ):
         super().__init__()
         self.tp_q_head_num = num_heads
         self.tp_k_head_num = num_kv_heads
         self.tp_v_head_num = num_kv_heads
         self.head_dim = head_dim
+        self.scaling = scaling
         self.layer_id = layer_id
 
-        from sglang.srt.managers.router.model_runner import global_server_args_dict
-
-        if global_server_args_dict.get("enable_flashinfer", False):
-            self.prefill_forward = self.prefill_forward_flashinfer
-            self.extend_forward = self.prefill_forward_flashinfer
+        if not global_server_args_dict.get("disable_flashinfer", False):
+            self.extend_forward = self.extend_forward_flashinfer
             self.decode_forward = self.decode_forward_flashinfer
         else:
-            self.prefill_forward = self.prefill_forward_triton
             self.extend_forward = self.extend_forward_triton
             self.decode_forward = self.decode_forward_triton
 
-    def prefill_forward_triton(self, q, k, v, input_metadata: InputMetadata):
-        o = torch.empty_like(q)
-
-        context_attention_fwd(
-            q.view(-1, self.tp_q_head_num, self.head_dim),
-            k,
-            v,
-            o.view(-1, self.tp_q_head_num, self.head_dim),
-            input_metadata.start_loc,
-            input_metadata.seq_lens,
-            input_metadata.max_seq_len,
-        )
-        self.store_kv_cache(k, v, input_metadata)
-
-        return o
+        self.logit_cap = logit_cap if logit_cap is not None and logit_cap > 0 else 0
 
     def extend_forward_triton(self, q, k, v, input_metadata: InputMetadata):
         o = torch.empty_like(q)
@@ -54,13 +50,15 @@ class RadixAttention(nn.Module):
             input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id),
             input_metadata.req_to_token_pool.req_to_token,
             input_metadata.req_pool_indices,
-            input_metadata.start_loc,
+            input_metadata.triton_start_loc,
             input_metadata.seq_lens,
-            input_metadata.prefix_lens,
+            input_metadata.triton_prefix_lens,
             input_metadata.extend_start_loc,
             input_metadata.extend_seq_lens,
-            input_metadata.max_seq_len,
-            input_metadata.max_extend_len,
+            input_metadata.triton_max_seq_len,
+            input_metadata.triton_max_extend_len,
+            sm_scale=self.scaling,
+            logit_cap=self.logit_cap,
         )
 
         return o
@@ -76,31 +74,54 @@ class RadixAttention(nn.Module):
             o.view(-1, self.tp_q_head_num, self.head_dim),
             input_metadata.req_to_token_pool.req_to_token,
             input_metadata.req_pool_indices,
-            input_metadata.start_loc,
+            input_metadata.triton_start_loc,
             input_metadata.seq_lens,
-            input_metadata.max_seq_len,
-            input_metadata.other_kv_index,
+            input_metadata.triton_max_seq_len,
             input_metadata.total_num_tokens,
+            sm_scale=self.scaling,
+            logit_cap=self.logit_cap,
         )
 
         return o
 
-    def prefill_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
-        self.store_kv_cache(k, v, input_metadata)
-
-        o = input_metadata.prefill_wrapper.forward(
+    def extend_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
+        o1, s1 = input_metadata.flashinfer_prefill_wrapper_ragged.forward_return_lse(
             q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
-            input_metadata.token_to_kv_pool.kv_data[self.layer_id],
+            k.contiguous().view(-1, self.tp_k_head_num, self.head_dim),
+            v.contiguous().view(-1, self.tp_v_head_num, self.head_dim),
+            causal=True,
+            sm_scale=self.scaling,
+            logits_soft_cap=self.logit_cap,
         )
 
+        if input_metadata.extend_no_prefix:
+            o = o1
+        else:
+            o2, s2 = input_metadata.flashinfer_prefill_wrapper_paged.forward_return_lse(
+                q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
+                input_metadata.token_to_kv_pool.kv_data[self.layer_id],
+                causal=False,
+                sm_scale=self.scaling,
+                logits_soft_cap=self.logit_cap,
+            )
+
+            o, _ = merge_state(o1, s1, o2, s2)
+
+        self.store_kv_cache(k, v, input_metadata)
+
+        if input_metadata.total_num_tokens >= global_config.layer_sync_threshold:
+            torch.cuda.synchronize()
+
         return o.view(-1, self.tp_q_head_num * self.head_dim)
 
     def decode_forward_flashinfer(self, q, k, v, input_metadata: InputMetadata):
         self.store_kv_cache(k, v, input_metadata)
 
-        o = input_metadata.decode_wrapper.forward(
+        o = input_metadata.flashinfer_decode_wrapper.forward(
             q.contiguous().view(-1, self.tp_q_head_num, self.head_dim),
             input_metadata.token_to_kv_pool.kv_data[self.layer_id],
+            sm_scale=self.scaling,
+            logits_soft_cap=self.logit_cap,
         )
 
         return o.view(-1, self.tp_q_head_num * self.head_dim)
@@ -109,25 +130,13 @@ class RadixAttention(nn.Module):
         k = k.view(-1, self.tp_k_head_num, self.head_dim)
         v = v.view(-1, self.tp_v_head_num, self.head_dim)
 
-        if input_metadata.forward_mode == ForwardMode.PREFILL:
-            return self.prefill_forward(q, k, v, input_metadata)
-        elif input_metadata.forward_mode == ForwardMode.EXTEND:
+        if input_metadata.forward_mode == ForwardMode.EXTEND:
             return self.extend_forward(q, k, v, input_metadata)
         elif input_metadata.forward_mode == ForwardMode.DECODE:
             return self.decode_forward(q, k, v, input_metadata)
 
     def store_kv_cache(self, cache_k, cache_v, input_metadata: InputMetadata):
         key_buffer = input_metadata.token_to_kv_pool.get_key_buffer(self.layer_id)
+        key_buffer[input_metadata.out_cache_loc] = cache_k
         value_buffer = input_metadata.token_to_kv_pool.get_value_buffer(self.layer_id)
-        if input_metadata.out_cache_loc is not None:
-            key_buffer[input_metadata.out_cache_loc] = cache_k
-            value_buffer[input_metadata.out_cache_loc] = cache_v
-        elif input_metadata.out_cache_cont_start is not None:
-            key_buffer[
-                input_metadata.out_cache_cont_start : input_metadata.out_cache_cont_end
-            ] = cache_k
-            value_buffer[
-                input_metadata.out_cache_cont_start : input_metadata.out_cache_cont_end
-            ] = cache_v
-        else:
-            raise RuntimeError()
+        value_buffer[input_metadata.out_cache_loc] = cache_v

sglang/srt/layers/token_attention.py

@@ -4,7 +4,8 @@
 import torch
 import triton
 import triton.language as tl
-from sglang.srt.managers.router.model_runner import global_server_args_dict
+
+from sglang.srt.managers.controller.model_runner import global_server_args_dict
 from sglang.srt.utils import wrap_kernel_launcher
 
 if global_server_args_dict.get("attention_reduce_in_fp32", False):
@@ -15,6 +16,12 @@ else:
     REDUCE_TORCH_TYPE = torch.float16
 
 
+@triton.jit
+def tanh(x):
+    # Tanh is just a scaled sigmoid
+    return 2 * tl.sigmoid(2 * x) - 1
+
+
 @triton.jit
 def _fwd_kernel_stage1(
     Q,
@@ -34,6 +41,7 @@ def _fwd_kernel_stage1(
     kv_group_num: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
+    logit_cap: tl.constexpr,
 ):
     cur_batch = tl.program_id(0)
     cur_head = tl.program_id(1)
@@ -76,6 +84,10 @@ def _fwd_kernel_stage1(
         ).to(REDUCE_TRITON_TYPE)
         att_value = tl.sum(q[None, :] * k, 1)
         att_value *= sm_scale
+
+        if logit_cap > 0:
+            att_value = logit_cap * tanh(att_value / logit_cap)
+
         off_o = cur_head * att_stride_h + (cur_batch_in_all_start_index + offs_n)
         tl.store(Att_Out + off_o, att_value, mask=offs_n_new < cur_batch_end_index)
 
@@ -95,7 +107,6 @@ def _fwd_kernel_stage2(
     stride_obs,
     stride_oh,
     stride_req_to_token_b,
-    other_kv_index,  # To fix a NAN issue
     kv_group_num: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
@@ -126,7 +137,7 @@ def _fwd_kernel_stage2(
             + cur_batch_req_idx * stride_req_to_token_b
             + (start_n + offs_n),
             mask=(start_n + offs_n) < cur_batch_seq_len,
-            other=other_kv_index,
+            other=0,
         )
 
         qk = tl.load(
@@ -164,13 +175,14 @@ def _token_att_m_fwd(
     B_Start_Loc,
     B_Seqlen,
     max_len_in_batch,
+    sm_scale,
+    logit_cap,
 ):
     BLOCK = 32
     # shape constraints
     Lq, Lk = q.shape[-1], k_buffer.shape[-1]
     assert Lq == Lk
     assert Lk in {16, 32, 64, 128, 256}
-    sm_scale = 1.0 / (Lk**0.5)
 
     batch, head_num = B_req_idx.shape[0], q.shape[1]
 
@@ -222,6 +234,7 @@ def _token_att_m_fwd(
         kv_group_num=kv_group_num,
         BLOCK_DMODEL=Lk,
         BLOCK_N=BLOCK,
+        logit_cap=logit_cap,
         num_warps=num_warps,
         num_stages=1,
     )
@@ -236,7 +249,6 @@ def _token_softmax_reducev_fwd(
     b_req_idx,
     b_start_loc,
     b_seq_len,
-    other_kv_index,
 ):
     BLOCK = 64
     batch, head = b_seq_len.shape[0], logics.shape[0]
@@ -263,7 +275,6 @@ def _token_softmax_reducev_fwd(
             o.stride(0),
             o.stride(1),
             req_to_tokens.stride(0),
-            other_kv_index,
         )
         return
 
@@ -281,7 +292,6 @@ def _token_softmax_reducev_fwd(
         o.stride(0),
         o.stride(1),
         req_to_tokens.stride(0),
-        other_kv_index,
         kv_group_num=kv_group_num,
         BLOCK_DMODEL=v_buffer.shape[-1],
         BLOCK_N=BLOCK,
@@ -301,8 +311,9 @@ def token_attention_fwd(
     b_start_loc,
     b_seq_len,
     max_len_in_batch,
-    other_kv_index,
     total_num_tokens,
+    sm_scale,
+    logit_cap=-1,
     att_m=None,
 ):
     if att_m is None:
@@ -319,6 +330,8 @@ def token_attention_fwd(
         b_start_loc,
         b_seq_len,
         max_len_in_batch,
+        sm_scale,
+        logit_cap,
     )
     _token_softmax_reducev_fwd(
         att_m,
@@ -328,5 +341,4 @@ def token_attention_fwd(
         b_req_idx,
         b_start_loc,
         b_seq_len,
-        other_kv_index,
     )