[DeepSeek V3.2] [Bugfix] slice indexer and padding fa3 when can not run cuda graph (#17076)

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(

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)
+    tokens = cal_padded_tokens(forward_batch)
-    pad_len = (tokens - 1) // attn_tp_size + 1 - nsa_cache_seqlens.shape[0]
+    pad_len = tokens - nsa_cache_seqlens.shape[0]
     if pad_len > 0:
         nsa_cache_seqlens = torch.cat(
             [

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
 

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 = []