diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index a213cc5ab..ec492bb4d 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -857,6 +857,11 @@ class AiterAttnBackend(AttentionBackend): seq_lens_cpu = ( forward_batch.seq_lens.cpu() if in_capture else forward_batch.seq_lens_cpu ) + verify_tokens_per_req = ( + forward_batch.input_ids.shape[0] // forward_batch.batch_size + if forward_batch.forward_mode.is_target_verify() + else None + ) self._apply_cuda_graph_metadata( bs=forward_batch.batch_size, req_pool_indices=forward_batch.req_pool_indices, @@ -866,6 +871,7 @@ class AiterAttnBackend(AttentionBackend): forward_mode=forward_batch.forward_mode, spec_info=forward_batch.spec_info, seq_lens_cpu=seq_lens_cpu, + verify_tokens_per_req=verify_tokens_per_req, ) # Refill the SWA write-target buffer from the live out_cache_loc and @@ -1196,8 +1202,8 @@ class AiterAttnBackend(AttentionBackend): run_graph=False, ) else: + draft_num = forward_batch.input_ids.shape[0] // bs bs = len(forward_batch.req_pool_indices) - draft_num = spec_info.draft_token_num if self._use_unified_verify: page_table, qo_indptr, max_q_len, swa_page_table = ( @@ -1500,6 +1506,7 @@ class AiterAttnBackend(AttentionBackend): forward_mode: ForwardMode, spec_info: Optional[SpecInput], seq_lens_cpu: Optional[torch.Tensor], + verify_tokens_per_req: Optional[int], ): num_kv_splits = None @@ -1652,11 +1659,17 @@ class AiterAttnBackend(AttentionBackend): elif forward_mode.is_target_verify(): bs = len(req_pool_indices) + assert verify_tokens_per_req is not None + # MLA uses a fixed draft length (num_draft_tokens); the non-MLA + # unified path derives it per batch from input_ids. + tokens_per_req = ( + self.num_draft_tokens if self.use_mla else verify_tokens_per_req + ) qo_indptr = self.qo_indptr[: bs + 1] qo_indptr[: bs + 1] = torch.arange( 0, - (1 + bs) * self.num_draft_tokens, - step=self.num_draft_tokens, + (1 + bs) * tokens_per_req, + step=tokens_per_req, dtype=torch.int32, device=self.device, ) @@ -1689,9 +1702,9 @@ class AiterAttnBackend(AttentionBackend): self.req_to_token.stride(0), ) kv_last_page_len = self.cuda_graph_kv_last_page_len[:bs] - max_q_len = self.num_draft_tokens if self.use_mla: + max_q_len = self.num_draft_tokens if _use_mla_ps_kernel: num_kv_splits = self.max_split_per_batch @@ -1735,6 +1748,7 @@ class AiterAttnBackend(AttentionBackend): num_kv_splits=num_kv_splits, ) else: + max_q_len = verify_tokens_per_req if self._use_unified_verify: max_num_blocks_per_seq = ( self.max_context_len + self.page_size - 1 @@ -1753,7 +1767,7 @@ class AiterAttnBackend(AttentionBackend): bs, seq_lens, req_pool_indices, - self.num_draft_tokens, + verify_tokens_per_req, page_table_dest=page_table, swa_page_table_dest=swa_page_table, ) @@ -2278,7 +2292,9 @@ class AiterAttnBackend(AttentionBackend): v=v_unified, out=o.view(-1, layer.tp_q_head_num, layer.v_head_dim), cu_seqlens_q=self.forward_metadata.qo_indptr, - seqused_k=forward_batch.seq_lens + self.num_draft_tokens, + seqused_k=( + forward_batch.seq_lens + self.forward_metadata.max_q_len + ), max_seqlen_q=self.forward_metadata.max_q_len, max_seqlen_k=max_kv_len, softmax_scale=layer.scaling,