Add Speculative Decoding Eagle3 topk > 1

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -20,9 +20,10 @@
 
 from sglang.srt.utils import is_flashinfer_available
 
-if is_flashinfer_available():
-    from flashinfer.cascade import merge_state
+# if is_flashinfer_available():
+    # from flashinfer.cascade import merge_state
 
+from vllm.v1.attention.backends.flash_attn import merge_attn_states
 
 @dataclass
 class FlashAttentionMetadata:
@@ -474,15 +475,21 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
                 ]
 
                 metadata_expand = FlashAttentionMetadata()
-                metadata_expand.cache_seqlens_int32 = torch.full(
-                    (
-                        forward_batch.seq_lens.numel()
-                        * self.speculative_num_draft_tokens,
-                    ),
-                    self.speculative_num_draft_tokens,
+                # metadata_expand.cache_seqlens_int32 = torch.full(
+                #     (
+                #         forward_batch.seq_lens.numel()
+                #         * self.speculative_num_draft_tokens,
+                #     ),
+                #     self.speculative_num_draft_tokens,
+                #     device=device,
+                #     dtype=torch.int32,
+                # )
+                metadata_expand.cache_seqlens_int32 = torch.arange(
+                    1,
+                    self.speculative_num_draft_tokens + 1,
                     device=device,
                     dtype=torch.int32,
-                )
+                ).repeat(forward_batch.seq_lens.numel())
                 metadata_expand.max_seq_len_q = 1
                 metadata_expand.max_seq_len_k = self.speculative_num_draft_tokens
                 metadata_expand.cu_seqlens_q = torch.arange(
@@ -491,14 +498,20 @@ def init_forward_metadata(self, forward_batch: ForwardBatch):
                     dtype=torch.int32,
                     device=device,
                 )
-                metadata_expand.cu_seqlens_k = torch.arange(
-                    0,
-                    metadata_expand.cache_seqlens_int32.numel()
-                    * self.speculative_num_draft_tokens
-                    + 1,
-                    step=self.speculative_num_draft_tokens,
-                    dtype=torch.int32,
-                    device=device,
+                # metadata_expand.cu_seqlens_k = torch.arange(
+                #     0,
+                #     metadata_expand.cache_seqlens_int32.numel()
+                #     * self.speculative_num_draft_tokens
+                #     + 1,
+                #     step=self.speculative_num_draft_tokens,
+                #     dtype=torch.int32,
+                #     device=device,
+                # )
+                metadata.cu_seqlens_k = torch.nn.functional.pad(
+                    torch.cumsum(
+                        metadata.cache_seqlens_int32, dim=0, dtype=torch.int32
+                    ),
+                    (1, 0),
                 )
                 offsets = torch.arange(
                     self.speculative_num_draft_tokens, device=device
@@ -776,12 +789,20 @@ def forward_extend(
                     v_descale=v_descale,
                     return_softmax_lse=True,
                 )
-                o, _ = merge_state(
-                    o,
-                    softmax_lse.T.contiguous(),
-                    o_expand,
-                    softmax_lse_expand.T.contiguous(),
-                )
+                # o, _ = merge_state(
+                #     o,
+                #     softmax_lse.T.contiguous(),
+                #     o_expand,
+                #     softmax_lse_expand.T.contiguous(),
+                # )
+
+                # Run the kernel.
+                output = torch.empty_like(o)
+                merge_attn_states(output, 
+                                o, softmax_lse.contiguous(),
+                                o_expand,
+                                softmax_lse_expand.contiguous())
+                o = output
                 # torch.distributed.breakpoint()
         else:
             # Do absorbed multi-latent attention
@@ -932,12 +953,18 @@ def forward_decode(
                     v_descale=v_descale,
                     return_softmax_lse=True,
                 )
-                o, _ = merge_state(
-                    o,
-                    softmax_lse.T.contiguous(),
-                    o_expand,
-                    softmax_lse_expand.T.contiguous(),
-                )
+                # o, _ = merge_state(
+                #     o,
+                #     softmax_lse.T.contiguous(),
+                #     o_expand,
+                #     softmax_lse_expand.T.contiguous(),
+                # )
+                output = torch.empty_like(o)
+                merge_attn_states(output, 
+                                o, softmax_lse.contiguous(),
+                                o_expand,
+                                softmax_lse_expand.contiguous())
+                o = output
                 # print("draft decode merge finish! id:", self.speculative_step_id)
                 # torch.distributed.breakpoint()
         else:
@@ -1121,35 +1148,48 @@ def init_cuda_graph_state(self, max_bs: int):
             }
 
             self.target_verify_metadata_topk_expand = {
-                "cache_seqlens": torch.full(
-                    (max_bs * self.speculative_num_draft_tokens,),
-                    self.speculative_num_draft_tokens,
-                    dtype=torch.int32,
+                # "cache_seqlens": torch.full(
+                #     (max_bs * self.speculative_num_draft_tokens,),
+                #     self.speculative_num_draft_tokens,
+                #     dtype=torch.int32,
+                #     device=self.device,
+                # ),
+                "cache_seqlens": torch.arange(
+                    1,
+                    self.speculative_num_draft_tokens + 1,
                     device=self.device,
-                ),
+                    dtype=torch.int32,
+                ).repeat(max_bs),
                 "cu_seqlens_q": torch.arange(
                     0,
                     max_bs * self.speculative_num_draft_tokens + 1,
                     dtype=torch.int32,
                     device=self.device,
                 ),
-                "cu_seqlens_k": torch.arange(
-                    0,
-                    max_bs
-                    * self.speculative_num_draft_tokens
-                    * self.speculative_num_draft_tokens
-                    + 1,
-                    step=self.speculative_num_draft_tokens,
-                    dtype=torch.int32,
-                    device=self.device,
-                ),
+                # "cu_seqlens_k": torch.arange(
+                #     0,
+                #     max_bs
+                #     * self.speculative_num_draft_tokens
+                #     * self.speculative_num_draft_tokens
+                #     + 1,
+                #     step=self.speculative_num_draft_tokens,
+                #     dtype=torch.int32,
+                #     device=self.device,
+                # ),
                 "page_table": torch.zeros(
                     max_bs * self.speculative_num_draft_tokens,
                     self.speculative_num_draft_tokens,
                     dtype=torch.int32,
                     device=self.device,
                 ),
             }
+            self.target_verify_metadata_topk_expand["cu_seqlens_k"] = torch.nn.functional.pad(
+                    torch.cumsum(
+                        self.target_verify_metadata_topk_expand["cache_seqlens"], dim=0, dtype=torch.int32
+                    ),
+                    (1, 0),
+            )
+
 
         # Biao's Note: Consolidate the encoder metadata
 
@@ -1442,7 +1482,12 @@ def init_forward_metadata_replay_cuda_graph(
                         cache_loc[:, :decode_length].contiguous().to(torch.int32)
                     )
                     # may not need
-                    metadata_expand.page_table[cache_loc.shape[0] :].fill_(0)
+                    # metadata_expand.page_table[cache_loc.shape[0] :].fill_(0)
+                    # metadata.page_table[cache_loc.shape[0] :].fill_(0)
+
+                    # print("decode draft step id", self.speculative_step_id)
+                    # print("metadata", metadata)
+                    # print("metadata_expand", metadata_expand)
             else:
                 metadata = self.decode_cuda_graph_metadata[bs]
                 # Normal Decode
@@ -1509,6 +1554,7 @@ def init_forward_metadata_replay_cuda_graph(
                     req_pool_indices, : metadata.max_seq_len_k
                 ]
                 metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table)
+                # metadata.page_table[spec_info.positions.numel() :].fill_(0)
 
                 # 2. The second half of metadata for draft tokens (per_batch_num_tokens = topk)
                 metadata_expand = self.target_verify_metadata_topk_expand[bs]
@@ -1555,31 +1601,35 @@ def init_forward_metadata_replay_cuda_graph(
                     (non_masked_page_table * mask).to(torch.int32)
                 )
                 # may not need
-                metadata_expand.page_table[spec_info.positions.numel() :].fill_(0)
-
-        if encoder_lens is not None:
-            # Only support encoder size 1 for now
-            metadata.encoder_max_seq_len_k = encoder_lens[0]
-            metadata.encoder_lens_int32.copy_(encoder_lens[:1])
-            metadata.encoder_cu_seqlens_k.copy_(
-                torch.nn.functional.pad(
-                    torch.cumsum(metadata.encoder_lens_int32, dim=0, dtype=torch.int32),
-                    (1, 0),
-                )
-            )
-
-            metadata.encoder_page_table[:, : metadata.encoder_max_seq_len_k].copy_(
-                self.req_to_token[req_pool_indices, : metadata.encoder_max_seq_len_k]
-            )
+                # metadata_expand.page_table[spec_info.positions.numel() :].fill_(0)
+
+                # print("target verify")
+                # print("metadata", metadata)
+                # print("metadata_expand", metadata_expand)
 
-            # Update the regular page table
-            page_table = self.req_to_token[
-                req_pool_indices,
-                metadata.encoder_max_seq_len_k : (
-                    metadata.encoder_max_seq_len_k + metadata.max_seq_len_k
-                ),
-            ]
-            metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table)
+        # if encoder_lens is not None:
+        #     # Only support encoder size 1 for now
+        #     metadata.encoder_max_seq_len_k = encoder_lens[0]
+        #     metadata.encoder_lens_int32.copy_(encoder_lens[:1])
+        #     metadata.encoder_cu_seqlens_k.copy_(
+        #         torch.nn.functional.pad(
+        #             torch.cumsum(metadata.encoder_lens_int32, dim=0, dtype=torch.int32),
+        #             (1, 0),
+        #         )
+        #     )
+
+        #     metadata.encoder_page_table[:, : metadata.encoder_max_seq_len_k].copy_(
+        #         self.req_to_token[req_pool_indices, : metadata.encoder_max_seq_len_k]
+        #     )
+
+        #     # Update the regular page table
+        #     page_table = self.req_to_token[
+        #         req_pool_indices,
+        #         metadata.encoder_max_seq_len_k : (
+        #             metadata.encoder_max_seq_len_k + metadata.max_seq_len_k
+        #         ),
+        #     ]
+        #     metadata.page_table[:, : metadata.max_seq_len_k].copy_(page_table)
 
         self.forward_metadata = metadata
         self.forward_metadata_spec_decode_expand = metadata_expand