diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 659236033..2fe14c368 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -86,6 +86,11 @@ class BaseIndexerMetadata(ABC): Return: seq lens for each batch. """ + def get_nsa_extend_len_cpu(self) -> List[int]: + """ + Return: extend seq lens for each batch. + """ + def get_token_to_batch_idx(self) -> torch.Tensor: """ Return: batch idx for each token. @@ -390,6 +395,9 @@ class Indexer(MultiPlatformOp): assert len(weights.shape) == 3 weights = weights.squeeze(2) + # When attn_tp_size > 1 or in the MAX_LEN padding mode, padding may exist in the hidden states, + # and it is necessary to extract the actual q length. + q_offset = sum(metadata.get_nsa_extend_len_cpu()) if _is_hip: from aiter.ops.triton.pa_mqa_logits import deepgemm_fp8_paged_mqa_logits @@ -416,9 +424,9 @@ class Indexer(MultiPlatformOp): ) else: logits = deep_gemm.fp8_paged_mqa_logits( - q_fp8, + q_fp8[:q_offset], kv_cache_fp8, - weights, + weights[:q_offset], seqlens_32, block_tables, schedule_metadata, @@ -428,6 +436,16 @@ class Indexer(MultiPlatformOp): # NOTE(dark): logits should be cleaned in topk_transform topk_result = metadata.topk_transform(logits, self.index_topk) + # Restore possible padding exist in the hidden states. + if not _is_hip and q_offset < q_fp8.shape[0]: + pad_len = q_fp8.shape[0] - q_offset + padding = torch.full( + (pad_len, topk_result.shape[1]), + -1, + dtype=topk_result.dtype, + device=topk_result.device, + ) + topk_result = torch.cat([topk_result, padding], dim=0) return topk_result def _should_chunk_mqa_logits( diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index b9a3c7c6b..e49f3ba01 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -9,12 +9,15 @@ import triton import triton.language as tl from sglang.srt.layers.dp_attention import ( + DpPaddingMode, attn_tp_all_gather_into_tensor, + get_attention_dp_rank, get_attention_tp_group, get_attention_tp_rank, get_attention_tp_size, ) from sglang.srt.server_args import get_global_server_args +from sglang.srt.utils.common import ceil_align, ceil_div if TYPE_CHECKING: from sglang.srt.model_executor.forward_batch_info import ForwardBatch @@ -81,12 +84,33 @@ def nsa_cp_round_robin_split_data(input_: Union[torch.Tensor, List]): return input_.view(-1, cp_size, *input_.shape[1:])[:, cp_rank].contiguous() -def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens): +def cal_padded_tokens(forward_batch: "ForwardBatch"): + # Consistent with the padding calculation logic in ForwardBatch.prepare_mlp_sync_batch, + # calculate the actual token length after padding when attn_tp_size > 1 or in the MAX_LEN padding mode. + global_num_tokens = forward_batch.global_num_tokens_cpu.copy() + sync_group_size = len(global_num_tokens) attn_tp_size = get_attention_tp_size() - if attn_tp_size == 1 or not can_nsa_prefill_cp_round_robin_split(forward_batch): + for i in range(sync_group_size): + global_num_tokens[i] = ceil_align(global_num_tokens[i], attn_tp_size) + dp_padding_mode = DpPaddingMode.get_dp_padding_mode( + forward_batch.is_extend_in_batch, global_num_tokens + ) + if dp_padding_mode.is_max_len(): + tokens = max(global_num_tokens) + elif len(global_num_tokens) > 1: + tokens = global_num_tokens[get_attention_dp_rank()] + else: + tokens = global_num_tokens[0] + if can_nsa_prefill_cp_round_robin_split(forward_batch): + tokens = ceil_div(tokens, attn_tp_size) + return tokens + + +def pad_nsa_cache_seqlens(forward_batch: "ForwardBatch", nsa_cache_seqlens): + if forward_batch.global_num_tokens_cpu is None: return nsa_cache_seqlens - tokens = sum(forward_batch.extend_seq_lens_cpu) - pad_len = (tokens - 1) // attn_tp_size + 1 - nsa_cache_seqlens.shape[0] + tokens = cal_padded_tokens(forward_batch) + pad_len = tokens - nsa_cache_seqlens.shape[0] if pad_len > 0: nsa_cache_seqlens = torch.cat( [ diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 0e95c28d2..b6bb7d08b 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -189,6 +189,9 @@ class NSAIndexerMetadata(BaseIndexerMetadata): def get_indexer_seq_len_cpu(self) -> torch.Tensor: return self.attn_metadata.indexer_seq_lens_cpu + def get_nsa_extend_len_cpu(self) -> List[int]: + return self.attn_metadata.nsa_extend_seq_lens_list + def get_token_to_batch_idx(self) -> torch.Tensor: return self.attn_metadata.token_to_batch_idx diff --git a/test/registered/kernels/test_nsa_indexer.py b/test/registered/kernels/test_nsa_indexer.py index 20488ecd6..07504a684 100644 --- a/test/registered/kernels/test_nsa_indexer.py +++ b/test/registered/kernels/test_nsa_indexer.py @@ -1,5 +1,5 @@ import unittest -from typing import Optional, Tuple +from typing import List, Optional, Tuple from unittest.mock import MagicMock, patch import torch @@ -132,6 +132,12 @@ class MockIndexerMetadata(BaseIndexerMetadata): """Return: seq lens for each batch.""" return torch.tensor(self.seq_lens, dtype=torch.int32, device="cpu") + def get_nsa_extend_len_cpu(self) -> List[int]: + """ + Return: extend seq lens for each batch. + """ + return list(self.seq_lens) + def get_token_to_batch_idx(self) -> torch.Tensor: """Return: batch idx for each token.""" result = []