sglang 0.3.0__py3-none-any.whl → 0.3.1.post1__py3-none-any.whl

sglang/srt/layers/{decode_attention.py → triton_attention/decode_attention.py}

@@ -15,24 +15,15 @@ limitations under the License.
 
 """
 Memory-efficient attention for decoding.
+It supports page size = 1.
 """
 
 # Adapted from
 # https://github.com/ModelTC/lightllm/blob/f2a54f0912293f683bf1d1695fd12c4098a5bf82/lightllm/models/llama/triton_kernel/token_attention_nopad_att1.py
 # https://github.com/ModelTC/lightllm/blob/f2a54f0912293f683bf1d1695fd12c4098a5bf82/lightllm/models/llama/triton_kernel/token_attention_softmax_and_reducev.py
-import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.managers.schedule_batch import global_server_args_dict
-
-if global_server_args_dict.get("triton_attention_reduce_in_fp32", False):
-    REDUCE_TRITON_TYPE = tl.float32
-    REDUCE_TORCH_TYPE = torch.float32
-else:
-    REDUCE_TRITON_TYPE = tl.float16
-    REDUCE_TORCH_TYPE = torch.float16
-
 
 @triton.jit
 def tanh(x):
@@ -60,11 +51,13 @@ def _fwd_kernel_stage1(
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
     logit_cap: tl.constexpr,
+    Lk: tl.constexpr,
 ):
     cur_batch = tl.program_id(0)
     cur_head = tl.program_id(1)
     start_n = tl.program_id(2)
 
+    reduce_dtype = Att_Out.dtype.element_ty
     cur_kv_head = cur_head // kv_group_num
 
     offs_d = tl.arange(0, BLOCK_DMODEL)
@@ -83,7 +76,7 @@ def _fwd_kernel_stage1(
     block_mask = tl.where(block_stard_index < cur_batch_seq_len, 1, 0)
 
     for start_mark in range(0, block_mask, 1):
-        q = tl.load(Q + off_q + start_mark).to(REDUCE_TRITON_TYPE)
+        q = tl.load(Q + off_q + start_mark).to(reduce_dtype)
         offs_n_new = cur_batch_start_index + offs_n
         k_loc = tl.load(
             Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_n_new,
@@ -97,9 +90,9 @@ def _fwd_kernel_stage1(
         )
         k = tl.load(
             K_Buffer + offs_buf_k,
-            mask=offs_n_new[:, None] < cur_batch_end_index,
+            mask=(offs_n_new[:, None] < cur_batch_end_index) & (offs_d[None, :] < Lk),
             other=0.0,
-        ).to(REDUCE_TRITON_TYPE)
+        ).to(reduce_dtype)
         att_value = tl.sum(q[None, :] * k, 1)
         att_value *= sm_scale
 
@@ -112,7 +105,7 @@ def _fwd_kernel_stage1(
 
 @triton.jit
 def _fwd_kernel_stage2(
-    Logics,
+    logits,
     V_Buffer,
     Out,
     Req_to_tokens,
@@ -128,6 +121,7 @@ def _fwd_kernel_stage2(
     kv_group_num: tl.constexpr,
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
+    Lv: tl.constexpr,
 ):
     cur_batch = tl.program_id(0)
     cur_head = tl.program_id(1)
@@ -159,7 +153,7 @@ def _fwd_kernel_stage2(
         )
 
         qk = tl.load(
-            Logics
+            logits
             + cur_head * stride_logic_h
             + (cur_batch_start_loc + start_n + offs_n),
             mask=start_n + offs_n < cur_batch_seq_len,
@@ -170,14 +164,16 @@ def _fwd_kernel_stage2(
         old_scale = tl.exp(e_max - n_e_max)
         p = tl.exp(qk - n_e_max)
         e_sum = e_sum * old_scale + tl.sum(p, 0)
-        v = tl.load(v_ptrs + v_index[:, None] * stride_buf_vbs)
+        v = tl.load(
+            v_ptrs + v_index[:, None] * stride_buf_vbs, mask=(offs_d[None, :] < Lv)
+        )
         acc = acc * old_scale + tl.sum(p[:, None] * v, 0)
         e_max = n_e_max
 
     acc = acc / e_sum
     off_o = cur_batch * stride_obs + cur_head * stride_oh + offs_d
     out_ptrs = Out + off_o
-    tl.store(out_ptrs, acc)
+    tl.store(out_ptrs, acc, mask=(offs_d < Lv))
 
 
 def _decode_att_m_fwd(
@@ -193,10 +189,7 @@ def _decode_att_m_fwd(
     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}
+    Lk = k_buffer.shape[-1]
 
     batch, head_num = B_req_idx.shape[0], q.shape[1]
 
@@ -208,6 +201,8 @@ def _decode_att_m_fwd(
     else:
         num_warps = 2
 
+    BLOCK_DMODEL = triton.next_power_of_2(Lk)
+
     _fwd_kernel_stage1[grid](
         q,
         k_buffer,
@@ -224,16 +219,17 @@ def _decode_att_m_fwd(
         k_buffer.stride(1),
         att_out.stride(0),
         kv_group_num=kv_group_num,
-        BLOCK_DMODEL=Lk,
+        BLOCK_DMODEL=BLOCK_DMODEL,
         BLOCK_N=BLOCK,
         logit_cap=logit_cap,
         num_warps=num_warps,
         num_stages=1,
+        Lk=Lk,
     )
 
 
 def _decode_softmax_reducev_fwd(
-    logics,
+    logits,
     v_buffer,
     o,
     req_to_tokens,
@@ -242,31 +238,35 @@ def _decode_softmax_reducev_fwd(
     b_seq_len,
 ):
     BLOCK = 64
-    batch, head = b_seq_len.shape[0], logics.shape[0]
+    batch, head = b_seq_len.shape[0], logits.shape[0]
     grid = (batch, head, 1)
-    kv_group_num = logics.shape[0] // v_buffer.shape[1]
+    kv_group_num = logits.shape[0] // v_buffer.shape[1]
 
     num_warps = 1
 
+    Lv = v_buffer.shape[-1]
+    BLOCK_DMODEL = triton.next_power_of_2(Lv)
+
     _fwd_kernel_stage2[grid](
-        logics,
+        logits,
         v_buffer,
         o,
         req_to_tokens,
         b_req_idx,
         b_start_loc,
         b_seq_len,
-        logics.stride(0),
+        logits.stride(0),
         v_buffer.stride(0),
         v_buffer.stride(1),
         o.stride(0),
         o.stride(1),
         req_to_tokens.stride(0),
         kv_group_num=kv_group_num,
-        BLOCK_DMODEL=v_buffer.shape[-1],
+        BLOCK_DMODEL=BLOCK_DMODEL,
         BLOCK_N=BLOCK,
         num_warps=num_warps,
         num_stages=3,
+        Lv=Lv,
     )
 
 
@@ -293,11 +293,13 @@ def _fwd_grouped_kernel_stage1(
     BLOCK_N: tl.constexpr,
     BLOCK_H: tl.constexpr,
     logit_cap: tl.constexpr,
+    Lk: tl.constexpr,
 ):
     cur_batch = tl.program_id(0)
     cur_kv_head = tl.program_id(1)
     start_n = tl.program_id(2)
 
+    reduce_dtype = Att_Out.dtype.element_ty
     cur_head = cur_kv_head * kv_group_num + tl.arange(0, BLOCK_H)
     mask_h = cur_head < (cur_kv_head + 1) * kv_group_num
     mask_h = mask_h & (cur_head < q_head_num)
@@ -324,9 +326,9 @@ def _fwd_grouped_kernel_stage1(
     block_mask = tl.where(block_stard_index < cur_batch_seq_len, 1, 0)
 
     for start_mark in range(0, block_mask, 1):
-        q = tl.load(Q + offs_q + start_mark, mask=mask_h[:, None]).to(
-            REDUCE_TRITON_TYPE
-        )
+        q = tl.load(
+            Q + offs_q + start_mark, mask=(mask_h[:, None]) & (offs_d[None, :] < Lk)
+        ).to(reduce_dtype)
         offs_n_new = cur_batch_start_index + offs_n
         k_loc = tl.load(
             Req_to_tokens + stride_req_to_tokens_b * cur_batch_req_idx + offs_n_new,
@@ -340,13 +342,13 @@ def _fwd_grouped_kernel_stage1(
         )
         k = tl.load(
             K_Buffer + offs_buf_k,
-            mask=offs_n_new[None, :] < cur_batch_end_index,
+            mask=(offs_n_new[None, :] < cur_batch_end_index) & (offs_d[:, None] < Lk),
             other=0.0,
-        ).to(REDUCE_TRITON_TYPE)
+        ).to(reduce_dtype)
         qk = tl.dot(q, k)
         if BLOCK_DPE > 0:
             qpe = tl.load(Q + off_qpe + start_mark, mask=mask_h[:, None]).to(
-                REDUCE_TRITON_TYPE
+                reduce_dtype
             )
             offs_buf_kpe = (
                 k_loc[None, :] * stride_buf_kbs
@@ -357,7 +359,7 @@ def _fwd_grouped_kernel_stage1(
                 K_Buffer + offs_buf_kpe,
                 mask=offs_n_new[None, :] < cur_batch_end_index,
                 other=0.0,
-            ).to(REDUCE_TRITON_TYPE)
+            ).to(reduce_dtype)
             qk += tl.dot(qpe, kpe)
         qk *= sm_scale
 
@@ -377,7 +379,7 @@ def _fwd_grouped_kernel_stage1(
 
 @triton.jit
 def _fwd_grouped_kernel_stage2(
-    Logics,
+    logits,
     V_Buffer,
     Out,
     Req_to_tokens,
@@ -395,6 +397,7 @@ def _fwd_grouped_kernel_stage2(
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_N: tl.constexpr,
     BLOCK_H: tl.constexpr,
+    Lv: tl.constexpr,
 ):
     cur_batch = tl.program_id(0)
     cur_kv_head = tl.program_id(1)
@@ -432,7 +435,7 @@ def _fwd_grouped_kernel_stage2(
         )
 
         qk = tl.load(
-            Logics + offs_qk,
+            logits + offs_qk,
             mask=mask_h[:, None] & (start_n + offs_n[None, :] < cur_batch_seq_len),
             other=float("-inf"),
         )
@@ -441,7 +444,9 @@ def _fwd_grouped_kernel_stage2(
         old_scale = tl.exp(e_max - n_e_max)
         p = tl.exp(qk - n_e_max[:, None])
         e_sum = e_sum * old_scale + tl.sum(p, 1)
-        v = tl.load(v_ptrs + v_index[:, None] * stride_buf_vbs)
+        v = tl.load(
+            v_ptrs + v_index[:, None] * stride_buf_vbs, mask=(offs_d[None, :] < Lv)
+        )
         p = p.to(v.dtype)
         acc = acc * old_scale[:, None] + tl.dot(p, v)
         e_max = n_e_max
@@ -449,7 +454,7 @@ def _fwd_grouped_kernel_stage2(
     acc = acc / e_sum[:, None]
     off_o = cur_batch * stride_obs + cur_head[:, None] * stride_oh + offs_d[None, :]
     out_ptrs = Out + off_o
-    tl.store(out_ptrs, acc, mask=mask_h[:, None])
+    tl.store(out_ptrs, acc, mask=(mask_h[:, None]) & (offs_d[None, :] < Lv))
 
 
 def _decode_grouped_att_m_fwd(
@@ -464,17 +469,17 @@ def _decode_grouped_att_m_fwd(
     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, 576}
+    BLOCK = 64
+    Lk = k_buffer.shape[-1]
 
     if Lk == 576:
         BLOCK_DMODEL = 512
         BLOCK_DPE = 64
+    elif Lk == 288:
+        BLOCK_DMODEL = 256
+        BLOCK_DPE = 32
     else:
-        BLOCK_DMODEL = Lk
+        BLOCK_DMODEL = triton.next_power_of_2(Lk)
         BLOCK_DPE = 0
 
     batch, head_num = B_req_idx.shape[0], q.shape[1]
@@ -513,11 +518,12 @@ def _decode_grouped_att_m_fwd(
         logit_cap=logit_cap,
         num_warps=num_warps,
         num_stages=1,
+        Lk=Lk,
     )
 
 
 def _decode_grouped_softmax_reducev_fwd(
-    logics,
+    logits,
     v_buffer,
     o,
     req_to_tokens,
@@ -526,22 +532,25 @@ def _decode_grouped_softmax_reducev_fwd(
     b_seq_len,
 ):
     BLOCK = 128
-    batch, head_num = b_seq_len.shape[0], logics.shape[0]
-    kv_group_num = logics.shape[0] // v_buffer.shape[1]
+    batch, head_num = b_seq_len.shape[0], logits.shape[0]
+    kv_group_num = logits.shape[0] // v_buffer.shape[1]
     BLOCK_H = max(16, triton.next_power_of_2(kv_group_num))
     grid = (batch, triton.cdiv(head_num, min(BLOCK_H, kv_group_num)), 1)
 
     num_warps = 8
 
+    Lv = v_buffer.shape[-1]
+    BLOCK_DMODEL = triton.next_power_of_2(Lv)
+
     _fwd_grouped_kernel_stage2[grid](
-        logics,
+        logits,
         v_buffer,
         o,
         req_to_tokens,
         b_req_idx,
         b_start_loc,
         b_seq_len,
-        logics.stride(0),
+        logits.stride(0),
         v_buffer.stride(0),
         v_buffer.stride(1),
         o.stride(0),
@@ -549,9 +558,10 @@ def _decode_grouped_softmax_reducev_fwd(
         req_to_tokens.stride(0),
         kv_group_num=kv_group_num,
         q_head_num=head_num,
-        BLOCK_DMODEL=v_buffer.shape[-1],
+        BLOCK_DMODEL=BLOCK_DMODEL,
         BLOCK_N=BLOCK,
         BLOCK_H=BLOCK_H,
+        Lv=Lv,
         num_warps=num_warps,
         num_stages=1,
     )
@@ -566,17 +576,11 @@ def decode_attention_fwd(
     b_req_idx,
     b_start_loc,
     b_seq_len,
+    attn_logits,
     max_len_in_batch,
-    total_num_tokens,
     sm_scale,
-    logit_cap=-1,
-    att_m=None,
+    logit_cap=0.0,
 ):
-    if att_m is None:
-        att_m = torch.empty(
-            (q.shape[-2], total_num_tokens), dtype=REDUCE_TORCH_TYPE, device="cuda"
-        )
-
     kv_group_num = q.shape[1] // v_buffer.shape[1]
 
     if kv_group_num == 1:
@@ -584,7 +588,7 @@ def decode_attention_fwd(
         _decode_att_m_fwd(
             q,
             k_buffer,
-            att_m,
+            attn_logits,
             req_to_token,
             b_req_idx,
             b_start_loc,
@@ -594,7 +598,7 @@ def decode_attention_fwd(
             logit_cap,
         )
         _decode_softmax_reducev_fwd(
-            att_m,
+            attn_logits,
             v_buffer,
             o,
             req_to_token,
@@ -607,7 +611,7 @@ def decode_attention_fwd(
         _decode_grouped_att_m_fwd(
             q,
             k_buffer,
-            att_m,
+            attn_logits,
             req_to_token,
             b_req_idx,
             b_start_loc,
@@ -617,7 +621,7 @@ def decode_attention_fwd(
             logit_cap,
         )
         _decode_grouped_softmax_reducev_fwd(
-            att_m,
+            attn_logits,
             v_buffer,
             o,
             req_to_token,