[AMD] Enable cudagraph for aiter nsa backend and add aiter impl for nsa pr… (#18526)

python/sglang/srt/layers/attention/nsa/triton_kernel.py

@@ -134,3 +134,63 @@ def act_quant(
     )
 
     return y, s
+
+
+@triton.jit
+def _get_valid_kv_indices_kernel(
+    page_table_ptr,  # [bs, topk]
+    kv_indptr_ptr,  # [bs + 1]
+    kv_indices_ptr,  # [bs * topk] output buffer
+    bs: tl.constexpr,
+    topk: tl.constexpr,
+):
+    """
+    Extract valid indices (non -1) from page_table into kv_indices.
+    Each program handles one batch.
+    """
+    batch_id = tl.program_id(0)
+
+    # Get the start position for this batch in kv_indices
+    dst_start = tl.load(kv_indptr_ptr + batch_id)
+
+    # Load all topk indices for this batch
+    src_offset = batch_id * topk
+    offsets = tl.arange(0, topk)
+    indices = tl.load(page_table_ptr + src_offset + offsets)
+
+    # Count valid indices and compact them
+    mask = indices != -1
+
+    # Use prefix sum to compute destination positions for valid elements
+    # For each position, count how many valid elements are before it
+    prefix_sum = tl.cumsum(mask.to(tl.int32), axis=0) - 1
+
+    # Store valid indices to their compacted positions
+    dst_positions = dst_start + prefix_sum
+    tl.store(kv_indices_ptr + dst_positions, indices, mask=mask)
+
+
+def get_valid_kv_indices(
+    page_table_1: torch.Tensor,
+    kv_indptr: torch.Tensor,
+    kv_indices: torch.Tensor,
+    bs: int,
+):
+    """
+    Extract valid indices from page_table_1 into kv_indices buffer.
+
+    Args:
+        page_table_1: [bs, topk] page table with -1 as invalid
+        kv_indptr: [bs + 1] cumulative count of valid indices per batch
+        kv_indices: [bs * topk] pre-allocated output buffer
+        bs: batch size
+    """
+    topk = page_table_1.shape[1]
+    grid = (bs,)
+    _get_valid_kv_indices_kernel[grid](
+        page_table_1,
+        kv_indptr,
+        kv_indices,
+        bs,
+        topk,
+    )

python/sglang/srt/layers/attention/nsa_backend.py

@@ -52,6 +52,8 @@ if TYPE_CHECKING:
 _is_hip = is_hip()
 
 if _is_hip:
+    from sglang.srt.layers.attention.nsa.triton_kernel import get_valid_kv_indices
+
     try:
         from aiter import (  # noqa: F401
             flash_attn_varlen_func,
@@ -333,6 +335,12 @@ class NativeSparseAttnBackend(
                 (max_bs + 1,), dtype=torch.int32, device=model_runner.device
             )
 
+            self.kv_indices = torch.zeros(
+                max_bs * self.nsa_index_topk,
+                dtype=torch.int32,
+                device=self.device,
+            )
+
         # Speculative decoding
         self.topk = model_runner.server_args.speculative_eagle_topk or 0
         self.speculative_num_steps = speculative_num_steps
@@ -1476,6 +1484,15 @@ class NativeSparseAttnBackend(
                 logit_cap=layer.logit_cap,
                 page_size=1,
             )
+        elif nsa_impl == "aiter":
+            if q_rope is not None:
+                q_all = torch.cat([q_nope, q_rope], dim=-1)
+            return self._forward_aiter_extend(
+                q_all=q_all,
+                kv_cache=kv_cache,
+                page_table_1=page_table_1,
+                layer=layer,
+            )
         else:
             raise ValueError(
                 f"Unsupported {nsa_impl = } for forward_extend. Consider using an other attention backend."
@@ -1861,7 +1878,8 @@ class NativeSparseAttnBackend(
         non_minus1_counts = non_minus1_mask.sum(dim=1)
         kv_indptr[1 : bs + 1] = torch.cumsum(non_minus1_counts, dim=0)
 
-        kv_indices = page_table_1[page_table_1 != -1]
+        kv_indices = self.kv_indices
+        get_valid_kv_indices(page_table_1, kv_indptr, kv_indices, bs)
 
         mla_decode_fwd(
             q.view(-1, layer.tp_q_head_num, layer.head_dim),
@@ -1872,12 +1890,61 @@ class NativeSparseAttnBackend(
             kv_indices,
             metadata.cu_seqlens_q,
             metadata.max_seq_len_q,
-            layer.scaling,
+            sm_scale=layer.scaling,
-            layer.logit_cap,
+            logit_cap=layer.logit_cap,
         )
         # kv_cache = kv_cache.view(-1, 1, layer.head_dim)
         return o
 
+    def _forward_aiter_extend(
+        self,
+        q_all: torch.Tensor,
+        kv_cache: torch.Tensor,
+        page_table_1: torch.Tensor,
+        layer: RadixAttention,
+    ) -> torch.Tensor:
+        num_tokens = q_all.shape[0]
+        q = q_all.reshape(-1, layer.tp_q_head_num * layer.head_dim)
+
+        if layer.head_dim != layer.v_head_dim:
+            o = q.new_empty((num_tokens, layer.tp_q_head_num * layer.v_head_dim))
+        else:
+            o = torch.empty_like(q)
+
+        non_minus1_mask = page_table_1 != -1
+        non_minus1_counts = non_minus1_mask.sum(dim=1)
+
+        kv_indptr = torch.zeros(num_tokens + 1, dtype=torch.int32, device=self.device)
+        kv_indptr[1:] = torch.cumsum(non_minus1_counts, dim=0)
+
+        # Allocate kv_indices with upper-bound size (num_tokens * topk)
+        topk = page_table_1.shape[1]
+        kv_indices = torch.zeros(
+            num_tokens * topk, dtype=torch.int32, device=self.device
+        )
+
+        # Use get_valid_kv_indices kernel to extract valid indices
+        get_valid_kv_indices(page_table_1, kv_indptr, kv_indices, num_tokens)
+
+        # Build cu_seqlens_q for extend: each token is treated as seq_len_q=1
+        cu_seqlens_q = torch.arange(
+            0, num_tokens + 1, dtype=torch.int32, device=self.device
+        )
+        # TODO support more forward_mode
+        mla_decode_fwd(
+            q.view(-1, layer.tp_q_head_num, layer.head_dim),
+            kv_cache.view(-1, 1, 1, layer.head_dim),
+            o.view(-1, layer.tp_q_head_num, layer.v_head_dim),
+            cu_seqlens_q,
+            kv_indptr,
+            kv_indices,
+            cu_seqlens_q,
+            1,  # max_seq_len_q = 1 for per-token attention
+            sm_scale=layer.scaling,
+            logit_cap=layer.logit_cap,
+        )
+        return o
+
     def _forward_trtllm(
         self,
         q: torch.Tensor,