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

sglang/srt/speculative/build_eagle_tree.py

@@ -1,122 +1,175 @@
-import cutex
+# NOTE: Please run this file to make sure the test cases are correct.
+
+from typing import List
+
 import torch
 
-# parent_table [bs,topk*depth+)]
-# selected_index [bs,draft_token_num-1)]
-# verified_seq_len [bs]
-# tree_mask [draft_token*(seq_len[0]+draft_token) | draft_token*(seq_len[1]+draft_token) | ..] = [sum(verified_seq_len)*draft_token+bs*draft_token*draft_token]
-# positions [bs*draft_token]
-# retrive_index [b, draft_token, depth+2]
-kernels = cutex.SourceModule(
-    """
-//cuda
-__global__ void build_tree(Tensor<long, 2> parent_list, Tensor<long, 2> selected_index, Tensor<int, 1> verified_seq_len,
-        Tensor<bool, 1> tree_mask, Tensor<long, 1> positions, Tensor<long, 3> retrive_index, int topk, int depth, int draft_token_num) {
-        int bid = blockIdx.x;
-        int tid = threadIdx.x;
-        if (tid >= draft_token_num){
-            return;
-        }
-        int seq_tree_idx = draft_token_num * draft_token_num * bid;
-        for(int i=0; i<bid; i++){
-            seq_tree_idx += verified_seq_len[i] * draft_token_num;
-        }
-        int seq_len = verified_seq_len[bid];
-        int token_tree_idx = seq_tree_idx + (seq_len+draft_token_num)*tid + seq_len + 1;
-        for(int i=0; i<draft_token_num-1; i++){
-            tree_mask[token_tree_idx+i] = false;
-        }
-
-        int position = 0;
-        if (tid==0){
-            positions[bid*draft_token_num] = seq_len;
-            retrive_index[bid][0][0] = bid * draft_token_num;
-            return;
-        }
-
-        int depends_order[10];
-
-        int cur_position = tid-1;
-        while(true){
-            depends_order[position] = cur_position+1;
-            position += 1;
-            tree_mask[token_tree_idx+cur_position] = true;
-            int parent_tb_idx = selected_index[bid][cur_position]/topk;
-            if(parent_tb_idx==0){
-                break;
-            }
-
-            int token_idx = parent_list[bid][parent_tb_idx];
-            for(cur_position=0; cur_position<draft_token_num;cur_position++){
-                if(selected_index[bid][cur_position]==token_idx){
-                    break;
-                }
-            }
-        }
-        positions[bid*draft_token_num+tid] = position + seq_len;
-
-        int is_leaf = 0;
-        for(int i=1;i<draft_token_num;i++){
-            if(tree_mask[seq_tree_idx + i * (draft_token_num+seq_len) + seq_len + tid])
-            {
-                is_leaf ++;
-            }
-        }
-        if(is_leaf==1){
-            for(int i=0; i<position; i++){
-                retrive_index[bid][tid][position-i] = depends_order[i] + bid * draft_token_num;
-            }
-            retrive_index[bid][tid][0] = bid*draft_token_num;
-        }
-
-
-
-}
-//!cuda
-""",
-    float_bits=16,  # change to 16 to use half precision as `float` type in the above source code.
-    boundscheck=True,  # turning on for debug and off for performance (to use full threads of a block), default is on.
-)
-
-
-def build_tree_kernel(parent_list, top_score_index, seq_lens, topk, depth, draft_token):
+from sglang.srt.utils import is_cuda_available
+
+if is_cuda_available():
+    from sgl_kernel import build_tree_kernel as sgl_build_tree_kernel
+    from sgl_kernel import (
+        build_tree_kernel_efficient as sgl_build_tree_kernel_efficient,
+    )
+
+
+def build_tree_kernel_efficient_preprocess(
+    verified_id: torch.Tensor,
+    score_list: List[torch.Tensor],
+    token_list: List[torch.Tensor],
+    parents_list: List[torch.Tensor],
+    num_verify_tokens: int,
+):
+    score_list = torch.cat(score_list, dim=1).flatten(
+        1
+    )  # b, n, topk; n= 1 + (num_steps-1) * self.topk
+    ss_token_list = torch.cat(
+        token_list, dim=1
+    )  # b, (self.topk + (num_steps-1) * self.topk)
+    top_scores = torch.topk(score_list, num_verify_tokens - 1, dim=-1)
+    top_scores_index = top_scores.indices
+    top_scores_index = torch.sort(top_scores_index).values
+
+    draft_tokens = torch.gather(ss_token_list, index=top_scores_index, dim=1)
+    draft_tokens = torch.cat((verified_id.unsqueeze(1), draft_tokens), dim=1).flatten()
+    parent_list = torch.cat(parents_list[:-1], dim=1)
+
+    return parent_list, top_scores_index, draft_tokens
+
+
+def build_tree_kernel_efficient(
+    verified_id: torch.Tensor,
+    score_list: List[torch.Tensor],
+    token_list: List[torch.Tensor],
+    parents_list: List[torch.Tensor],
+    seq_lens: torch.Tensor,
+    seq_lens_sum: int,
+    topk: int,
+    spec_steps: int,
+    num_verify_tokens: int,
+):
+    parent_list, top_scores_index, draft_tokens = (
+        build_tree_kernel_efficient_preprocess(
+            verified_id,
+            score_list,
+            token_list,
+            parents_list,
+            num_verify_tokens,
+        )
+    )
+
+    # seq_lens_sum == sum(seq_lens); seq_lens: sequence length without draft tokens
     bs = seq_lens.numel()
-    device = parent_list.device
+    device = seq_lens.device
+    # e.g. for bs=1, tree_mask: num_draft_token, seq_lens_sum + num_draft_token (flattened)
+    # where each row indicates the attending pattern of each draft token
+    # TODO: make them torch.empty and fuse them into `sgl_build_tree_kernel`
     tree_mask = torch.full(
-        (torch.sum(seq_lens).item() * draft_token + draft_token * draft_token * bs,),
+        (
+            seq_lens_sum * num_verify_tokens
+            + num_verify_tokens * num_verify_tokens * bs,
+        ),
         True,
         device=device,
     )
     retrive_index = torch.full(
-        (bs, draft_token, depth + 2), -1, device=device, dtype=torch.long
+        (bs, num_verify_tokens), -1, device=device, dtype=torch.long
+    )
+    retrive_next_token = torch.full(
+        (bs, num_verify_tokens), -1, device=device, dtype=torch.long
     )
-    positions = torch.empty((bs * draft_token,), device=device, dtype=torch.long)
+    retrive_next_sibling = torch.full(
+        (bs, num_verify_tokens), -1, device=device, dtype=torch.long
+    )
+    # position: where each token belongs to
+    # e.g. if depth of each draft token is [0, 1, 1, 2] and the prompt length is 7
+    # then, positions = [7, 8, 8, 9]
+    positions = torch.empty((bs * num_verify_tokens,), device=device, dtype=torch.long)
 
-    kernels.build_tree(
+    sgl_build_tree_kernel_efficient(
         parent_list,
-        top_score_index,
+        top_scores_index,
         seq_lens.to(torch.int32),
         tree_mask,
         positions,
         retrive_index,
+        retrive_next_token,
+        retrive_next_sibling,
         topk,
-        depth,
-        draft_token,
-        grid=(bs, 1, 1),
-        block=(64, 1, 1),
+        spec_steps,
+        num_verify_tokens,
+    )
+    return (
+        tree_mask,
+        positions,
+        retrive_index,
+        retrive_next_token,
+        retrive_next_sibling,
+        draft_tokens,
+    )
+
+
+def build_tree_kernel(
+    verified_id: torch.Tensor,
+    score_list: List[torch.Tensor],
+    token_list: List[torch.Tensor],
+    parents_list: List[torch.Tensor],
+    seq_lens: torch.Tensor,
+    seq_lens_sum: int,
+    topk: int,
+    spec_steps: int,
+    num_verify_tokens: int,
+):
+    parent_list, top_scores_index, draft_tokens = (
+        build_tree_kernel_efficient_preprocess(
+            verified_id,
+            score_list,
+            token_list,
+            parents_list,
+            num_verify_tokens,
+        )
+    )
+
+    bs = seq_lens.numel()
+    device = seq_lens.device
+
+    tree_mask = torch.full(
+        (
+            seq_lens_sum * num_verify_tokens
+            + num_verify_tokens * num_verify_tokens * bs,
+        ),
+        True,
+        device=device,
+    )
+    retrive_index = torch.full(
+        (bs, num_verify_tokens, spec_steps + 2), -1, device=device, dtype=torch.long
+    )
+    positions = torch.empty((bs * num_verify_tokens,), device=device, dtype=torch.long)
+
+    sgl_build_tree_kernel(
+        parent_list,
+        top_scores_index,
+        seq_lens.to(torch.int32),
+        tree_mask,
+        positions,
+        retrive_index,
+        topk,
+        spec_steps,
+        num_verify_tokens,
     )
-    index = retrive_index.sum(dim=-1) != -depth - 2
+
+    index = retrive_index.sum(dim=-1) != -spec_steps - 2
     cum_len = torch.cumsum(torch.sum(index, dim=-1), dim=-1)
     retrive_cum_len = torch.zeros(
         (cum_len.numel() + 1,), dtype=torch.int32, device="cuda"
     )
     retrive_cum_len[1:] = cum_len
+    # TODO: this indexing cause a synchronization, optimize this
     retrive_index = retrive_index[index]
-    return tree_mask, positions, retrive_index, retrive_cum_len
+    return tree_mask, positions, retrive_index, retrive_cum_len, draft_tokens
 
 
-if __name__ == "__main__":
-
+def test_build_tree_kernel():
     def findp(p_i, index, parent_list):
         pos = index // 10
         index_list = index.tolist()
@@ -309,21 +362,21 @@ if __name__ == "__main__":
     bs = verified_seq_len.shape[0]
     topk = 10
     depth = 5  # depth <= 10
-    draft_token = 64
+    num_draft_token = 64
 
     tree_mask = torch.full(
         (
-            torch.sum(verified_seq_len).item() * draft_token
-            + draft_token * draft_token * bs,
+            torch.sum(verified_seq_len).item() * num_draft_token
+            + num_draft_token * num_draft_token * bs,
         ),
         True,
     ).cuda()
     retrive_index = torch.full(
-        (bs, draft_token, depth + 2), -1, device="cuda", dtype=torch.long
+        (bs, num_draft_token, depth + 2), -1, device="cuda", dtype=torch.long
     )
-    positions = torch.empty((bs * draft_token,), device="cuda", dtype=torch.long)
+    positions = torch.empty((bs * num_draft_token,), device="cuda", dtype=torch.long)
 
-    kernels.build_tree(
+    sgl_build_tree_kernel(
         parent_list.unsqueeze(0),
         index.unsqueeze(0),
         verified_seq_len,
@@ -332,16 +385,345 @@ if __name__ == "__main__":
         retrive_index,
         topk,
         depth,
-        draft_token,
-        grid=(bs, 1, 1),
-        block=(64, 1, 1),
+        num_draft_token,
     )
+
     retrive_index = retrive_index[retrive_index.sum(dim=-1) != -depth - 2]
 
     c_mask, c_positions, c_retive_index = create_mask(
-        verified_seq_len, draft_token, index, parent_list, depth
+        verified_seq_len, num_draft_token, index, parent_list, depth
     )
 
     assert torch.allclose(tree_mask, c_mask), "tree mask has error."
     assert torch.allclose(positions, c_positions), "positions has error."
     assert torch.allclose(retrive_index, c_retive_index), "retrive_index has error."
+
+
+def test_build_tree_kernel_efficient():
+    verified_id = torch.tensor([29974, 13], device="cuda", dtype=torch.int32)
+    score_list = [
+        torch.tensor(
+            [
+                [[7.1127e-01, 2.8292e-01, 2.2995e-03, 1.7357e-03]],
+                [[9.7476e-01, 2.2219e-02, 6.5031e-04, 1.3212e-04]],
+            ],
+            dtype=torch.float32,
+            device="cuda",
+        ),
+        torch.tensor(
+            [
+                [
+                    [6.9142e-01, 1.2863e-02, 1.6873e-03, 1.1871e-03],
+                    [2.4787e-01, 1.8818e-02, 1.4204e-02, 9.2235e-04],
+                    [2.2971e-03, 1.6700e-06, 1.8737e-07, 8.3146e-08],
+                    [1.2771e-03, 2.4374e-04, 1.7832e-04, 1.1947e-05],
+                ],
+                [
+                    [8.4832e-02, 6.6068e-02, 5.8304e-02, 5.7851e-02],
+                    [2.3616e-03, 1.1243e-03, 5.4368e-04, 2.7768e-04],
+                    [2.5286e-04, 1.5578e-04, 2.8817e-05, 1.2888e-05],
+                    [1.2834e-04, 2.5417e-06, 1.1279e-06, 1.6088e-08],
+                ],
+            ],
+            dtype=torch.float32,
+            device="cuda",
+        ),
+        torch.tensor(
+            [
+                [
+                    [6.6438e-01, 2.6997e-02, 2.4236e-05, 4.0821e-06],
+                    [2.4402e-01, 2.8409e-03, 5.0935e-04, 2.9022e-04],
+                    [1.6178e-02, 2.0567e-03, 4.5892e-04, 3.0034e-05],
+                    [1.3023e-02, 5.0497e-04, 3.6371e-04, 8.7750e-05],
+                ],
+                [
+                    [2.3263e-02, 2.0054e-02, 9.3990e-03, 2.7783e-03],
+                    [6.4156e-02, 5.5506e-04, 1.0429e-04, 9.7211e-05],
+                    [4.9950e-02, 5.0630e-03, 9.0068e-04, 3.3656e-04],
+                    [7.5817e-03, 8.5731e-04, 6.9972e-04, 6.0793e-04],
+                ],
+            ],
+            dtype=torch.float32,
+            device="cuda",
+        ),
+        torch.tensor(
+            [
+                [
+                    [6.6420e-01, 1.0525e-04, 6.5864e-05, 1.2253e-06],
+                    [1.3019e-01, 1.0461e-01, 5.2083e-03, 1.6777e-03],
+                    [2.0103e-02, 6.7335e-03, 1.2625e-04, 1.0364e-05],
+                    [1.5142e-02, 7.0819e-04, 9.6595e-05, 8.7951e-05],
+                ],
+                [
+                    [5.8608e-02, 1.8840e-03, 7.8535e-04, 4.4400e-04],
+                    [1.2185e-02, 2.0684e-03, 1.7418e-03, 1.4327e-03],
+                    [6.2455e-03, 6.1487e-03, 2.6862e-03, 1.8034e-03],
+                    [1.8590e-03, 1.6151e-03, 1.2481e-03, 3.6038e-04],
+                ],
+            ],
+            dtype=torch.float32,
+            device="cuda",
+        ),
+    ]
+    token_list = [
+        torch.tensor(
+            [[29896, 29906, 29900, 29945], [13, 2, 29871, 28956]],
+            dtype=torch.int64,
+            device="cuda",
+        ),
+        torch.tensor(
+            [
+                [
+                    29889,
+                    29974,
+                    29945,
+                    29900,
+                    29974,
+                    29922,
+                    29930,
+                    29958,
+                    29889,
+                    29974,
+                    29930,
+                    29945,
+                    29974,
+                    29922,
+                    29930,
+                    29958,
+                ],
+                [
+                    22550,
+                    4136,
+                    16492,
+                    8439,
+                    29871,
+                    2,
+                    3001,
+                    13,
+                    2,
+                    13,
+                    29906,
+                    29946,
+                    2,
+                    13,
+                    29871,
+                    259,
+                ],
+            ],
+            device="cuda",
+        ),
+        torch.tensor(
+            [
+                [
+                    29946,
+                    29945,
+                    29953,
+                    29906,
+                    29896,
+                    29945,
+                    29900,
+                    29906,
+                    29896,
+                    29945,
+                    29906,
+                    29953,
+                    29896,
+                    29945,
+                    29906,
+                    29946,
+                ],
+                [
+                    29871,
+                    2,
+                    29901,
+                    29889,
+                    29871,
+                    2,
+                    395,
+                    259,
+                    29901,
+                    29871,
+                    2,
+                    29889,
+                    3001,
+                    1234,
+                    7146,
+                    2186,
+                ],
+            ],
+            device="cuda",
+        ),
+        torch.tensor(
+            [
+                [
+                    29946,
+                    29974,
+                    29945,
+                    29930,
+                    29889,
+                    29922,
+                    29974,
+                    29930,
+                    29974,
+                    29946,
+                    29930,
+                    29922,
+                    29889,
+                    29974,
+                    29945,
+                    29922,
+                ],
+                [
+                    29941,
+                    29906,
+                    2,
+                    29946,
+                    29871,
+                    450,
+                    319,
+                    14990,
+                    29946,
+                    29941,
+                    2,
+                    29906,
+                    29871,
+                    2,
+                    3001,
+                    13,
+                ],
+            ],
+            device="cuda",
+        ),
+    ]
+    parents_list = [
+        torch.tensor(
+            [[-1, 0, 1, 2, 3], [-1, 0, 1, 2, 3]], dtype=torch.int64, device="cuda"
+        ),
+        torch.tensor([[4, 8, 9, 10], [4, 5, 6, 7]], dtype=torch.int64, device="cuda"),
+        torch.tensor(
+            [[20, 24, 21, 28], [24, 28, 20, 21]], dtype=torch.int64, device="cuda"
+        ),
+        torch.tensor(
+            [[36, 40, 41, 44], [36, 40, 44, 45]], dtype=torch.int64, device="cuda"
+        ),
+    ]
+    seq_lens = torch.tensor([5, 10], dtype=torch.int64, device="cuda")
+    topk = 4
+    depth = 4
+    num_draft_token = 8
+
+    tree_mask, position, retrive_index, retrive_cum_len, draft_tokens = (
+        build_tree_kernel(
+            verified_id=verified_id,
+            score_list=score_list,
+            token_list=token_list,
+            parents_list=parents_list,
+            seq_lens=seq_lens,
+            seq_lens_sum=torch.sum(seq_lens).item(),
+            topk=topk,
+            spec_steps=depth,
+            num_verify_tokens=num_draft_token,
+        )
+    )
+
+    from sglang.srt.utils import first_rank_print
+
+    first_rank_print("=========== build tree kernel ==========")
+    # first_rank_print(f"{tree_mask=}", flush=True)
+    first_rank_print(f"{position=}", flush=True)
+    first_rank_print(f"{retrive_index=}", flush=True)
+    first_rank_print(f"{retrive_cum_len=}", flush=True)
+    first_rank_print(f"{draft_tokens=}", flush=True)
+    assert position.tolist() == [5, 6, 6, 7, 7, 8, 8, 9, 10, 11, 12, 12, 12, 12, 13, 14]
+    assert retrive_index.tolist() == [
+        [0, -1, -1, -1, -1, -1],
+        [0, 2, 4, 6, -1, -1],
+        [0, 1, 3, 5, 7, -1],
+        [8, -1, -1, -1, -1, -1],
+        [8, 9, 10, -1, -1, -1],
+        [8, 9, 12, -1, -1, -1],
+        [8, 9, 13, -1, -1, -1],
+        [8, 9, 11, 14, 15, -1],
+    ]
+    assert retrive_cum_len.tolist() == [0, 3, 8]
+    assert draft_tokens.tolist() == [
+        29974,
+        29896,
+        29906,
+        29889,
+        29974,
+        29946,
+        29896,
+        29946,
+        13,
+        13,
+        22550,
+        4136,
+        16492,
+        8439,
+        29871,
+        29941,
+    ]
+
+    (
+        tree_mask,
+        position,
+        retrive_index,
+        retrive_next_token,
+        retrive_next_sibling,
+        draft_tokens,
+    ) = build_tree_kernel_efficient(
+        verified_id=verified_id,
+        score_list=score_list,
+        token_list=token_list,
+        parents_list=parents_list,
+        seq_lens=seq_lens,
+        seq_lens_sum=torch.sum(seq_lens).item(),
+        topk=topk,
+        spec_steps=depth,
+        num_verify_tokens=num_draft_token,
+    )
+
+    first_rank_print("=========== build tree kernel efficient ==========")
+    # first_rank_print(f"{tree_mask=}", flush=True)
+    first_rank_print(f"{position=}", flush=True)
+    first_rank_print(f"{retrive_index=}", flush=True)
+    first_rank_print(f"{retrive_next_token=}", flush=True)
+    first_rank_print(f"{retrive_next_sibling=}", flush=True)
+    first_rank_print(f"{draft_tokens=}", flush=True)
+    assert position.tolist() == [5, 6, 6, 7, 7, 8, 8, 9, 10, 11, 12, 12, 12, 12, 13, 14]
+    assert retrive_index.tolist() == [
+        [0, 1, 2, 3, 4, 5, 6, 7],
+        [8, 9, 10, 11, 12, 13, 14, 15],
+    ]
+    assert retrive_next_token.tolist() == [
+        [1, 3, 4, 5, 6, 7, -1, -1],
+        [1, 2, -1, 6, -1, -1, 7, -1],
+    ]
+    assert retrive_next_sibling.tolist() == [
+        [-1, 2, -1, -1, -1, -1, -1, -1],
+        [-1, -1, 3, 4, 5, -1, -1, -1],
+    ]
+    assert draft_tokens.tolist() == [
+        29974,
+        29896,
+        29906,
+        29889,
+        29974,
+        29946,
+        29896,
+        29946,
+        13,
+        13,
+        22550,
+        4136,
+        16492,
+        8439,
+        29871,
+        29941,
+    ]
+
+
+if __name__ == "__main__":
+    test_build_tree_kernel_efficient()
+    test_build_tree_kernel()