rebase code

Author: BJWang-ant

python/sglang/srt/layers/attention/nsa/index_buf_accessor.py

@@ -186,32 +186,6 @@ def triton(
             index_head_dim=pool.index_head_dim,
         )
 
-class GetKEAndKS:
-    @classmethod
-    def execute(cls, *args, **kwargs):
-        return cls.triton(*args, **kwargs)
-
-    @classmethod
-    def triton(
-        cls, pool: "NSATokenToKVPool",extend_sum_seq_len: int, seq_lens_tensor: torch.Tensor, 
-        extend_seq_lens_tensor: torch.Tensor, seq_lens_expanded_tensor: torch.Tensor,
-    ):
-        """
-        Triton implementation for generate ke and ks data for all batch in a single call.
-
-        :param page_indices: (num_pages,), int32/int64
-        :return: tuple of (k_fp8, k_scale) where
-                 k_fp8: (seq_len, index_head_dim), uint8
-                 k_scale: (seq_len, 4), uint8
-        """
-        return _get_ke_and_ks_triton(
-            extend_sum_seq_len=extend_sum_seq_len,
-            seq_lens=seq_lens_tensor,
-            extend_seq_lens=extend_seq_lens_tensor,
-            seq_lens_expanded=seq_lens_expanded_tensor,
-        )
-
-
 class SetK:
     @classmethod
     def execute(cls, *args, buf, **kwargs):
@@ -740,108 +714,4 @@ def _get_k_and_s_triton_kernel(
 
     # Store S to output
     s_dst_offset = token_id * 4
-    tl.store(s_out_ptr + s_dst_offset + s_offsets + s_offset_batch, s_data, mask=s_mask)
-
-
-def _get_ke_and_ks_triton(
-    extend_sum_seq_len: int,
-    seq_lens: torch.Tensor,
-    extend_seq_lens: torch.Tensor,
-    seq_lens_expanded: torch.Tensor
-):
-    """
-    Fused gather of both K (key) and S (scale) data from paged buffer using Triton.
-    This is more efficient than calling GetK and GetS separately.
-    for example:
-    seq_lens = [20, 30, 40, 50]
-    extend_seq_lens = [10, 15, 20, 10]
-    seq_lens_expanded = [10,11,...,19(bs0),15,16,...,29(bs1),40,41,...,49]
-
-    prefix_sum = [0,10,25,45,55]
-    seq_lens_sum = [0, 20, 50, 90]
-
-    :param extend_sum_seq_len: sum of all extend sequence len, int32
-    :param seq_lens: (num_pages, page_size * 128 + page_size * 4), int32
-    :param extend_seq_lens: (num_pages,), int32
-    :param seq_lens_expanded: int, number of tokens to gather
-    :return: tuple of (ks, ke) where
-             ks: (sum_extend_seq_len,), int32
-             ke: (sum_extend_seq_len,), int32
-    """
-    
-    ks = torch.empty((extend_sum_seq_len), dtype=torch.int32, device="cuda")
-    ke = torch.empty((extend_sum_seq_len), dtype=torch.int32, device="cuda")
-
-    max_iter = math.ceil(math.log2(extend_sum_seq_len)) + 1 if extend_sum_seq_len > 0 else 1
-
-    BLOCK_SIZE = 256
-    grid = lambda meta: (triton.cdiv(extend_sum_seq_len, meta['BLOCK_SIZE']),)
-    _get_ke_ks_triton_kernel[grid](
-        seq_lens_ptr=seq_lens,
-        extend_seq_lens_ptr=extend_seq_lens,
-        seq_lens_expanded=seq_lens_expanded,
-        ks_out_ptr=ks,
-        ke_out_ptr=ke,
-        seq_num=extend_seq_lens.shape[0],
-        extend_seq_lens_sum=extend_sum_seq_len,
-        iter_num=max_iter,
-        BLOCK_SIZE=BLOCK_SIZE,
-    )
-
-    return ks, ke
-
-
-@triton.jit
-def _get_ke_ks_triton_kernel(
-    seq_lens_ptr,
-    extend_seq_lens_ptr,
-    seq_lens_expanded,
-    ks_out_ptr,
-    ke_out_ptr,
-    seq_num: tl.constexpr,
-    extend_seq_lens_sum: tl.constexpr,
-    iter_num: tl.constexpr,
-    BLOCK_SIZE: tl.constexpr,
-):
-    '''
-    Get ke and ks fuse kernel.
-    '''
-    pid = tl.program_id(axis=0)
-    if pid >= extend_seq_lens_sum:
-        return
-    
-    start_pos = pid * BLOCK_SIZE
-    out_pos = tl.arange(0, BLOCK_SIZE) + start_pos
-    pos_mask = out_pos < extend_seq_lens_sum
-    
-    low = tl.zeros((BLOCK_SIZE,), dtype=tl.int32)
-    high = tl.full((BLOCK_SIZE,), seq_num, dtype=tl.int32)
-    for _ in range(iter_num):
-        mid = (low + high) // 2
-        
-        prefix_mid = tl.zeros((BLOCK_SIZE,), dtype=tl.int32)
-        for j in range(seq_num):
-            j_lt_mid = (j < mid) & pos_mask
-            extend_seq_len_j = tl.load(extend_seq_lens_ptr + j)
-            prefix_mid = tl.where(j_lt_mid, prefix_mid + extend_seq_len_j, prefix_mid)
-        
-        cond = out_pos >= prefix_mid
-        low = tl.where(cond, mid, low)
-        high = tl.where(~cond, mid, high)
-    
-    i = low
-    out_mask = (i >= 0) & (i < seq_num) & pos_mask
-    
-    seq_lens_sum_val = tl.zeros((BLOCK_SIZE,), dtype=tl.int32)
-    for j in range(seq_num):
-        j_lt_i = (j < i) & out_mask
-
-        seq_len_j = tl.load(seq_lens_ptr + j)
-        seq_len_j = tl.cast(seq_len_j, tl.int32)
-        seq_lens_sum_val = tl.where(j_lt_i, seq_lens_sum_val + seq_len_j, seq_lens_sum_val)
-    
-    D_val = tl.load(seq_lens_expanded + out_pos, mask=out_mask)
-    
-    tl.store(ks_out_ptr + out_pos, seq_lens_sum_val, mask=out_mask)
-    store_val = seq_lens_sum_val + D_val
-    tl.store(ke_out_ptr + out_pos, store_val, mask=out_mask)
+    tl.store(s_out_ptr + s_dst_offset + s_offsets + s_offset_batch, s_data, mask=s_mask)

python/sglang/srt/mem_cache/memory_pool.py

@@ -1753,31 +1753,6 @@ def get_index_k_scale_buffer(
             self, buf, page_indices=page_indices, seq_len_tensor=seq_len_tensor,
             seq_len_sum=seq_len_sum, max_seq_len=max_seq_len,
         )
-    
-    def get_ks_ke_buffer(
-        self,
-        extend_sum_seq_len: int,
-        seq_lens_tensor: torch.Tensor, 
-        extend_seq_lens_tensor: torch.Tensor, 
-        seq_lens_expanded_tensor: torch.Tensor,
-    ):
-        """
-        Fused method to get both index K and scale data in a single call using Triton.
-        More efficient than calling get_index_k_continuous and get_index_k_scale_continuous separately.
-
-        :param layer_id: Layer index
-        :param seq_len: Sequence length
-        :param page_indices: Page indices tensor
-        :return: tuple of (k_fp8, k_scale) where
-                 k_fp8: (seq_len, index_head_dim), uint8
-                 k_scale: (seq_len, 4), uint8
-        """
-        return index_buf_accessor.GetKEAndKS.execute(
-            self,extend_sum_seq_len=extend_sum_seq_len,
-            seq_lens_tensor=seq_lens_tensor,
-            extend_seq_lens_tensor=extend_seq_lens_tensor,
-            seq_lens_expanded_tensor=seq_lens_expanded_tensor,
-        )
 
     def set_index_k_scale_buffer(
         self,