minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Refactor] Rename NSA → DSA: user-facing aliases, file/class/import rename (#25821)

Browse Source 
Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-05-20 00:18:04 -07:00
committed by GitHub
co-authored by Claude Sonnet 4.6
parent da6d549ab2
commit 8131641bc6
162 changed files with 11298 additions and 10740 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/layers/attention/nsa/test_act_quant_triton.py → test/manual/layers/attention/dsa/test_act_quant_triton.py
+2 -2
View File
@@ -9,8 +9,8 @@ from typing import Tuple
import pytest
import torch
from sglang.srt.layers.attention.nsa.tilelang_kernel import act_quant
from sglang.srt.layers.attention.nsa.triton_kernel import act_quant as act_quant_triton
from sglang.srt.layers.attention.dsa.tilelang_kernel import act_quant
from sglang.srt.layers.attention.dsa.triton_kernel import act_quant as act_quant_triton
def benchmark_kernel(
test/manual/layers/attention/nsa/test_get_k_scale_triton_kernel.py → test/manual/layers/attention/dsa/test_get_k_scale_triton_kernel.py
+1 -1
View File
@@ -1,6 +1,6 @@
import torch
from sglang.srt.layers.attention.nsa.index_buf_accessor import (
from sglang.srt.layers.attention.dsa.index_buf_accessor import (
_get_k_and_s_triton_kernel,
)
test/manual/layers/attention/nsa/test_index_buf_accessor.py → test/manual/layers/attention/dsa/test_index_buf_accessor.py
+18 -18
View File
@@ -1,5 +1,5 @@
"""
Correctness tests for NSA Indexer K/S Buffer Access with Fused Triton Kernels.
Correctness tests for DSA Indexer K/S Buffer Access with Fused Triton Kernels.
This test verifies that the optimized Triton implementations (GetK, GetS, GetKAndS)
produce identical results to the torch_fast baseline implementations.
@@ -13,11 +13,11 @@ Test coverage:
import pytest
import torch
from sglang.srt.layers.attention.nsa.index_buf_accessor import GetK, GetKAndS, GetS
from sglang.srt.layers.attention.dsa.index_buf_accessor import GetK, GetKAndS, GetS
class MockNSATokenToKVPool:
"""Mock pool object that mimics NSATokenToKVPool for testing."""
class MockDSATokenToKVPool:
"""Mock pool object that mimics DSATokenToKVPool for testing."""
def __init__(
self,
@@ -78,7 +78,7 @@ class TestGetK:
seq_len = min(seq_len, max_seq_len)
# Create mock pool
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
@@ -119,7 +119,7 @@ class TestGetK:
num_pages = 10
seq_len = 320 # 5 pages
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
@@ -140,7 +140,7 @@ class TestGetK:
num_pages = 5
seq_len = 192 # 3 pages
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
@@ -171,7 +171,7 @@ class TestGetS:
seq_len = min(seq_len, max_seq_len)
# Create mock pool
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
@@ -212,7 +212,7 @@ class TestGetS:
num_pages = 10
seq_len = 320 # 5 pages
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
@@ -233,7 +233,7 @@ class TestGetS:
num_pages = 5
seq_len = 192 # 3 pages
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
@@ -267,7 +267,7 @@ class TestGetKAndS:
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
# Create mock pool
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
@@ -326,7 +326,7 @@ class TestGetKAndS:
seq_len = 320 # 5 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
@@ -356,7 +356,7 @@ class TestGetKAndS:
seq_len = 192 # 3 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
@@ -386,7 +386,7 @@ class TestGetKAndS:
seq_len = 100 # Not a multiple of 64
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
@@ -422,7 +422,7 @@ class TestEdgeCases:
seq_len = 1
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
@@ -455,7 +455,7 @@ class TestEdgeCases:
seq_len = 192 # Exactly 3 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
@@ -488,7 +488,7 @@ class TestEdgeCases:
seq_len = 4096 # 64 pages
seq_len_tensor = torch.tensor([seq_len], dtype=torch.int64, device=device)
pool = MockNSATokenToKVPool(
pool = MockDSATokenToKVPool(
page_size=page_size, index_head_dim=index_head_dim, device=device
)
buf = create_test_buffer(num_pages, page_size, index_head_dim, device)
@@ -520,7 +520,7 @@ class TestEdgeCases:
def print_test_summary():
"""Print a summary message about the test suite."""
print("\n" + "=" * 80)
print("NSA Indexer K/S Buffer Accessor Correctness Tests")
print("DSA Indexer K/S Buffer Accessor Correctness Tests")
print("=" * 80)
print("Testing Triton implementations against torch_fast baseline:")
print(" - GetK.triton() vs GetK.torch_fast()")