diff --git a/python/sglang/srt/layers/attention/nsa/triton_kernel.py b/python/sglang/srt/layers/attention/nsa/triton_kernel.py index 9d970b83a..0d2969804 100644 --- a/python/sglang/srt/layers/attention/nsa/triton_kernel.py +++ b/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, + ) diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 649ca9e50..637f0d7ca 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/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, - layer.logit_cap, + sm_scale=layer.scaling, + 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,