chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

[Deepseek V3.2] Fix accuracy bug in the Indexer (#12583)

Browse Source 
Signed-off-by: Hao Lu <14827759+hlu1@users.noreply.github.com>
This commit is contained in:
hlu1
2025-11-11 16:15:26 -08:00
committed by GitHub
parent 2ca25a8aab
commit 0d4a418424
6 changed files with 96 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/pr-test.yml
+1 -1
View File
@@ -868,7 +868,7 @@ jobs:
CUSTOM_BUILD_SGL_KERNEL=${{needs.check-changes.outputs.sgl_kernel}} bash scripts/ci/ci_install_dependency.sh
- name: Run test
timeout-minutes: 20
timeout-minutes: 30
run: |
cd test/srt
python3 run_suite.py --suite per-commit-8-gpu-h200-deepseek-v32
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
+32 -8
View File
@@ -345,7 +345,10 @@ class Indexer(CustomOp):
k_scale_list = []
ks_list = []
ke_list = []
offset = 0
q_offset = 0
k_offset = 0
seq_lens_expanded = metadata.get_seqlens_expanded()
block_tables = metadata.get_page_table_64()
@@ -368,13 +371,17 @@ class Indexer(CustomOp):
block_tables[i],
)
extend_seq_len = forward_batch.extend_seq_lens_cpu[i]
ks = torch.full((extend_seq_len,), offset, dtype=torch.int32, device="cuda")
ke = ks + seq_lens_expanded[offset : offset + extend_seq_len]
ks = torch.full(
(extend_seq_len,), k_offset, dtype=torch.int32, device="cuda"
)
ke = ks + seq_lens_expanded[q_offset : q_offset + extend_seq_len]
k_fp8_list.append(k_fp8)
k_scale_list.append(k_scale)
ks_list.append(ks)
ke_list.append(ke)
offset += extend_seq_len
q_offset += extend_seq_len
k_offset += seq_len
k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn)
k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1)
@@ -382,21 +389,38 @@ class Indexer(CustomOp):
ks = torch.cat(ks_list, dim=0)
ke = torch.cat(ke_list, dim=0)
# Suppose there are two requests, with extend_seq_len = [3, 2]
# and seq_lens = [10, 4]
# The logits matrix looks like this, with * representing the valid logits
# and - representing the invalid logits:
#
# ********--|----
# *********-|----
# **********|----
# ----------|***-
# ----------|****
#
# ks = [0, 0, 0, 10, 10]
# ke = [8, 9, 10, 13, 14]
logits = deep_gemm.fp8_mqa_logits(
q_fp8[:offset],
q_fp8[:q_offset],
kv_fp8,
weights[:offset],
weights[:q_offset],
ks,
ke,
clean_logits=False,
)
token_nums, _, _ = q_fp8.shape
assert logits.shape[0] == len(seq_lens_expanded)
raw_topk_result = metadata.topk_transform(logits, self.index_topk)
assert logits.shape[1] == k_offset
raw_topk_result = metadata.topk_transform(logits, self.index_topk, ks=ks)
topk_result = torch.full(
(token_nums, self.index_topk), -1, device=q_fp8.device, dtype=torch.int32
)
topk_result[:offset] = raw_topk_result
topk_result[:q_offset] = raw_topk_result
return topk_result
def _forward_cuda_k_only(
python/sglang/srt/layers/attention/nsa_backend.py
+6 -1
View File
@@ -140,6 +140,7 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
self,
logits: torch.Tensor,
topk: int,
ks: Optional[torch.Tensor] = None,
) -> torch.Tensor:
from sgl_kernel import (
fast_topk_transform_fused,
@@ -148,7 +149,9 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
)
if not NSA_FUSE_TOPK:
return fast_topk_v2(logits, self.get_seqlens_expanded(), topk)
return fast_topk_v2(
logits, self.get_seqlens_expanded(), topk, row_starts=ks
)
elif self.topk_transform_method == TopkTransformMethod.PAGED:
# NOTE(dark): if fused, we return a transformed page table directly
return fast_topk_transform_fused(
@@ -157,6 +160,7 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
page_table_size_1=self.attn_metadata.page_table_1,
cu_seqlens_q=self.attn_metadata.cu_seqlens_q,
topk=topk,
row_starts=ks,
)
elif self.topk_transform_method == TopkTransformMethod.RAGGED:
return fast_topk_transform_ragged_fused(
@@ -164,6 +168,7 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
lengths=self.get_seqlens_expanded(),
topk_indices_offset=self.attn_metadata.topk_indices_offset,
topk=topk,
row_starts=ks,
)
else:
assert False, f"Unsupported {self.topk_transform_method = }"
test/srt/test_deepseek_v32_basic.py
+1 -1
View File
@@ -44,7 +44,7 @@ class TestDeepseekV32Basic(CustomTestCase):
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=8,
num_shots=20,
data_path=None,
num_questions=1400,
parallel=1400,
test/srt/test_deepseek_v32_mtp.py
+4 -4
View File
@@ -15,13 +15,13 @@ from sglang.test.test_utils import (
write_github_step_summary,
)
FULL_DEEPSEEK_V3_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
FULL_DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2-Exp"
class TestDeepseekV32MTP(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = FULL_DEEPSEEK_V3_MODEL_PATH
cls.model = FULL_DEEPSEEK_V32_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
@@ -82,7 +82,7 @@ class TestDeepseekV32MTP(CustomTestCase):
f"{avg_spec_accept_length=:.2f}\n"
)
self.assertGreater(metrics["accuracy"], 0.935)
self.assertGreater(avg_spec_accept_length, 2.9)
self.assertGreater(avg_spec_accept_length, 2.7)
def test_bs_1_speed(self):
args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048)
@@ -97,7 +97,7 @@ class TestDeepseekV32MTP(CustomTestCase):
f"{speed=:.2f} token/s\n"
)
self.assertGreater(acc_length, 2.9)
self.assertGreater(acc_length, 2.7)
self.assertGreater(speed, 75)
test/srt/test_deepseek_v32_nsabackend.py
+52 -2
View File
@@ -50,7 +50,7 @@ class TestDeepseekV32NasBackend_flashmla(CustomTestCase):
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=8,
num_shots=20,
data_path=None,
num_questions=1400,
parallel=1400,
@@ -102,7 +102,57 @@ class TestDeepseekV32NasBackend_fa3(CustomTestCase):
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=8,
num_shots=20,
data_path=None,
num_questions=1400,
parallel=1400,
max_new_tokens=512,
host="http://127.0.0.1",
port=int(self.base_url.split(":")[-1]),
)
metrics = run_eval_few_shot_gsm8k(args)
print(f"{metrics=}")
if is_in_ci():
write_github_step_summary(
f"### test_gsm8k (deepseek-v3)\n" f'{metrics["accuracy"]=:.3f}\n'
)
self.assertGreater(metrics["accuracy"], 0.935)
class TestDeepseekV32NasBackend_fp8kvcache(CustomTestCase):
@classmethod
def setUpClass(cls):
cls.model = DEEPSEEK_V32_MODEL_PATH
cls.base_url = DEFAULT_URL_FOR_TEST
other_args = [
"--trust-remote-code",
"--attention-backend",
"nsa",
"--kv-cache-dtype",
"fp8_e4m3",
"--tp",
"8",
"--dp",
"8",
"--enable-dp-attention",
]
cls.process = popen_launch_server(
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=other_args,
)
@classmethod
def tearDownClass(cls):
kill_process_tree(cls.process.pid)
def test_a_gsm8k(
self,
): # Append an "a" to make this test run first (alphabetically) to warm up the server
args = SimpleNamespace(
num_shots=20,
data_path=None,
num_questions=1400,
parallel=1400,