chenchenghao/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
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

55
python/sglang/srt/layers/attention/trtllm_mha_backend.py
View File

@@ -377,6 +377,7 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
]
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 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
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)