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

[AMD] Fix eager metadata for AITER EAGLE draft extend (#36915)

Browse Source 
Co-authored-by: jacky.cheng <yichiche@amd.com>
This commit is contained in:
zijiec
2026-08-28 22:56:38 -07:00
committed by GitHub
co-authored by jacky.cheng
parent 24c9251ac5
commit 7f2ee22b70
1 changed files with 13 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/aiter_backend.py
+10 -10
View File
@@ -1270,21 +1270,21 @@ class AiterAttnBackend(AttentionBackend):
run_graph=False, run_graph=False,
) )
else: else:
self.indices_updater_prefill.update( kv_indices, kv_indptr, qo_indptr, _ = (
forward_batch.spec_info.generate_attn_arg_prefill(
forward_batch.req_pool_indices, forward_batch.req_pool_indices,
forward_batch.seq_lens, forward_batch.seq_lens,
forward_batch.seq_lens_sum, forward_batch.seq_lens_sum,
prefix_lens=None, self.req_to_token,
encoder_lens=forward_batch.encoder_lens, )
spec_info=forward_batch.spec_info,
) )
self.forward_metadata = ForwardMetadata( self.forward_metadata = ForwardMetadata(
self.indices_updater_prefill.kv_indptr, kv_indptr,
self.indices_updater_prefill.kv_indices, kv_indices,
qo_indptr,
None, None,
None, forward_batch.spec_info.num_tokens_per_req,
self.indices_updater_prefill.max_q_len, max_kv_len,
self.indices_updater_prefill.max_kv_len,
) )
elif forward_batch.forward_mode.is_target_verify(): elif forward_batch.forward_mode.is_target_verify():
if self.use_mla: if self.use_mla:
@@ -2592,7 +2592,7 @@ class AiterAttnBackend(AttentionBackend):
-1, self.page_size, layer.tp_v_head_num, layer.v_head_dim -1, self.page_size, layer.tp_v_head_num, layer.v_head_dim
), ),
out=o.view(-1, layer.tp_q_head_num, layer.v_head_dim), out=o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
cu_seqlens_q=self.qo_indptr[: bs + 1], cu_seqlens_q=self.forward_metadata.qo_indptr,
seqused_k=seqused_k, seqused_k=seqused_k,
max_seqlen_q=self.forward_metadata.max_q_len, max_seqlen_q=self.forward_metadata.max_q_len,
max_seqlen_k=pt.shape[1] * self.page_size, max_seqlen_k=pt.shape[1] * self.page_size,