diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 0d557694d..3b02fad76 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -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 diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index ae9ff38a2..ac086238a 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -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( diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index b77fb1a1a..0c8077427 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -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 = }" diff --git a/test/srt/test_deepseek_v32_basic.py b/test/srt/test_deepseek_v32_basic.py index de56d9585..ce2d5afae 100644 --- a/test/srt/test_deepseek_v32_basic.py +++ b/test/srt/test_deepseek_v32_basic.py @@ -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, diff --git a/test/srt/test_deepseek_v32_mtp.py b/test/srt/test_deepseek_v32_mtp.py index 1681b56ce..3fbc523df 100644 --- a/test/srt/test_deepseek_v32_mtp.py +++ b/test/srt/test_deepseek_v32_mtp.py @@ -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) diff --git a/test/srt/test_deepseek_v32_nsabackend.py b/test/srt/test_deepseek_v32_nsabackend.py index 1f9975f45..cd7702f28 100644 --- a/test/srt/test_deepseek_v32_nsabackend.py +++ b/test/srt/test_deepseek_v32_nsabackend.py @@ -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,