Use trtllm mha decode kernel for target_verify in speculative decoding (sgl-project#13976)

Co-authored-by: Kangyan-Zhou <[email protected]>

Author: Qiaolin-Yu

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

@@ -377,6 +377,7 @@ def init_forward_metadata_replay_cuda_graph(
             ]
             page_indices //= self.page_size
             metadata.page_table[:, :max_seq_pages].copy_(page_indices)
+            metadata.max_seq_len_q = self.speculative_num_draft_tokens
         elif forward_mode.is_draft_extend():
             metadata = self.draft_extend_metadata[bs]
             metadata.cache_seqlens_int32.copy_(seq_lens)
@@ -614,24 +615,42 @@ def forward_extend(
         bmm1_scale = q_scale * k_scale * layer.scaling
         bmm2_scale = 1.0
 
-        o = flashinfer.prefill.trtllm_batch_context_with_kv_cache(
-            query=q,
-            kv_cache=kv_cache,
-            workspace_buffer=self.workspace_buffer,
-            block_tables=self.forward_metadata.page_table,
-            seq_lens=self.forward_metadata.cache_seqlens_int32,
-            max_q_len=self.forward_metadata.max_seq_len_q,
-            max_kv_len=self.max_context_len,
-            bmm1_scale=bmm1_scale,
-            bmm2_scale=bmm2_scale,
-            batch_size=forward_batch.batch_size,
-            cum_seq_lens_q=self.forward_metadata.cu_seqlens_q,
-            cum_seq_lens_kv=self.forward_metadata.cu_seqlens_k,
-            window_left=layer.sliding_window_size,
-            # TODO: add attention_sink operation or nvfp4 scale factor if needed
-            sinks=attention_sink,
-            out_dtype=self.q_data_type,  # model_runner.dtype
-        )
+        if forward_batch.forward_mode.is_target_verify():
+            o = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
+                query=q,
+                kv_cache=kv_cache,
+                workspace_buffer=self.workspace_buffer,
+                block_tables=self.forward_metadata.page_table,
+                seq_lens=self.forward_metadata.cache_seqlens_int32,
+                max_seq_len=self.max_context_len,
+                bmm1_scale=bmm1_scale,
+                bmm2_scale=bmm2_scale,
+                window_left=layer.sliding_window_size,
+                # TODO: add attention_sink operation or nvfp4 scale factor if needed
+                sinks=attention_sink,
+                out_dtype=self.q_data_type,  # model_runner.dtype
+                q_len_per_req=self.forward_metadata.max_seq_len_q,
+            )
+        else:
+
+            o = flashinfer.prefill.trtllm_batch_context_with_kv_cache(
+                query=q,
+                kv_cache=kv_cache,
+                workspace_buffer=self.workspace_buffer,
+                block_tables=self.forward_metadata.page_table,
+                seq_lens=self.forward_metadata.cache_seqlens_int32,
+                max_q_len=self.forward_metadata.max_seq_len_q,
+                max_kv_len=self.max_context_len,
+                bmm1_scale=bmm1_scale,
+                bmm2_scale=bmm2_scale,
+                batch_size=forward_batch.batch_size,
+                cum_seq_lens_q=self.forward_metadata.cu_seqlens_q,
+                cum_seq_lens_kv=self.forward_metadata.cu_seqlens_k,
+                window_left=layer.sliding_window_size,
+                # TODO: add attention_sink operation or nvfp4 scale factor if needed
+                sinks=attention_sink,
+                out_dtype=self.q_data_type,  # model_runner.dtype
+            )
 
         return o.view(-1, layer.tp_q_head_num * layer.head_dim)