sglang 0.4.4.post1__py3-none-any.whl → 0.4.4.post3__py3-none-any.whl

sglang/srt/layers/attention/triton_ops/decode_attention.py

@@ -37,6 +37,9 @@ logger.warning(
 )
 
 
+_MIN_BLOCK_KV = 32
+
+
 @triton.jit
 def tanh(x):
     # Tanh is just a scaled sigmoid
@@ -52,6 +55,8 @@ def _fwd_kernel_stage1(
     kv_indptr,
     kv_indices,
     Att_Out,
+    Att_Lse,
+    num_kv_splits,
     stride_qbs,
     stride_qh,
     stride_buf_kbs,
@@ -65,7 +70,7 @@ def _fwd_kernel_stage1(
     BLOCK_DMODEL: tl.constexpr,
     BLOCK_DV: tl.constexpr,
     BLOCK_N: tl.constexpr,
-    NUM_KV_SPLITS: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
     logit_cap: tl.constexpr,
     Lk: tl.constexpr,
     Lv: tl.constexpr,
@@ -83,11 +88,13 @@ def _fwd_kernel_stage1(
 
     cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
     cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+    kv_splits = tl.load(num_kv_splits + cur_batch)
 
     off_q = cur_batch * stride_qbs + cur_head * stride_qh + offs_d
-    q = tl.load(Q + off_q, mask=mask_d, other=0.0)
 
-    kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
     split_kv_start = kv_len_per_split * split_kv_id
     split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
 
@@ -96,6 +103,7 @@ def _fwd_kernel_stage1(
     acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
 
     if split_kv_end > split_kv_start:
+        q = tl.load(Q + off_q, mask=mask_d, other=0.0)
         for start_n in range(split_kv_start, split_kv_end, BLOCK_N):
             offs_n = start_n + tl.arange(0, BLOCK_N)
             kv_loc = tl.load(
@@ -158,11 +166,10 @@ def _fwd_kernel_stage1(
             cur_batch * stride_mid_ob
             + cur_head * stride_mid_oh
             + split_kv_id * stride_mid_os
-            + Lv
-        )
+        ) // Lv
 
         tl.store(
-            Att_Out + offs_mid_o_1,
+            Att_Lse + offs_mid_o_1,
             e_max + tl.log(e_sum),
         )
 
@@ -172,9 +179,11 @@ def _decode_att_m_fwd(
     k_buffer,
     v_buffer,
     att_out,
+    att_lse,
     kv_indptr,
     kv_indices,
     num_kv_splits,
+    max_kv_splits,
     sm_scale,
     logit_cap,
 ):
@@ -182,13 +191,13 @@ def _decode_att_m_fwd(
     # [TODO] work around SGPR limit on MI3xx
     if _is_hip:
         BLOCK = 8
-    NUM_KV_SPLITS = num_kv_splits
+    MAX_KV_SPLITS = max_kv_splits
     Lk = k_buffer.shape[-1]
     Lv = v_buffer.shape[-1]
 
     batch, head_num = kv_indptr.shape[0] - 1, q.shape[1]
 
-    grid = (batch, head_num, NUM_KV_SPLITS)
+    grid = (batch, head_num, MAX_KV_SPLITS)
     kv_group_num = q.shape[1] // k_buffer.shape[1]
 
     if kv_group_num == 1:
@@ -209,6 +218,8 @@ def _decode_att_m_fwd(
         kv_indptr,
         kv_indices,
         att_out,
+        att_lse,
+        num_kv_splits,
         q.stride(0),
         q.stride(1),
         k_buffer.stride(0),
@@ -222,7 +233,7 @@ def _decode_att_m_fwd(
         BLOCK_DMODEL=BLOCK_DMODEL,
         BLOCK_DV=BLOCK_DV,
         BLOCK_N=BLOCK,
-        NUM_KV_SPLITS=NUM_KV_SPLITS,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
         logit_cap=logit_cap,
         num_warps=num_warps,
         num_stages=2,
@@ -240,6 +251,8 @@ def _fwd_grouped_kernel_stage1(
     kv_indptr,
     kv_indices,
     Att_Out,
+    Att_Lse,
+    num_kv_splits,
     stride_qbs,
     stride_qh,
     stride_buf_kbs,
@@ -256,7 +269,7 @@ def _fwd_grouped_kernel_stage1(
     BLOCK_DV: tl.constexpr,
     BLOCK_N: tl.constexpr,
     BLOCK_H: tl.constexpr,
-    NUM_KV_SPLITS: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
     logit_cap: tl.constexpr,
     Lk: tl.constexpr,
     Lv: tl.constexpr,
@@ -281,9 +294,9 @@ def _fwd_grouped_kernel_stage1(
 
     cur_batch_kv_start_idx = tl.load(kv_indptr + cur_batch)
     cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - cur_batch_kv_start_idx
+    kv_splits = tl.load(num_kv_splits + cur_batch)
 
     offs_q = cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_d[None, :]
-    q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_d[None, :]), other=0.0)
 
     if BLOCK_DPE > 0:
         offs_dpe = BLOCK_DMODEL + tl.arange(0, BLOCK_DPE)
@@ -291,11 +304,10 @@ def _fwd_grouped_kernel_stage1(
         off_qpe = (
             cur_batch * stride_qbs + cur_head[:, None] * stride_qh + offs_dpe[None, :]
         )
-        qpe = tl.load(
-            Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0
-        )
 
-    kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
     split_kv_start = kv_len_per_split * split_kv_id
     split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
 
@@ -304,6 +316,11 @@ def _fwd_grouped_kernel_stage1(
     acc = tl.zeros([BLOCK_H, BLOCK_DV], dtype=tl.float32)
 
     if split_kv_end > split_kv_start:
+        q = tl.load(Q + offs_q, mask=(mask_h[:, None]) & (mask_d[None, :]), other=0.0)
+        if BLOCK_DPE > 0:
+            qpe = tl.load(
+                Q + off_qpe, mask=(mask_h[:, None]) & (mask_dpe[None, :]), other=0.0
+            )
         for start_n in range(split_kv_start, split_kv_end, BLOCK_N):
             offs_n = start_n + tl.arange(0, BLOCK_N)
             kv_loc = tl.load(
@@ -380,11 +397,10 @@ def _fwd_grouped_kernel_stage1(
             cur_batch * stride_mid_ob
             + cur_head * stride_mid_oh
             + split_kv_id * stride_mid_os
-            + Lv
-        )
+        ) // Lv
 
         tl.store(
-            Att_Out + offs_mid_o_1,
+            Att_Lse + offs_mid_o_1,
             e_max + tl.log(e_sum),
             mask=mask_h,
         )
@@ -395,9 +411,11 @@ def _decode_grouped_att_m_fwd(
     k_buffer,
     v_buffer,
     att_out,
+    att_lse,
     kv_indptr,
     kv_indices,
     num_kv_splits,
+    max_kv_splits,
     sm_scale,
     logit_cap,
 ):
@@ -424,11 +442,11 @@ def _decode_grouped_att_m_fwd(
     kv_group_num = q.shape[1] // k_buffer.shape[1]
 
     BLOCK_H = 16
-    NUM_KV_SPLITS = num_kv_splits
+    MAX_KV_SPLITS = max_kv_splits
     grid = (
         batch,
         triton.cdiv(head_num, min(BLOCK_H, kv_group_num)),
-        NUM_KV_SPLITS,
+        MAX_KV_SPLITS,
     )
 
     extra_kargs = {}
@@ -447,6 +465,8 @@ def _decode_grouped_att_m_fwd(
         kv_indptr,
         kv_indices,
         att_out,
+        att_lse,
+        num_kv_splits,
         q.stride(0),
         q.stride(1),
         k_buffer.stride(0),
@@ -463,7 +483,7 @@ def _decode_grouped_att_m_fwd(
         BLOCK_DV=BLOCK_DV,
         BLOCK_N=BLOCK,
         BLOCK_H=BLOCK_H,
-        NUM_KV_SPLITS=NUM_KV_SPLITS,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
         logit_cap=logit_cap,
         num_warps=4,
         num_stages=num_stages,
@@ -476,14 +496,17 @@ def _decode_grouped_att_m_fwd(
 @triton.jit
 def _fwd_kernel_stage2(
     Mid_O,
+    Mid_O_1,
     O,
     kv_indptr,
+    num_kv_splits,
     stride_mid_ob,
     stride_mid_oh,
     stride_mid_os,
     stride_obs,
     stride_oh,
-    NUM_KV_SPLITS: tl.constexpr,
+    MAX_KV_SPLITS: tl.constexpr,
+    MIN_BLOCK_KV: tl.constexpr,
     BLOCK_DV: tl.constexpr,
     Lv: tl.constexpr,
 ):
@@ -493,6 +516,7 @@ def _fwd_kernel_stage2(
     cur_batch_seq_len = tl.load(kv_indptr + cur_batch + 1) - tl.load(
         kv_indptr + cur_batch
     )
+    kv_splits = tl.load(num_kv_splits + cur_batch)
 
     offs_d = tl.arange(0, BLOCK_DV)
     mask_d = offs_d < Lv
@@ -502,10 +526,12 @@ def _fwd_kernel_stage2(
     acc = tl.zeros([BLOCK_DV], dtype=tl.float32)
 
     offs_v = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + offs_d
-    offs_logic = cur_batch * stride_mid_ob + cur_head * stride_mid_oh + Lv
+    offs_logic = (cur_batch * stride_mid_ob + cur_head * stride_mid_oh) // Lv
+    kv_len_per_split = (
+        tl.cdiv(tl.cdiv(cur_batch_seq_len, kv_splits), MIN_BLOCK_KV) * MIN_BLOCK_KV
+    )
 
-    for split_kv_id in range(0, NUM_KV_SPLITS):
-        kv_len_per_split = tl.cdiv(cur_batch_seq_len, NUM_KV_SPLITS)
+    for split_kv_id in range(0, MAX_KV_SPLITS):
         split_kv_start = kv_len_per_split * split_kv_id
         split_kv_end = tl.minimum(split_kv_start + kv_len_per_split, cur_batch_seq_len)
 
@@ -513,7 +539,7 @@ def _fwd_kernel_stage2(
             tv = tl.load(
                 Mid_O + offs_v + split_kv_id * stride_mid_os, mask=mask_d, other=0.0
             )
-            tlogic = tl.load(Mid_O + offs_logic + split_kv_id * stride_mid_os)
+            tlogic = tl.load(Mid_O_1 + offs_logic + split_kv_id * stride_mid_os // Lv)
             n_e_max = tl.maximum(tlogic, e_max)
 
             old_scale = tl.exp(e_max - n_e_max)
@@ -533,17 +559,19 @@ def _fwd_kernel_stage2(
 
 def _decode_softmax_reducev_fwd(
     logits,
+    lse,
     q,
     o,
     v_buffer,
     kv_indptr,
     num_kv_splits,
+    max_kv_splits,
 ):
     batch, head_num = q.shape[0], q.shape[1]
     Lv = v_buffer.shape[-1]
     BLOCK_DV = triton.next_power_of_2(Lv)
 
-    NUM_KV_SPLITS = num_kv_splits
+    MAX_KV_SPLITS = max_kv_splits
 
     extra_kargs = {}
     if _is_hip:
@@ -554,14 +582,17 @@ def _decode_softmax_reducev_fwd(
     grid = (batch, head_num)
     _fwd_kernel_stage2[grid](
         logits,
+        lse,
         o,
         kv_indptr,
+        num_kv_splits,
         logits.stride(0),
         logits.stride(1),
         logits.stride(2),
         o.stride(0),
         o.stride(1),
-        NUM_KV_SPLITS=NUM_KV_SPLITS,
+        MAX_KV_SPLITS=MAX_KV_SPLITS,
+        MIN_BLOCK_KV=_MIN_BLOCK_KV,
         BLOCK_DV=BLOCK_DV,
         Lv=Lv,
         num_warps=4,
@@ -578,7 +609,9 @@ def decode_attention_fwd_normal(
     kv_indptr,
     kv_indices,
     attn_logits,
+    attn_lse,
     num_kv_splits,
+    max_kv_splits,
     sm_scale,
     logit_cap=0.0,
 ):
@@ -587,13 +620,24 @@ def decode_attention_fwd_normal(
         k_buffer,
         v_buffer,
         attn_logits,
+        attn_lse,
         kv_indptr,
         kv_indices,
         num_kv_splits,
+        max_kv_splits,
         sm_scale,
         logit_cap,
     )
-    _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, kv_indptr, num_kv_splits)
+    _decode_softmax_reducev_fwd(
+        attn_logits,
+        attn_lse,
+        q,
+        o,
+        v_buffer,
+        kv_indptr,
+        num_kv_splits,
+        max_kv_splits,
+    )
 
 
 def decode_attention_fwd_grouped(
@@ -604,7 +648,9 @@ def decode_attention_fwd_grouped(
     kv_indptr,
     kv_indices,
     attn_logits,
+    attn_lse,
     num_kv_splits,
+    max_kv_splits,
     sm_scale,
     logit_cap=0.0,
 ):
@@ -613,13 +659,24 @@ def decode_attention_fwd_grouped(
         k_buffer,
         v_buffer,
         attn_logits,
+        attn_lse,
         kv_indptr,
         kv_indices,
         num_kv_splits,
+        max_kv_splits,
         sm_scale,
         logit_cap,
     )
-    _decode_softmax_reducev_fwd(attn_logits, q, o, v_buffer, kv_indptr, num_kv_splits)
+    _decode_softmax_reducev_fwd(
+        attn_logits,
+        attn_lse,
+        q,
+        o,
+        v_buffer,
+        kv_indptr,
+        num_kv_splits,
+        max_kv_splits,
+    )
 
 
 def decode_attention_fwd(
@@ -630,11 +687,13 @@ def decode_attention_fwd(
     kv_indptr,
     kv_indices,
     attn_logits,
+    attn_lse,
     num_kv_splits,
+    max_kv_splits,
     sm_scale,
     logit_cap=0.0,
 ):
-    assert num_kv_splits == attn_logits.shape[2]
+    assert max_kv_splits == attn_logits.shape[2]
     assert q.shape[0] <= kv_indptr.shape[0] - 1
     assert q.shape[0] <= attn_logits.shape[0]
 
@@ -650,7 +709,9 @@ def decode_attention_fwd(
             kv_indptr,
             kv_indices,
             attn_logits,
+            attn_lse,
             num_kv_splits,
+            max_kv_splits,
             sm_scale,
             logit_cap,
         )
@@ -664,7 +725,9 @@ def decode_attention_fwd(
             kv_indptr,
             kv_indices,
             attn_logits,
+            attn_lse,
             num_kv_splits,
+            max_kv_splits,
             sm_scale,
             logit_cap,
         )

sglang/srt/layers/attention/triton_ops/extend_attention.py

@@ -341,12 +341,21 @@ def extend_attention_fwd(
             else:
                 BLOCK_M, BLOCK_N = (32, 64)
         elif is_cuda_available and CUDA_CAPABILITY[0] >= 8:
-            if Lq <= 128:
-                BLOCK_M, BLOCK_N = (128, 128)
-            elif Lq <= 256:
-                BLOCK_M, BLOCK_N = (64, 64)
+            # sm86/sm89 has a much smaller shared memory size (100K) than sm80 (160K)
+            if CUDA_CAPABILITY[1] == 9 or CUDA_CAPABILITY[1] == 6:
+                if Lq <= 128:
+                    BLOCK_M, BLOCK_N = (64, 128)
+                elif Lq <= 256:
+                    BLOCK_M, BLOCK_N = (64, 64)
+                else:
+                    BLOCK_M, BLOCK_N = (32, 32)
             else:
-                BLOCK_M, BLOCK_N = (32, 64)
+                if Lq <= 128:
+                    BLOCK_M, BLOCK_N = (128, 128)
+                elif Lq <= 256:
+                    BLOCK_M, BLOCK_N = (64, 64)
+                else:
+                    BLOCK_M, BLOCK_N = (32, 64)
         else:
             BLOCK_M, BLOCK_N = (64, 64) if Lq <= 128 else (32, 32)
 

sglang/srt/layers/attention/utils.py

@@ -15,6 +15,7 @@ def create_flashinfer_kv_indices_triton(
     BLOCK_SIZE: tl.constexpr = 512
     pid = tl.program_id(axis=0)
 
+    # find the req pool idx, this is for batch to token
     req_pool_index = tl.load(req_pool_indices_ptr + pid)
     kv_indices_offset = tl.load(kv_indptr + pid)
 
@@ -37,3 +38,55 @@ def create_flashinfer_kv_indices_triton(
             mask=mask,
         )
         tl.store(kv_indices_ptr + kv_indices_offset + offset, data, mask=mask)
+
+
+@triton.jit
+def create_flashmla_kv_indices_triton(
+    req_to_token_ptr,  # [max_batch, max_context_len]
+    req_pool_indices_ptr,
+    page_kernel_lens_ptr,
+    kv_start_idx,
+    kv_indices_ptr,
+    req_to_token_ptr_stride: tl.constexpr,
+    kv_indices_ptr_stride: tl.constexpr,
+):
+    PAGED_SIZE: tl.constexpr = 64
+    BLOCK_SIZE: tl.constexpr = 4096
+    NUM_PAGE_PER_BLOCK: tl.constexpr = 64
+    pid = tl.program_id(axis=0)
+
+    # find the req pool idx, this is for batch to token
+    req_pool_index = tl.load(req_pool_indices_ptr + pid)
+
+    kv_start = 0
+    kv_end = 0
+    if kv_start_idx:
+        kv_start = tl.load(kv_start_idx + pid).to(tl.int32)
+        kv_end = kv_start
+
+    kv_end += tl.load(page_kernel_lens_ptr + pid).to(tl.int32)
+
+    num_paged = tl.cdiv(kv_end - kv_start, PAGED_SIZE)
+    num_pages_loop = tl.cdiv(kv_end - kv_start, BLOCK_SIZE)
+
+    for i in range(num_pages_loop):
+        paged_offset = (
+            tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK
+        ) * PAGED_SIZE
+        paged_offset_out = tl.arange(0, NUM_PAGE_PER_BLOCK) + i * NUM_PAGE_PER_BLOCK
+
+        mask = paged_offset <= num_paged * PAGED_SIZE
+        mask_out = paged_offset_out <= num_paged
+
+        data = tl.load(
+            req_to_token_ptr
+            + req_pool_index * req_to_token_ptr_stride
+            + kv_start
+            + paged_offset,
+            mask=mask,
+        )
+        tl.store(
+            kv_indices_ptr + pid * kv_indices_ptr_stride + paged_offset_out,
+            data // PAGED_SIZE,
+            mask=mask_out,
+        )

sglang/srt/layers/attention/vision.py

@@ -19,34 +19,10 @@ from sglang.srt.layers.linear import (
     RowParallelLinear,
 )
 from sglang.srt.layers.quantization import QuantizationConfig
+from sglang.srt.layers.rotary_embedding import apply_rotary_pos_emb, rotate_half
 from sglang.srt.utils import add_prefix
 
 
-# Copied from transformers, modeling_qwen2_vl.py
-def rotate_half(x):
-    """Rotates half the hidden dims of the input."""
-    x1 = x[..., : x.shape[-1] // 2]
-    x2 = x[..., x.shape[-1] // 2 :]
-    return torch.cat((-x2, x1), dim=-1)
-
-
-def apply_rotary_pos_emb_vision(
-    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
-) -> Tuple[torch.Tensor, torch.Tensor]:
-    orig_q_dtype = q.dtype
-    orig_k_dtype = k.dtype
-    q, k = q.float(), k.float()
-
-    cos, sin = cos.unsqueeze(-2).float(), sin.unsqueeze(-2).float()
-    q_embed = (q * cos) + (rotate_half(q) * sin)
-    k_embed = (k * cos) + (rotate_half(k) * sin)
-
-    q_embed = q_embed.to(orig_q_dtype)
-    k_embed = k_embed.to(orig_k_dtype)
-
-    return q_embed, k_embed
-
-
 class VisionAttention(nn.Module):
     r"""
         Multi-headed attention without any cache, mostly used for ViT.
@@ -167,9 +143,14 @@ class VisionAttention(nn.Module):
         if position_embeddings is not None:
             cos, sin = position_embeddings
             original_shape = q.shape
-            q, k = q.view(s, head, -1), k.view(s, head, -1)
-            q, k = apply_rotary_pos_emb_vision(q, k, cos, sin)
-            q, k = q.reshape(original_shape), k.reshape(original_shape)
+            # [total_tokens, head, head_size]
+            q = q.view(-1, head, self.head_size)
+            k = k.view(-1, head, self.head_size)
+
+            q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+            q = q.view(original_shape)
+            k = k.view(original_shape)
 
         if self.use_qkv_parallel:
             pass

sglang/srt/layers/dp_attention.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import functools
 import logging
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING, List
 
 import torch
 import triton
@@ -38,7 +38,12 @@ def compute_dp_attention_world_info(enable_dp_attention, tp_rank, tp_size, dp_si
     return attn_tp_rank, attn_tp_size, dp_rank
 
 
-def initialize_dp_attention(enable_dp_attention, tp_rank, tp_size, dp_size):
+def initialize_dp_attention(
+    enable_dp_attention: bool,
+    tp_rank: int,
+    tp_size: int,
+    dp_size: int,
+):
     global _ATTN_TP_GROUP, _ATTN_TP_RANK, _ATTN_TP_SIZE, _DP_RANK, _DP_SIZE
 
     from sglang.srt.layers.sampler import SYNC_TOKEN_IDS_ACROSS_TP
@@ -46,7 +51,11 @@ def initialize_dp_attention(enable_dp_attention, tp_rank, tp_size, dp_size):
     _ATTN_TP_RANK, _ATTN_TP_SIZE, _DP_RANK = compute_dp_attention_world_info(
         enable_dp_attention, tp_rank, tp_size, dp_size
     )
-    _DP_SIZE = dp_size
+
+    if enable_dp_attention:
+        _DP_SIZE = dp_size
+    else:
+        _DP_SIZE = 1
 
     tp_group = get_tp_group()
     _ATTN_TP_GROUP = GroupCoordinator(
@@ -54,7 +63,7 @@ def initialize_dp_attention(enable_dp_attention, tp_rank, tp_size, dp_size):
             list(range(head, head + _ATTN_TP_SIZE))
             for head in range(0, tp_size, _ATTN_TP_SIZE)
         ],
-        tp_rank,
+        tp_group.local_rank,
         torch.distributed.get_backend(tp_group.device_group),
         SYNC_TOKEN_IDS_ACROSS_TP,
         False,
@@ -169,20 +178,19 @@ def memcpy_triton(dst, src, dim, offset, sz, offset_src):
     memcpy_triton_kernel[grid](dst, src, offset, sz, offset_src, chunk_size, BLOCK_SIZE)
 
 
-def dp_gather(
+def _dp_gather(
     global_tokens: torch.Tensor,
     local_tokens: torch.Tensor,
     forward_batch: ForwardBatch,
-    layer_id: Union[str, int],
+    is_partial: bool,
 ):
     local_start_pos, local_num_tokens = get_dp_local_info(forward_batch)
 
     global_tokens.fill_(0)
     assert local_tokens.is_contiguous()
     assert global_tokens.is_contiguous()
-    if local_tokens.shape[0] > 0 and (
-        layer_id != "embedding" or get_attention_tp_rank() == 0
-    ):
+
+    if local_tokens.shape[0] > 0 and (is_partial or get_attention_tp_rank() == 0):
         assert (
             global_tokens.untyped_storage().data_ptr()
             != local_tokens.untyped_storage().data_ptr()
@@ -205,6 +213,22 @@ def dp_gather(
         global_tokens[:] = tensor_model_parallel_all_reduce(global_tokens)
 
 
+def dp_gather_partial(
+    global_tokens: torch.Tensor,
+    local_tokens: torch.Tensor,
+    forward_batch: ForwardBatch,
+):
+    _dp_gather(global_tokens, local_tokens, forward_batch, is_partial=True)
+
+
+def dp_gather_replicate(
+    global_tokens: torch.Tensor,
+    local_tokens: torch.Tensor,
+    forward_batch: ForwardBatch,
+):
+    _dp_gather(global_tokens, local_tokens, forward_batch, is_partial=False)
+
+
 def dp_scatter(
     local_tokens: torch.Tensor,  # output
     global_tokens: torch.Tensor,  # input
@@ -227,14 +251,12 @@ def dp_scatter(
         )
 
 
-def get_do_logits_dp_scatter(forward_batch: ForwardBatch):
-    def do_logits_dp_scatter(logits: torch.Tensor):
-        local_logits = torch.empty(
-            (forward_batch.input_ids.shape[0], *logits.shape[1:]),
-            dtype=logits.dtype,
-            device=logits.device,
-        )
-        dp_scatter(local_logits, logits, forward_batch)
-        return local_logits
+def tp_reduce_scatter(
+    output: torch.Tensor,
+    input_list: List[torch.Tensor],
+):
+    return get_attention_tp_group().reduce_scatter(output, input_list)
+
 
-    return do_logits_dp_scatter
+def tp_all_gather(output_list: List[torch.Tensor], input_: torch.Tensor):
+    return get_attention_tp_group().all_gather(input_, tensor_list=output_list)