minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Kangyan-Zhou <zky314343421@gmail.com>
This commit is contained in:
Qiaolin Yu
2025-11-26 20:40:34 -08:00
committed by GitHub
co-authored by Kangyan-Zhou
parent 5443db8759
commit 7cb04dc0e5
1 changed files with 37 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/trtllm_mha_backend.py
+37 -18
View File
@@ -377,6 +377,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
] ]
page_indices //= self.page_size page_indices //= self.page_size
metadata.page_table[:, :max_seq_pages].copy_(page_indices) 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(): elif forward_mode.is_draft_extend():
metadata = self.draft_extend_metadata[bs] metadata = self.draft_extend_metadata[bs]
metadata.cache_seqlens_int32.copy_(seq_lens) metadata.cache_seqlens_int32.copy_(seq_lens)
@@ -614,24 +615,42 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
bmm1_scale = q_scale * k_scale * layer.scaling bmm1_scale = q_scale * k_scale * layer.scaling
bmm2_scale = 1.0 bmm2_scale = 1.0
o = flashinfer.prefill.trtllm_batch_context_with_kv_cache( if forward_batch.forward_mode.is_target_verify():
query=q, o = flashinfer.decode.trtllm_batch_decode_with_kv_cache(
kv_cache=kv_cache, query=q,
workspace_buffer=self.workspace_buffer, kv_cache=kv_cache,
block_tables=self.forward_metadata.page_table, workspace_buffer=self.workspace_buffer,
seq_lens=self.forward_metadata.cache_seqlens_int32, block_tables=self.forward_metadata.page_table,
max_q_len=self.forward_metadata.max_seq_len_q, seq_lens=self.forward_metadata.cache_seqlens_int32,
max_kv_len=self.max_context_len, max_seq_len=self.max_context_len,
bmm1_scale=bmm1_scale, bmm1_scale=bmm1_scale,
bmm2_scale=bmm2_scale, bmm2_scale=bmm2_scale,
batch_size=forward_batch.batch_size, window_left=layer.sliding_window_size,
cum_seq_lens_q=self.forward_metadata.cu_seqlens_q, # TODO: add attention_sink operation or nvfp4 scale factor if needed
cum_seq_lens_kv=self.forward_metadata.cu_seqlens_k, sinks=attention_sink,
window_left=layer.sliding_window_size, out_dtype=self.q_data_type, # model_runner.dtype
# TODO: add attention_sink operation or nvfp4 scale factor if needed q_len_per_req=self.forward_metadata.max_seq_len_q,
sinks=attention_sink, )
out_dtype=self.q_data_type, # model_runner.dtype 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) return o.view(-1, layer.tp_q_head_num * layer.head_dim)