diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index f18c57070..33345593f 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -628,7 +628,7 @@ class AscendAttnBackend(AttentionBackend): if self.forward_metadata.actual_seq_lengths_q is not None: actual_seq_qlen = self.forward_metadata.actual_seq_lengths_q else: - actual_seq_qlen = torch.cumsum(forward_batch.seq_lens, dim=0) + actual_seq_qlen = torch.cumsum(forward_batch.extend_seq_lens, dim=0) else: if self.forward_metadata.actual_seq_lengths_q is None: if ( diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index 4968c03e8..82e7ce693 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -221,6 +221,7 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): dtype=self.store_dtype, device=self.device, ) + self.index_k_buffer = None if self.index_head_dim is not None: self.index_k_buffer = torch.zeros( ( diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 2fe14c368..cf362e712 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1244,7 +1244,7 @@ class Indexer(MultiPlatformOp): ) else: actual_seq_lengths_kv = forward_batch.seq_lens - actual_seq_lengths_q = forward_batch.seq_lens.cumsum(dim=0) + actual_seq_lengths_q = forward_batch.extend_seq_lens.cumsum(dim=0) else: if forward_batch.attn_backend.forward_metadata.actual_seq_lengths_q is None: if ( diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py index 548c9d9f1..f40ddb17d 100644 --- a/python/sglang/srt/mem_cache/memory_pool_host.py +++ b/python/sglang/srt/mem_cache/memory_pool_host.py @@ -769,6 +769,15 @@ class MLATokenToKVPoolHost(HostKVCache): pin_memory=self.pin_memory, allocator=self.allocator, ) + self.index_k_buffer = None + if self.device_pool.index_head_dim is not None: + self.index_k_buffer = alloc_func( + (*base_dims, self.device_pool.index_head_dim), + dtype=self.dtype, + device=self.device, + pin_memory=self.pin_memory, + allocator=self.allocator, + ) # Return k_buffer to preserve original kv_buffer and data_refs init logic, # though Ascend doesn't use these parameters. return self.k_buffer @@ -844,6 +853,8 @@ class MLATokenToKVPoolHost(HostKVCache): host_k=self.k_buffer, device_v=device_pool.v_buffer, host_v=self.v_buffer, + device_index_k=device_pool.index_k_buffer, + host_index_k=self.index_k_buffer, page_size=self.page_size, direction=TransferDirection.H2D, ) @@ -905,6 +916,8 @@ class MLATokenToKVPoolHost(HostKVCache): host_k=self.k_buffer, device_v=device_pool.v_buffer, host_v=self.v_buffer, + device_index_k=device_pool.index_k_buffer, + host_index_k=self.index_k_buffer, page_size=self.page_size, direction=TransferDirection.D2H, )