1b99de7459
The bs>1 prefill path has multiple coupled stages: scheduler admission, page-aligned batch planning, tensor splitting, direct cache writes, index top-k, MLA reuse, and disaggregated KV handoff. Add a default-off, rate-limited debug channel so production ETE runs can identify where batching or metadata semantics diverge without permanently increasing hot-path log volume. Constraint: Logs must be default-off and rate-limited because these paths execute per-rank and often per-layer. Rejected: Always-on INFO logs | would flood logs and add CPU overhead during normal prefill. Rejected: Only scheduler-side logging | insufficient to distinguish planner, index, MLA, and transfer handoff failures. Confidence: medium Scope-risk: moderate Directive: Keep bs>1 debug evidence env-gated; do not add unconditional per-layer or per-token logs in these paths. Tested: Local py_compile for touched files Tested: git diff --check Tested: Remote py_compile and targeted NSA CP utility tests: 5 passed Not-tested: Full ETE correctness with debug disabled
3938 lines
142 KiB
Python
3938 lines
142 KiB
Python
import ast
|
|
import os
|
|
from pathlib import Path
|
|
import unittest
|
|
import sys
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from sglang.srt.environ import envs
|
|
from sglang.srt.layers.attention.nsa import utils as nsa_utils
|
|
from sglang.srt.layers.attention.nsa.utils import (
|
|
NSAContextParallelMetadata,
|
|
PageAlignedCacheExtent,
|
|
build_flat_page_owner_plan,
|
|
build_batch_page_aligned_in_seq_split_plan,
|
|
build_page_aligned_cache_extent,
|
|
_get_in_seq_last_token_owner_and_offset,
|
|
_build_batch_metadata_from_plan,
|
|
build_page_aligned_in_seq_split_list,
|
|
build_token_balanced_in_seq_split_list,
|
|
can_cp_split,
|
|
cp_all_gather_rerange_output,
|
|
cp_collect_last_token_hidden,
|
|
cp_split_and_rebuild_1d,
|
|
cp_split_and_rebuild_data,
|
|
cp_split_and_rebuild_position,
|
|
_torch_batch_in_seq_all_gather_rerange,
|
|
get_cp_shared_kv_batch_plan,
|
|
get_cp_shared_kv_local_out_cache_loc,
|
|
get_cp_shared_kv_local_physical_out_cache_loc,
|
|
get_cp_local_embedding_padded_token_count,
|
|
nsa_use_prefill_cp,
|
|
pad_cp_local_input_ids_for_embedding,
|
|
prepare_input_dp_with_cp_dsa,
|
|
select_cp_local_valid_rows_for_cache_write,
|
|
split_tensor_by_cp_batch_plan,
|
|
split_in_seq_cp_local_pair,
|
|
)
|
|
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
|
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
|
|
from sglang.srt.models.deepseek_nextn import DeepseekModelNextN
|
|
from sglang.srt.speculative.eagle_info import EagleDraftInput
|
|
from sglang.test.ci.ci_register import register_cpu_ci
|
|
|
|
register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
|
|
|
|
|
|
class TestPageAlignedCacheExtent(unittest.TestCase):
|
|
def test_extent_uses_page_boundary_not_cp_size(self):
|
|
extent = build_page_aligned_cache_extent(valid_tokens=100, page_size=64)
|
|
|
|
self.assertEqual(extent.valid_tokens, 100)
|
|
self.assertEqual(extent.padded_pages, 2)
|
|
self.assertEqual(extent.padded_tokens, 128)
|
|
self.assertEqual(extent.padding_tokens, 28)
|
|
|
|
def test_extent_handles_empty_and_aligned_lengths(self):
|
|
self.assertEqual(
|
|
build_page_aligned_cache_extent(valid_tokens=0, page_size=64),
|
|
PageAlignedCacheExtent(
|
|
valid_tokens=0,
|
|
padded_pages=0,
|
|
padded_tokens=0,
|
|
padding_tokens=0,
|
|
),
|
|
)
|
|
self.assertEqual(
|
|
build_page_aligned_cache_extent(valid_tokens=128, page_size=64),
|
|
PageAlignedCacheExtent(
|
|
valid_tokens=128,
|
|
padded_pages=2,
|
|
padded_tokens=128,
|
|
padding_tokens=0,
|
|
),
|
|
)
|
|
|
|
|
|
class TestNSAInSeqCPUtils(unittest.TestCase):
|
|
def test_contiguous_valid_cp_query_count(self):
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import (
|
|
_compute_contiguous_valid_cp_query_count,
|
|
)
|
|
|
|
self.assertEqual(
|
|
_compute_contiguous_valid_cp_query_count(
|
|
cp_kv_end=1024,
|
|
actual_seq_q=128,
|
|
logical_kv_limit=1024,
|
|
),
|
|
128,
|
|
)
|
|
self.assertEqual(
|
|
_compute_contiguous_valid_cp_query_count(
|
|
cp_kv_end=1100,
|
|
actual_seq_q=128,
|
|
logical_kv_limit=1024,
|
|
),
|
|
52,
|
|
)
|
|
self.assertEqual(
|
|
_compute_contiguous_valid_cp_query_count(
|
|
cp_kv_end=1100,
|
|
actual_seq_q=64,
|
|
logical_kv_limit=1000,
|
|
),
|
|
0,
|
|
)
|
|
self.assertEqual(
|
|
_compute_contiguous_valid_cp_query_count(
|
|
cp_kv_end=100,
|
|
actual_seq_q=0,
|
|
logical_kv_limit=100,
|
|
),
|
|
0,
|
|
)
|
|
|
|
def assert_page_aligned_boundaries(
|
|
self, split_list, *, extend_prefix_len, extend_len, page_size
|
|
):
|
|
cursor = 0
|
|
for segment_len in split_list[:-1]:
|
|
cursor += segment_len
|
|
if cursor < extend_len:
|
|
self.assertEqual((extend_prefix_len + cursor) % page_size, 0)
|
|
|
|
def test_page_aligned_split_keeps_boundaries_on_pages(self):
|
|
split_list, info = build_page_aligned_in_seq_split_list(
|
|
total_len=32768,
|
|
extend_len=32768,
|
|
extend_prefix_len=0,
|
|
page_size=64,
|
|
cp_size=8,
|
|
)
|
|
|
|
self.assertTrue(info.page_aligned)
|
|
self.assertEqual(sum(split_list), 32768)
|
|
self.assertEqual(len(split_list), 16)
|
|
self.assertTrue(all(segment_len > 0 for segment_len in split_list))
|
|
self.assert_page_aligned_boundaries(
|
|
split_list, extend_prefix_len=0, extend_len=32768, page_size=64
|
|
)
|
|
|
|
def test_page_aligned_split_uses_prefix_for_boundary_alignment(self):
|
|
split_list, info = build_page_aligned_in_seq_split_list(
|
|
total_len=1024,
|
|
extend_len=1024,
|
|
extend_prefix_len=128,
|
|
page_size=64,
|
|
cp_size=8,
|
|
)
|
|
|
|
self.assertTrue(info.page_aligned)
|
|
self.assertEqual(sum(split_list), 1024)
|
|
self.assert_page_aligned_boundaries(
|
|
split_list, extend_prefix_len=128, extend_len=1024, page_size=64
|
|
)
|
|
|
|
def test_page_aligned_split_keeps_tail_partial_page_unsplit(self):
|
|
split_list, info = build_page_aligned_in_seq_split_list(
|
|
total_len=1100,
|
|
extend_len=1100,
|
|
extend_prefix_len=0,
|
|
page_size=64,
|
|
cp_size=8,
|
|
)
|
|
|
|
self.assertTrue(info.page_aligned)
|
|
self.assertEqual(sum(split_list), 1100)
|
|
self.assertEqual(split_list[-1], 12)
|
|
self.assert_page_aligned_boundaries(
|
|
split_list, extend_prefix_len=0, extend_len=1100, page_size=64
|
|
)
|
|
|
|
def test_page_aligned_split_exposes_padded_extent_without_padding_split_list(self):
|
|
split_list, info = build_page_aligned_in_seq_split_list(
|
|
total_len=100,
|
|
extend_len=100,
|
|
extend_prefix_len=0,
|
|
page_size=64,
|
|
cp_size=8,
|
|
)
|
|
|
|
self.assertTrue(info.page_aligned)
|
|
self.assertEqual(sum(split_list), 100)
|
|
self.assertEqual(split_list[:2], [64, 36])
|
|
self.assertEqual(split_list[2:], [0] * 14)
|
|
self.assertEqual(info.extend_valid_tokens, 100)
|
|
self.assertEqual(info.extend_padded_pages, 2)
|
|
self.assertEqual(info.extend_padded_tokens, 128)
|
|
self.assertEqual(info.extend_padding_tokens, 28)
|
|
|
|
def test_page_aligned_split_falls_back_when_prefix_is_not_page_aligned(self):
|
|
split_list, info = build_page_aligned_in_seq_split_list(
|
|
total_len=1024,
|
|
extend_len=1024,
|
|
extend_prefix_len=1,
|
|
page_size=64,
|
|
cp_size=8,
|
|
)
|
|
|
|
self.assertFalse(info.page_aligned)
|
|
self.assertEqual(split_list, build_token_balanced_in_seq_split_list(1024, 8))
|
|
|
|
def test_page_aligned_split_pads_zero_segments_when_page_units_are_short(self):
|
|
split_list, info = build_page_aligned_in_seq_split_list(
|
|
total_len=512,
|
|
extend_len=512,
|
|
extend_prefix_len=0,
|
|
page_size=64,
|
|
cp_size=8,
|
|
)
|
|
|
|
self.assertTrue(info.page_aligned)
|
|
self.assertEqual(sum(split_list), 512)
|
|
self.assertEqual(split_list[:8], [64] * 8)
|
|
self.assertEqual(split_list[8:], [0] * 8)
|
|
self.assert_page_aligned_boundaries(
|
|
split_list, extend_prefix_len=0, extend_len=512, page_size=64
|
|
)
|
|
|
|
def test_page_aligned_split_allows_radix_hit_suffix_with_one_page_per_rank(self):
|
|
split_list, info = build_page_aligned_in_seq_split_list(
|
|
total_len=512,
|
|
extend_len=512,
|
|
extend_prefix_len=54464,
|
|
page_size=64,
|
|
cp_size=8,
|
|
)
|
|
|
|
self.assertTrue(info.page_aligned)
|
|
self.assertEqual(sum(split_list), 512)
|
|
self.assertEqual(split_list[:8], [64] * 8)
|
|
self.assertEqual(split_list[8:], [0] * 8)
|
|
self.assert_page_aligned_boundaries(
|
|
split_list, extend_prefix_len=54464, extend_len=512, page_size=64
|
|
)
|
|
|
|
def test_page_aligned_split_keeps_short_radix_hit_suffix_page_aligned(self):
|
|
split_list, info = build_page_aligned_in_seq_split_list(
|
|
total_len=256,
|
|
extend_len=256,
|
|
extend_prefix_len=54464,
|
|
page_size=64,
|
|
cp_size=8,
|
|
)
|
|
|
|
self.assertTrue(info.page_aligned)
|
|
self.assertEqual(sum(split_list), 256)
|
|
self.assertEqual(split_list[:4], [64] * 4)
|
|
self.assertEqual(split_list[4:], [0] * 12)
|
|
self.assert_page_aligned_boundaries(
|
|
split_list, extend_prefix_len=54464, extend_len=256, page_size=64
|
|
)
|
|
|
|
def test_can_cp_split_uses_compute_padding_for_short_radix_hit_suffix(self):
|
|
class Mode:
|
|
def is_context_parallel_extend(self):
|
|
return True
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_seq_lens_cpu=[65],
|
|
extend_prefix_lens_cpu=[54464],
|
|
token_to_kv_pool=SimpleNamespace(page_size=64),
|
|
forward_mode=Mode(),
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
|
|
return_value=True,
|
|
),
|
|
):
|
|
self.assertTrue(can_cp_split(128, 8, True, forward_batch))
|
|
|
|
def test_can_cp_split_uses_compute_padding_when_page_units_do_not_cover_all_lanes(
|
|
self,
|
|
):
|
|
class Mode:
|
|
def is_context_parallel_extend(self):
|
|
return True
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_seq_lens_cpu=[256],
|
|
extend_prefix_lens_cpu=[54464],
|
|
token_to_kv_pool=SimpleNamespace(page_size=64),
|
|
forward_mode=Mode(),
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
|
|
return_value=True,
|
|
),
|
|
):
|
|
self.assertTrue(can_cp_split(256, 8, True, forward_batch))
|
|
|
|
def test_can_cp_split_uses_compute_padding_for_current_only_one_page_suffix(
|
|
self,
|
|
):
|
|
class Mode:
|
|
def is_context_parallel_extend(self):
|
|
return True
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_seq_lens_cpu=[64],
|
|
extend_prefix_lens_cpu=[0],
|
|
token_to_kv_pool=SimpleNamespace(page_size=64),
|
|
forward_mode=Mode(),
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
|
|
return_value=True,
|
|
),
|
|
):
|
|
self.assertTrue(can_cp_split(64, 8, True, forward_batch))
|
|
|
|
def test_can_cp_split_enables_cp_draft_shared_kv_draft_extend(self):
|
|
class DraftMode:
|
|
def is_context_parallel_extend(self, include_draft_extend_v2=False):
|
|
return False
|
|
|
|
def is_draft_extend(self, include_v2=False):
|
|
return True
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_seq_lens_cpu=[4],
|
|
extend_prefix_lens_cpu=[0],
|
|
token_to_kv_pool=SimpleNamespace(page_size=64),
|
|
forward_mode=DraftMode(),
|
|
)
|
|
|
|
with (
|
|
patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
|
|
return_value=True,
|
|
),
|
|
):
|
|
self.assertTrue(can_cp_split(8, 8, True, forward_batch))
|
|
|
|
def test_nsa_use_prefill_cp_enables_cp_draft_shared_kv_draft_extend(self):
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
forward_mode=ForwardMode.DRAFT_EXTEND,
|
|
nsa_cp_metadata=NSAContextParallelMetadata(batch_size=1),
|
|
)
|
|
|
|
with patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}):
|
|
self.assertTrue(nsa_use_prefill_cp(forward_batch, True))
|
|
|
|
def test_can_cp_split_uses_compute_padding_per_request_for_batched_tiny_suffix(
|
|
self,
|
|
):
|
|
class Mode:
|
|
def is_context_parallel_extend(self):
|
|
return True
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_seq_lens_cpu=[65, 64],
|
|
extend_prefix_lens_cpu=[54464, 8192],
|
|
token_to_kv_pool=SimpleNamespace(page_size=64),
|
|
forward_mode=Mode(),
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
|
|
return_value=True,
|
|
),
|
|
):
|
|
self.assertTrue(can_cp_split(129, 8, True, forward_batch))
|
|
|
|
def test_can_cp_split_fails_on_non_page_aligned_cp_shared_prefix(self):
|
|
class Mode:
|
|
def is_context_parallel_extend(self):
|
|
return True
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_seq_lens_cpu=[1024],
|
|
extend_prefix_lens_cpu=[65],
|
|
token_to_kv_pool=SimpleNamespace(page_size=64),
|
|
forward_mode=Mode(),
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
|
|
return_value=True,
|
|
),
|
|
self.assertRaisesRegex(
|
|
RuntimeError,
|
|
r"\[CP_SHARED_KV_FAIL_FAST\]\[cp_split_non_page_aligned_prefix\]",
|
|
),
|
|
):
|
|
can_cp_split(1089, 8, True, forward_batch)
|
|
|
|
def test_can_cp_split_keeps_cp_for_radix_hit_suffix_with_one_page_per_rank(self):
|
|
class Mode:
|
|
def is_context_parallel_extend(self):
|
|
return True
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_seq_lens_cpu=[512],
|
|
extend_prefix_lens_cpu=[54464],
|
|
token_to_kv_pool=SimpleNamespace(page_size=64),
|
|
forward_mode=Mode(),
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
|
|
return_value=True,
|
|
),
|
|
):
|
|
self.assertTrue(can_cp_split(512, 8, True, forward_batch))
|
|
|
|
def test_page_aligned_split_adds_padding_tokens_to_last_segment(self):
|
|
split_list, info = build_page_aligned_in_seq_split_list(
|
|
total_len=1040,
|
|
extend_len=1024,
|
|
extend_prefix_len=0,
|
|
page_size=64,
|
|
cp_size=8,
|
|
)
|
|
|
|
self.assertTrue(info.page_aligned)
|
|
self.assertEqual(sum(split_list), 1040)
|
|
self.assertEqual(split_list[-1], 80)
|
|
self.assert_page_aligned_boundaries(
|
|
split_list, extend_prefix_len=0, extend_len=1024, page_size=64
|
|
)
|
|
|
|
def test_last_token_owner_uses_actual_token_count_when_batch_is_padded(self):
|
|
# Padded prefill can have 64 model tokens while the real prompt has only
|
|
# 11 tokens. In in-seq split with cp_size=8, the real last token is in
|
|
# segment 2, not in rank 0's trailing padded segment.
|
|
split_list = [4] * 16
|
|
|
|
owner, local_offset = _get_in_seq_last_token_owner_and_offset(
|
|
split_list=split_list,
|
|
cp_size=8,
|
|
actual_token_count=11,
|
|
)
|
|
|
|
self.assertEqual(owner, 2)
|
|
self.assertEqual(local_offset, 2)
|
|
|
|
def test_last_token_owner_keeps_existing_unpadded_fast_path_location(self):
|
|
split_list = [4] * 16
|
|
|
|
owner, local_offset = _get_in_seq_last_token_owner_and_offset(
|
|
split_list=split_list,
|
|
cp_size=8,
|
|
actual_token_count=64,
|
|
)
|
|
|
|
self.assertEqual(owner, 0)
|
|
self.assertEqual(local_offset, 7)
|
|
|
|
def test_batch_page_aligned_plan_keeps_request_boundaries_and_last_token_owners(
|
|
self,
|
|
):
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[4, 9],
|
|
prefix_lens=[0, 8],
|
|
page_size=4,
|
|
cp_size=2,
|
|
cp_rank=0,
|
|
)
|
|
|
|
self.assertEqual(plan.batch_size, 2)
|
|
self.assertEqual(plan.request_split_lists, [[4, 0, 0, 0], [4, 4, 1, 0]])
|
|
self.assertEqual(plan.request_padded_pages, [1, 3])
|
|
self.assertEqual(plan.request_padded_tokens, [4, 12])
|
|
self.assertEqual(plan.request_token_offsets, [0, 4])
|
|
self.assertEqual(plan.request_padded_token_offsets, [0, 4])
|
|
self.assertEqual(plan.request_page_offsets, [0, 1])
|
|
self.assertEqual(plan.request_last_token_owner, [0, 1])
|
|
self.assertEqual(plan.request_last_token_local_offset, [3, 4])
|
|
self.assertEqual(plan.request_rank_local_tokens, [4, 4])
|
|
self.assertEqual(plan.request_rank_local_offsets, [0, 4])
|
|
self.assertEqual(plan.request_kv_len_prev, [4, 4])
|
|
self.assertEqual(plan.request_kv_len_next, [4, 9])
|
|
self.assertEqual(plan.request_actual_seq_q_prev, [4, 4])
|
|
self.assertEqual(plan.request_actual_seq_q_next, [0, 0])
|
|
self.assertEqual(plan.flat_segment_request_ids, [0, 0, 0, 0, 1, 1, 1, 1])
|
|
self.assertEqual(plan.flat_segment_offsets, [0, 4, 4, 4, 0, 4, 8, 9])
|
|
|
|
def test_batch_plan_exposes_compute_padding_without_inflating_valid_cache_extent(
|
|
self,
|
|
):
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[40320],
|
|
page_size=64,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
|
|
self.assertTrue(plan.compute_padding_enabled)
|
|
self.assertEqual(
|
|
plan.request_valid_split_lists,
|
|
[[64, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
|
|
)
|
|
self.assertEqual(
|
|
plan.request_compute_split_lists,
|
|
[[64, 64, 64, 64, 64, 64, 64, 64, 0, 0, 0, 0, 0, 0, 0, 0]],
|
|
)
|
|
self.assertEqual(plan.request_valid_padded_pages, [2])
|
|
self.assertEqual(plan.request_valid_padded_tokens, [128])
|
|
self.assertEqual(plan.request_compute_padded_pages, [8])
|
|
self.assertEqual(plan.request_compute_padded_tokens, [512])
|
|
self.assertEqual(plan.request_compute_padding_tokens, [447])
|
|
self.assertEqual(plan.request_compute_rank_local_tokens, [64])
|
|
self.assertEqual(plan.request_compute_rank_local_offsets, [0])
|
|
self.assertEqual(plan.request_valid_rank_local_tokens, [1])
|
|
self.assertEqual(plan.request_valid_rank_local_offsets, [0])
|
|
self.assertEqual(plan.request_last_token_owner, [1])
|
|
self.assertEqual(plan.request_last_token_local_offset, [0])
|
|
|
|
# Compatibility aliases for cache/page accounting stay valid-token
|
|
# based. Query-length metadata exposes both valid and compute rows:
|
|
# consumers must choose the view that matches their actual q layout.
|
|
self.assertEqual(plan.request_split_lists, plan.request_valid_split_lists)
|
|
self.assertEqual(plan.request_padded_pages, plan.request_valid_padded_pages)
|
|
self.assertEqual(plan.request_actual_seq_q_prev, [64])
|
|
self.assertEqual(plan.request_actual_seq_q_next, [0])
|
|
self.assertEqual(plan.request_valid_seq_q_prev, [1])
|
|
self.assertEqual(plan.request_valid_seq_q_next, [0])
|
|
self.assertEqual(plan.request_compute_seq_q_prev, [64])
|
|
self.assertEqual(plan.request_compute_seq_q_next, [0])
|
|
|
|
def test_bs_gt1_debug_log_is_env_gated_and_limited(self):
|
|
nsa_utils._CP_SHARED_KV_BS_GT1_DEBUG_COUNTS.clear()
|
|
|
|
with envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.override(False):
|
|
with patch.object(nsa_utils.logger, "info") as info:
|
|
nsa_utils.log_cp_shared_kv_bs_gt1_debug(
|
|
"unit_test",
|
|
"bs=%s",
|
|
2,
|
|
)
|
|
self.assertFalse(info.called)
|
|
|
|
with envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.override(True):
|
|
with envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG_LIMIT.override(1):
|
|
with patch.object(nsa_utils.logger, "info") as info:
|
|
nsa_utils.log_cp_shared_kv_bs_gt1_debug(
|
|
"unit_test",
|
|
"bs=%s",
|
|
2,
|
|
)
|
|
nsa_utils.log_cp_shared_kv_bs_gt1_debug(
|
|
"unit_test",
|
|
"bs=%s",
|
|
3,
|
|
)
|
|
self.assertEqual(info.call_count, 1)
|
|
self.assertIn(
|
|
"[CP_SHARED_KV_BS_GT1_DEBUG]",
|
|
info.call_args.args[0],
|
|
)
|
|
|
|
def test_index_topk_batch_lengths_follow_actual_q_rows_not_compute_alias(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import (
|
|
_select_batch_topk_query_lengths,
|
|
)
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[40387],
|
|
prefix_lens=[0],
|
|
page_size=64,
|
|
cp_size=8,
|
|
cp_rank=0,
|
|
)
|
|
valid_local_rows = (
|
|
plan.request_valid_seq_q_prev[0] + plan.request_valid_seq_q_next[0]
|
|
)
|
|
compute_local_rows = (
|
|
plan.request_compute_seq_q_prev[0] + plan.request_compute_seq_q_next[0]
|
|
)
|
|
self.assertEqual(valid_local_rows, 4995)
|
|
self.assertEqual(compute_local_rows, valid_local_rows)
|
|
self.assertFalse(plan.compute_padding_enabled)
|
|
|
|
local = split_tensor_by_cp_batch_plan(
|
|
torch.arange(40387, dtype=torch.int64),
|
|
plan,
|
|
mode="1d",
|
|
)
|
|
self.assertEqual(local.numel(), valid_local_rows)
|
|
self.assertEqual(local[:2560].tolist(), list(range(2560)))
|
|
self.assertEqual(local[2560:4995].tolist(), list(range(37952, 40387)))
|
|
|
|
local_valid = select_cp_local_valid_rows_for_cache_write(
|
|
SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
)
|
|
),
|
|
local,
|
|
)
|
|
self.assertEqual(local_valid.numel(), valid_local_rows)
|
|
self.assertEqual(local_valid.tolist(), local.tolist())
|
|
|
|
selected = _select_batch_topk_query_lengths(
|
|
cp_metadata=NSAContextParallelMetadata(batch_size=1, batch_plan=plan),
|
|
batch_plan=plan,
|
|
batch_size=1,
|
|
q_tokens=valid_local_rows,
|
|
weights_tokens=valid_local_rows,
|
|
)
|
|
|
|
self.assertFalse(selected.uses_compute_query_rows)
|
|
self.assertEqual(selected.request_seq_q_prev, plan.request_valid_seq_q_prev)
|
|
self.assertEqual(selected.request_seq_q_next, plan.request_valid_seq_q_next)
|
|
self.assertEqual(
|
|
selected.request_valid_seq_q_prev, plan.request_valid_seq_q_prev
|
|
)
|
|
self.assertEqual(
|
|
selected.request_valid_seq_q_next, plan.request_valid_seq_q_next
|
|
)
|
|
|
|
selected_compute_alias = _select_batch_topk_query_lengths(
|
|
cp_metadata=NSAContextParallelMetadata(batch_size=1, batch_plan=plan),
|
|
batch_plan=plan,
|
|
batch_size=1,
|
|
q_tokens=compute_local_rows,
|
|
weights_tokens=compute_local_rows,
|
|
)
|
|
|
|
self.assertFalse(selected_compute_alias.uses_compute_query_rows)
|
|
self.assertEqual(
|
|
selected_compute_alias.request_seq_q_prev, plan.request_compute_seq_q_prev
|
|
)
|
|
self.assertEqual(
|
|
selected_compute_alias.request_seq_q_next, plan.request_compute_seq_q_next
|
|
)
|
|
self.assertEqual(
|
|
selected_compute_alias.request_valid_seq_q_prev, plan.request_valid_seq_q_prev
|
|
)
|
|
self.assertEqual(
|
|
selected_compute_alias.request_valid_seq_q_next, plan.request_valid_seq_q_next
|
|
)
|
|
|
|
def test_batch_plan_keeps_long_page_tail_out_of_compute_padding(self):
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[40387],
|
|
prefix_lens=[0],
|
|
page_size=64,
|
|
cp_size=8,
|
|
cp_rank=0,
|
|
)
|
|
|
|
self.assertFalse(plan.compute_padding_enabled)
|
|
self.assertEqual(plan.request_valid_padded_pages, [632])
|
|
self.assertEqual(plan.request_valid_padded_tokens, [40448])
|
|
self.assertEqual(plan.request_compute_padded_pages, [632])
|
|
self.assertEqual(plan.request_compute_padded_tokens, [40387])
|
|
self.assertEqual(plan.request_compute_padding_tokens, [0])
|
|
self.assertEqual(
|
|
plan.request_compute_split_lists,
|
|
plan.request_valid_split_lists,
|
|
)
|
|
|
|
def test_batch_plan_compute_padding_only_pads_tiny_request_in_mixed_batch(self):
|
|
import torch
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65, 40387],
|
|
prefix_lens=[40320, 0],
|
|
page_size=64,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
|
|
self.assertTrue(plan.compute_padding_enabled)
|
|
self.assertEqual(plan.request_valid_padded_pages, [2, 632])
|
|
self.assertEqual(plan.request_compute_padded_pages, [8, 632])
|
|
self.assertEqual(plan.request_compute_padded_tokens, [512, 40387])
|
|
self.assertEqual(plan.request_compute_padding_tokens, [447, 0])
|
|
self.assertEqual(
|
|
plan.request_compute_split_lists[1],
|
|
plan.request_valid_split_lists[1],
|
|
)
|
|
|
|
local = split_tensor_by_cp_batch_plan(
|
|
torch.arange(65 + 40387, dtype=torch.int64),
|
|
plan,
|
|
mode="1d",
|
|
)
|
|
self.assertEqual(local.numel(), sum(plan.request_compute_rank_local_tokens))
|
|
|
|
valid = select_cp_local_valid_rows_for_cache_write(
|
|
SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
batch_plan=plan,
|
|
)
|
|
),
|
|
local,
|
|
)
|
|
self.assertEqual(valid.numel(), sum(plan.request_valid_rank_local_tokens))
|
|
|
|
def test_batch_plan_compute_padding_is_per_request_not_batch_total(self):
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65, 1024],
|
|
prefix_lens=[40320, 8192],
|
|
page_size=64,
|
|
cp_size=8,
|
|
cp_rank=0,
|
|
)
|
|
|
|
self.assertTrue(plan.compute_padding_enabled)
|
|
self.assertEqual(plan.request_valid_padded_pages, [2, 16])
|
|
self.assertEqual(plan.request_compute_padded_pages, [8, 16])
|
|
self.assertEqual(plan.request_compute_padded_tokens, [512, 1024])
|
|
self.assertEqual(plan.request_compute_padding_tokens, [447, 0])
|
|
self.assertEqual(plan.request_compute_rank_local_tokens, [64, 128])
|
|
self.assertEqual(plan.request_compute_rank_local_offsets, [0, 64])
|
|
self.assertEqual(plan.request_valid_rank_local_tokens, [64, 128])
|
|
self.assertEqual(plan.request_valid_rank_local_offsets, [0, 64])
|
|
self.assertEqual(plan.request_last_token_owner, [1, 0])
|
|
|
|
def test_batch_plan_stable_helpers_split_and_build_page_owner_plan(self):
|
|
import torch
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[4, 9],
|
|
prefix_lens=[0, 8],
|
|
page_size=4,
|
|
cp_size=2,
|
|
cp_rank=1,
|
|
)
|
|
forward_batch = SimpleNamespace(nsa_cp_metadata=SimpleNamespace(batch_plan=plan))
|
|
|
|
self.assertIs(get_cp_shared_kv_batch_plan(forward_batch), plan)
|
|
self.assertEqual(build_flat_page_owner_plan(plan), [0, 0, 1, 1])
|
|
|
|
local_1d = split_tensor_by_cp_batch_plan(torch.arange(13), plan, mode="1d")
|
|
self.assertEqual(local_1d.tolist(), [0, 0, 0, 0, 8, 9, 10, 11, 12])
|
|
|
|
local_data = split_tensor_by_cp_batch_plan(
|
|
torch.arange(13 * 2).view(13, 2), plan, mode="data"
|
|
)
|
|
self.assertEqual(
|
|
local_data[:, 0].tolist(),
|
|
[0, 0, 0, 0, 16, 18, 20, 22, 24],
|
|
)
|
|
|
|
def test_collect_last_token_hidden_uses_batch_owner_metadata(self):
|
|
import torch
|
|
|
|
hidden_states = torch.tensor(
|
|
[[10.0], [11.0], [12.0], [13.0], [20.0], [21.0], [22.0], [23.0]]
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
extend_seq_lens_cpu=[4, 9],
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
request_last_token_owner=[0, 1],
|
|
request_last_token_local_offset=[3, 4],
|
|
request_rank_local_offsets=[0, 4],
|
|
),
|
|
)
|
|
|
|
def fake_all_gather(output, local_last):
|
|
self.assertEqual(local_last.tolist(), [[13.0], [0.0]])
|
|
output.copy_(torch.tensor([[13.0], [0.0], [0.0], [99.0]]))
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank",
|
|
return_value=0,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor",
|
|
side_effect=fake_all_gather,
|
|
),
|
|
):
|
|
collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 2)
|
|
|
|
self.assertEqual(collected.tolist(), [[13.0], [99.0]])
|
|
|
|
def test_collect_last_token_hidden_uses_compute_padding_for_single_request(self):
|
|
import torch
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[40320],
|
|
page_size=64,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
hidden_states = torch.zeros((64, 1), dtype=torch.float32)
|
|
hidden_states[0] = 123.0
|
|
hidden_states[1] = 999.0
|
|
forward_batch = SimpleNamespace(
|
|
extend_seq_lens_cpu=[65],
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
),
|
|
)
|
|
|
|
def fake_all_gather(output, local_last):
|
|
self.assertEqual(local_last.tolist(), [[123.0]])
|
|
output.zero_()
|
|
output[1] = local_last[0]
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank",
|
|
return_value=1,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor",
|
|
side_effect=fake_all_gather,
|
|
),
|
|
):
|
|
collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 8)
|
|
|
|
self.assertEqual(collected.tolist(), [[123.0]])
|
|
|
|
def test_collect_last_token_hidden_uses_compute_rank_offsets_for_batch(self):
|
|
import torch
|
|
|
|
hidden_states = torch.zeros((8, 1), dtype=torch.float32)
|
|
hidden_states[0] = 10.0
|
|
hidden_states[1] = 99.0
|
|
hidden_states[4] = 20.0
|
|
forward_batch = SimpleNamespace(
|
|
extend_seq_lens_cpu=[5, 5],
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
request_last_token_owner=[1, 1],
|
|
request_last_token_local_offset=[0, 0],
|
|
request_rank_local_offsets=[0, 1],
|
|
request_compute_rank_local_offsets=[0, 4],
|
|
compute_padding_enabled=True,
|
|
),
|
|
)
|
|
|
|
def fake_all_gather(output, local_last):
|
|
self.assertEqual(local_last.tolist(), [[10.0], [20.0]])
|
|
output.copy_(torch.tensor([[0.0], [0.0], [10.0], [20.0]]))
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank",
|
|
return_value=1,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor",
|
|
side_effect=fake_all_gather,
|
|
),
|
|
):
|
|
collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 2)
|
|
|
|
self.assertEqual(collected.tolist(), [[10.0], [20.0]])
|
|
|
|
def test_collect_last_token_hidden_fails_fast_without_batch_owner_metadata(self):
|
|
import torch
|
|
|
|
forward_batch = SimpleNamespace(
|
|
extend_seq_lens_cpu=[4, 9],
|
|
nsa_cp_metadata=NSAContextParallelMetadata(batch_size=2),
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank",
|
|
return_value=0,
|
|
),
|
|
self.assertRaisesRegex(
|
|
RuntimeError,
|
|
r"\[CP_SHARED_KV_FAIL_FAST\]\[batch_gt1_missing_last_token_metadata\]",
|
|
),
|
|
):
|
|
cp_collect_last_token_hidden(torch.zeros((8, 1)), forward_batch, 2)
|
|
|
|
def test_deepseek_nextn_cp_draft_bs_gt1_fails_fast_on_hidden_shape_fallback(
|
|
self,
|
|
):
|
|
import torch
|
|
|
|
model = DeepseekModelNextN.__new__(DeepseekModelNextN)
|
|
model._debug_cp_draft_shared_kv = lambda _message: None
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
batch_size=2,
|
|
extend_seq_lens_cpu=[4, 9],
|
|
nsa_cp_metadata=NSAContextParallelMetadata(batch_size=2),
|
|
)
|
|
|
|
with (
|
|
patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}),
|
|
patch(
|
|
"sglang.srt.models.deepseek_nextn.get_attention_cp_rank",
|
|
return_value=0,
|
|
),
|
|
patch(
|
|
"sglang.srt.models.deepseek_nextn.get_attention_cp_size",
|
|
return_value=8,
|
|
),
|
|
self.assertRaisesRegex(
|
|
RuntimeError,
|
|
r"\[CP_SHARED_KV_FAIL_FAST\]\[draft_batch_gt1_spec_hidden_shape_mismatch\]",
|
|
),
|
|
):
|
|
model._get_cp_local_spec_hidden_states(
|
|
forward_batch,
|
|
torch.zeros((3, 2)),
|
|
full_num_tokens=13,
|
|
local_num_tokens=8,
|
|
)
|
|
|
|
def test_cp_draft_padding_keeps_local_hidden_when_static_tokens_are_shorter(self):
|
|
import torch
|
|
|
|
class FakeAttnBackend:
|
|
def get_cuda_graph_seq_len_fill_value(self):
|
|
return 0
|
|
|
|
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
|
|
spec_info = EagleDraftInput(
|
|
hidden_states=torch.ones((64, 2), dtype=torch.float32),
|
|
verified_id=torch.tensor([1], dtype=torch.int64),
|
|
num_tokens_per_req=1,
|
|
num_tokens_for_logprob_per_req=1,
|
|
cp_local_hidden_states=True,
|
|
)
|
|
forward_batch = ForwardBatch(
|
|
forward_mode=ForwardMode.DRAFT_EXTEND,
|
|
batch_size=1,
|
|
input_ids=torch.arange(8, dtype=torch.int64),
|
|
req_pool_indices=torch.tensor([0], dtype=torch.int64),
|
|
seq_lens=torch.tensor([4], dtype=torch.int32),
|
|
out_cache_loc=torch.arange(8, dtype=torch.int64),
|
|
seq_lens_sum=4,
|
|
positions=torch.arange(8, dtype=torch.int64),
|
|
lora_ids=[None],
|
|
spec_info=spec_info,
|
|
uses_cp_shared_kv=True,
|
|
)
|
|
|
|
with patch.dict(os.environ, {"SGLANG_CP_DRAFT_SHARED_KV": "1"}):
|
|
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
|
|
|
|
self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2))
|
|
self.assertTrue(
|
|
torch.equal(
|
|
forward_batch.hidden_states_backup,
|
|
torch.ones((64, 2), dtype=torch.float32),
|
|
)
|
|
)
|
|
|
|
def test_cp_draft_padding_keeps_marked_cp_local_hidden_before_cp_flags_are_visible(
|
|
self,
|
|
):
|
|
import torch
|
|
|
|
class FakeAttnBackend:
|
|
def get_cuda_graph_seq_len_fill_value(self):
|
|
return 0
|
|
|
|
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
|
|
spec_info = EagleDraftInput(
|
|
hidden_states=torch.ones((64, 2), dtype=torch.float32),
|
|
verified_id=torch.tensor([1], dtype=torch.int64),
|
|
num_tokens_per_req=1,
|
|
num_tokens_for_logprob_per_req=1,
|
|
cp_local_hidden_states=True,
|
|
)
|
|
forward_batch = ForwardBatch(
|
|
forward_mode=ForwardMode.DRAFT_EXTEND,
|
|
batch_size=1,
|
|
input_ids=torch.arange(8, dtype=torch.int64),
|
|
req_pool_indices=torch.tensor([0], dtype=torch.int64),
|
|
seq_lens=torch.tensor([4], dtype=torch.int32),
|
|
out_cache_loc=torch.arange(8, dtype=torch.int64),
|
|
seq_lens_sum=4,
|
|
positions=torch.arange(8, dtype=torch.int64),
|
|
lora_ids=[None],
|
|
spec_info=spec_info,
|
|
uses_cp_shared_kv=False,
|
|
)
|
|
|
|
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
|
|
|
|
self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2))
|
|
self.assertTrue(
|
|
torch.equal(
|
|
forward_batch.hidden_states_backup,
|
|
torch.ones((64, 2), dtype=torch.float32),
|
|
)
|
|
)
|
|
|
|
def test_cp_draft_padding_keeps_marked_cp_local_hidden_after_forward_mode_rewrite(
|
|
self,
|
|
):
|
|
import torch
|
|
|
|
class FakeAttnBackend:
|
|
def get_cuda_graph_seq_len_fill_value(self):
|
|
return 0
|
|
|
|
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
|
|
spec_info = EagleDraftInput(
|
|
hidden_states=torch.ones((64, 2), dtype=torch.float32),
|
|
verified_id=torch.tensor([1], dtype=torch.int64),
|
|
num_tokens_per_req=1,
|
|
num_tokens_for_logprob_per_req=1,
|
|
cp_local_hidden_states=True,
|
|
)
|
|
forward_batch = ForwardBatch(
|
|
# prepare_mlp_sync_batch can temporarily rewrite draft extend to
|
|
# EXTEND while static DP padding is being prepared. The draft
|
|
# side-channel contract must therefore be carried by spec_info, not
|
|
# inferred from forward_mode.
|
|
forward_mode=ForwardMode.EXTEND,
|
|
batch_size=1,
|
|
input_ids=torch.arange(8, dtype=torch.int64),
|
|
req_pool_indices=torch.tensor([0], dtype=torch.int64),
|
|
seq_lens=torch.tensor([4], dtype=torch.int32),
|
|
out_cache_loc=torch.arange(8, dtype=torch.int64),
|
|
seq_lens_sum=4,
|
|
positions=torch.arange(8, dtype=torch.int64),
|
|
lora_ids=[None],
|
|
spec_info=spec_info,
|
|
uses_cp_shared_kv=True,
|
|
)
|
|
|
|
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
|
|
|
|
self.assertEqual(tuple(forward_batch.spec_info.hidden_states.shape), (64, 2))
|
|
self.assertTrue(
|
|
torch.equal(
|
|
forward_batch.hidden_states_backup,
|
|
torch.ones((64, 2), dtype=torch.float32),
|
|
)
|
|
)
|
|
|
|
def test_cp_draft_padding_rejects_unmarked_oversized_hidden(self):
|
|
import torch
|
|
|
|
class FakeAttnBackend:
|
|
def get_cuda_graph_seq_len_fill_value(self):
|
|
return 0
|
|
|
|
model_runner = SimpleNamespace(attn_backend=FakeAttnBackend())
|
|
spec_info = EagleDraftInput(
|
|
hidden_states=torch.ones((64, 2), dtype=torch.float32),
|
|
verified_id=torch.tensor([1], dtype=torch.int64),
|
|
num_tokens_per_req=1,
|
|
num_tokens_for_logprob_per_req=1,
|
|
)
|
|
forward_batch = ForwardBatch(
|
|
forward_mode=ForwardMode.DRAFT_EXTEND,
|
|
batch_size=1,
|
|
input_ids=torch.arange(8, dtype=torch.int64),
|
|
req_pool_indices=torch.tensor([0], dtype=torch.int64),
|
|
seq_lens=torch.tensor([4], dtype=torch.int32),
|
|
out_cache_loc=torch.arange(8, dtype=torch.int64),
|
|
seq_lens_sum=4,
|
|
positions=torch.arange(8, dtype=torch.int64),
|
|
lora_ids=[None],
|
|
spec_info=spec_info,
|
|
uses_cp_shared_kv=False,
|
|
)
|
|
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
r"\[CP_SHARED_KV_FAIL_FAST\]\[draft_hidden_static_padding_mismatch\]",
|
|
):
|
|
forward_batch._pad_inputs_to_size(model_runner, num_tokens=8, bs=1)
|
|
|
|
def test_full_rerange_fails_fast_for_batch_metadata(self):
|
|
import torch
|
|
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(batch_size=2)
|
|
)
|
|
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
r"\[CP_SHARED_KV_FAIL_FAST\]\[batch_gt1_full_rerange_unsupported\]",
|
|
):
|
|
cp_all_gather_rerange_output(
|
|
torch.zeros((8, 1)), 2, forward_batch, stream=None
|
|
)
|
|
|
|
def test_cp_shared_kv_prepare_rejects_round_robin_mode(self):
|
|
class Mode:
|
|
def is_context_parallel_extend(self):
|
|
return True
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_seq_lens_cpu=[128],
|
|
extend_prefix_lens_cpu=[0],
|
|
token_to_kv_pool=SimpleNamespace(page_size=64),
|
|
forward_mode=Mode(),
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=True,
|
|
),
|
|
self.assertRaisesRegex(
|
|
RuntimeError,
|
|
r"\[CP_SHARED_KV_FAIL_FAST\]\[round_robin_unsupported\]",
|
|
),
|
|
):
|
|
prepare_input_dp_with_cp_dsa(
|
|
128,
|
|
cp_rank=0,
|
|
cp_size=2,
|
|
seqs_len=[128],
|
|
forward_batch=forward_batch,
|
|
page_size=64,
|
|
)
|
|
|
|
def test_cp_shared_kv_prepare_uses_batch_plan_for_bs1_compute_padding(self):
|
|
class Mode:
|
|
def is_context_parallel_extend(self):
|
|
return True
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_seq_lens_cpu=[65],
|
|
extend_prefix_lens_cpu=[0],
|
|
token_to_kv_pool=SimpleNamespace(page_size=64),
|
|
forward_mode=Mode(),
|
|
)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
metadata = prepare_input_dp_with_cp_dsa(
|
|
65,
|
|
cp_rank=1,
|
|
cp_size=8,
|
|
seqs_len=[65],
|
|
forward_batch=forward_batch,
|
|
page_size=64,
|
|
)
|
|
|
|
self.assertIsNotNone(metadata.batch_plan)
|
|
self.assertEqual(metadata.batch_size, 1)
|
|
self.assertTrue(metadata.compute_padding_enabled)
|
|
self.assertEqual(
|
|
metadata.request_valid_split_lists,
|
|
[[64, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
|
|
)
|
|
self.assertEqual(
|
|
metadata.request_compute_split_lists,
|
|
[[64, 64, 64, 64, 64, 64, 64, 64, 0, 0, 0, 0, 0, 0, 0, 0]],
|
|
)
|
|
self.assertEqual(metadata.split_list, metadata.request_compute_split_lists[0])
|
|
self.assertEqual(metadata.max_rank_len, [64] * 8)
|
|
self.assertEqual(metadata.per_rank_actual_token, [64] * 8)
|
|
self.assertEqual(metadata.actual_seq_q_prev, 64)
|
|
self.assertEqual(metadata.actual_seq_q_next, 0)
|
|
self.assertEqual(metadata.request_valid_seq_q_prev, [1])
|
|
self.assertEqual(metadata.request_valid_seq_q_next, [0])
|
|
|
|
def test_cp_shared_kv_all_gather_rejects_round_robin_mode(self):
|
|
import torch
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
split_list=[2, 2, 2, 2],
|
|
zigzag_index=[0, 3],
|
|
reverse_split_len=[2, 2, 2, 2],
|
|
cp_reverse_index=[0, 2, 3, 1],
|
|
total_seq_lens=torch.tensor(8),
|
|
),
|
|
)
|
|
|
|
with (
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=True,
|
|
),
|
|
self.assertRaisesRegex(
|
|
RuntimeError,
|
|
r"\[CP_SHARED_KV_FAIL_FAST\]\[round_robin_unsupported\]",
|
|
),
|
|
):
|
|
cp_all_gather_rerange_output(
|
|
torch.zeros((4, 1)), 2, forward_batch, stream=None
|
|
)
|
|
|
|
def test_batch_in_seq_all_gather_rerange_restores_request_order_bf16(self):
|
|
import torch
|
|
|
|
cp_size = 2
|
|
request_split_lists = [
|
|
[2, 1, 3, 0],
|
|
[1, 2, 0, 1],
|
|
]
|
|
input_tensor_all, expected = self._build_batch_rerange_case(
|
|
cp_size=cp_size,
|
|
request_split_lists=request_split_lists,
|
|
row_width=3,
|
|
dtype=torch.bfloat16,
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
request_split_lists=request_split_lists,
|
|
max_rank_len=[6, 6],
|
|
)
|
|
)
|
|
|
|
actual = _torch_batch_in_seq_all_gather_rerange(
|
|
input_tensor_all,
|
|
forward_batch,
|
|
cp_size=cp_size,
|
|
)
|
|
|
|
self.assertEqual(actual.dtype, torch.bfloat16)
|
|
self.assertTrue(torch.equal(actual, expected))
|
|
|
|
def test_batch_in_seq_all_gather_rerange_treats_fp8_payload_as_opaque_rows(self):
|
|
import torch
|
|
|
|
cp_size = 2
|
|
request_split_lists = [
|
|
[1, 2, 1, 0],
|
|
[2, 0, 1, 1],
|
|
]
|
|
input_tensor_all, expected = self._build_batch_rerange_case(
|
|
cp_size=cp_size,
|
|
request_split_lists=request_split_lists,
|
|
row_width=5,
|
|
dtype=torch.uint8,
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
request_split_lists=request_split_lists,
|
|
max_rank_len=[4, 4],
|
|
)
|
|
)
|
|
|
|
actual = _torch_batch_in_seq_all_gather_rerange(
|
|
input_tensor_all,
|
|
forward_batch,
|
|
cp_size=cp_size,
|
|
)
|
|
|
|
self.assertEqual(actual.dtype, torch.uint8)
|
|
self.assertTrue(torch.equal(actual, expected))
|
|
|
|
def _build_batch_rerange_case(
|
|
self,
|
|
*,
|
|
cp_size,
|
|
request_split_lists,
|
|
row_width,
|
|
dtype,
|
|
):
|
|
import torch
|
|
|
|
rank_tokens = []
|
|
for rank in range(cp_size):
|
|
mirror = cp_size * 2 - rank - 1
|
|
rank_tokens.append(
|
|
sum(split[rank] + split[mirror] for split in request_split_lists)
|
|
)
|
|
max_rank_token = max(rank_tokens)
|
|
total_tokens = sum(sum(split) for split in request_split_lists)
|
|
input_tensor_all = torch.zeros(
|
|
(max_rank_token * cp_size, row_width),
|
|
dtype=dtype,
|
|
)
|
|
expected = torch.empty((total_tokens, row_width), dtype=dtype)
|
|
|
|
next_value = 1
|
|
request_segments = []
|
|
for split in request_split_lists:
|
|
segments = []
|
|
for segment_len in split:
|
|
if dtype == torch.uint8:
|
|
rows = (
|
|
torch.arange(
|
|
next_value,
|
|
next_value + segment_len * row_width,
|
|
dtype=torch.int64,
|
|
)
|
|
.remainder(251)
|
|
.to(torch.uint8)
|
|
.view(segment_len, row_width)
|
|
)
|
|
else:
|
|
rows = (
|
|
torch.arange(
|
|
next_value,
|
|
next_value + segment_len * row_width,
|
|
dtype=torch.float32,
|
|
)
|
|
.view(segment_len, row_width)
|
|
.to(dtype)
|
|
)
|
|
next_value += segment_len * row_width
|
|
segments.append(rows)
|
|
request_segments.append(segments)
|
|
|
|
output_cursor = 0
|
|
for segments in request_segments:
|
|
for rows in segments:
|
|
expected[output_cursor : output_cursor + rows.shape[0]] = rows
|
|
output_cursor += rows.shape[0]
|
|
|
|
for rank in range(cp_size):
|
|
mirror = cp_size * 2 - rank - 1
|
|
rank_cursor = rank * max_rank_token
|
|
for segments in request_segments:
|
|
for segment_id in (rank, mirror):
|
|
rows = segments[segment_id]
|
|
input_tensor_all[
|
|
rank_cursor : rank_cursor + rows.shape[0]
|
|
] = rows
|
|
rank_cursor += rows.shape[0]
|
|
|
|
return input_tensor_all, expected
|
|
|
|
def test_local_pair_split_uses_metadata_lengths_not_half_split(self):
|
|
import torch
|
|
|
|
tensor = torch.arange(9)
|
|
|
|
prev, next_ = split_in_seq_cp_local_pair(tensor, 6, 3)
|
|
|
|
self.assertEqual(prev.tolist(), [0, 1, 2, 3, 4, 5])
|
|
self.assertEqual(next_.tolist(), [6, 7, 8])
|
|
|
|
def test_local_pair_split_rejects_stale_metadata(self):
|
|
import torch
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "local in-seq CP length mismatch"):
|
|
split_in_seq_cp_local_pair(torch.arange(9), 5, 5, name="q_fp8")
|
|
|
|
def test_cp_split_and_rebuild_1d_matches_in_seq_zigzag_order(self):
|
|
import torch
|
|
from types import SimpleNamespace
|
|
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
split_list=[2, 2, 2, 2, 2, 2, 2, 2],
|
|
zigzag_index=[1, 6],
|
|
)
|
|
)
|
|
|
|
local_locs = cp_split_and_rebuild_1d(forward_batch, torch.arange(16))
|
|
|
|
self.assertEqual(local_locs.tolist(), [2, 3, 12, 13])
|
|
|
|
def test_cp_split_and_rebuild_data_keeps_batch_request_boundaries(self):
|
|
import torch
|
|
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
request_extend_lens=[4, 9],
|
|
request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]],
|
|
request_zigzag_indices=[[0, 3], [0, 3]],
|
|
)
|
|
)
|
|
tensor = torch.arange(13 * 2).view(13, 2)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local = cp_split_and_rebuild_data(forward_batch, tensor)
|
|
|
|
self.assertEqual(local[:, 0].tolist(), list(range(0, 16, 2)))
|
|
|
|
def test_cp_split_and_rebuild_data_preserves_non_token_dimensions(self):
|
|
import torch
|
|
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
request_extend_lens=[4, 9],
|
|
request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]],
|
|
request_zigzag_indices=[[0, 3], [0, 3]],
|
|
)
|
|
)
|
|
tensor = torch.arange(13 * 2 * 3, dtype=torch.float32).view(13, 2, 3)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local = cp_split_and_rebuild_data(forward_batch, tensor)
|
|
|
|
self.assertEqual(tuple(local.shape), (8, 2, 3))
|
|
self.assertTrue(torch.equal(local[0], tensor[0]))
|
|
|
|
def test_cp_split_and_rebuild_data_uses_compute_padding_rows(self):
|
|
import torch
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[40320],
|
|
page_size=64,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
)
|
|
)
|
|
tensor = torch.arange(65 * 2, dtype=torch.float32).view(65, 2)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local = cp_split_and_rebuild_data(forward_batch, tensor)
|
|
|
|
self.assertEqual(local.shape, (64, 2))
|
|
self.assertEqual(local[0].tolist(), [128.0, 129.0])
|
|
self.assertTrue(torch.equal(local[1:], torch.zeros((63, 2))))
|
|
|
|
def test_cp_split_and_rebuild_data_ignores_trailing_static_padding_rows(self):
|
|
import torch
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[4],
|
|
prefix_lens=[0],
|
|
page_size=4,
|
|
cp_size=2,
|
|
cp_rank=1,
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
)
|
|
)
|
|
tensor = torch.arange(8 * 2, dtype=torch.float32).view(8, 2)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local = cp_split_and_rebuild_data(forward_batch, tensor)
|
|
|
|
self.assertEqual(local.shape, (4, 2))
|
|
self.assertTrue(torch.equal(local, torch.zeros((4, 2))))
|
|
|
|
def test_cp_split_and_rebuild_data_ignores_mlp_sync_static_padding_without_compute_padding(
|
|
self,
|
|
):
|
|
import torch
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[8],
|
|
prefix_lens=[0],
|
|
page_size=4,
|
|
cp_size=2,
|
|
cp_rank=1,
|
|
)
|
|
self.assertFalse(plan.compute_padding_enabled)
|
|
forward_batch = SimpleNamespace(
|
|
extend_num_tokens=9,
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
),
|
|
)
|
|
tensor = torch.arange(9 * 2, dtype=torch.float32).view(9, 2)
|
|
expected = split_tensor_by_cp_batch_plan(
|
|
tensor[:8],
|
|
plan,
|
|
mode="data",
|
|
)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local = cp_split_and_rebuild_data(forward_batch, tensor)
|
|
|
|
self.assertTrue(torch.equal(local, expected))
|
|
|
|
def test_cp_split_valid_kind_rejects_trailing_padding_rows(self):
|
|
import torch
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[4],
|
|
prefix_lens=[0],
|
|
page_size=4,
|
|
cp_size=2,
|
|
cp_rank=1,
|
|
)
|
|
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
"batch_gt1_split_input_len_mismatch",
|
|
):
|
|
split_tensor_by_cp_batch_plan(
|
|
torch.arange(8),
|
|
plan,
|
|
mode="1d",
|
|
split_kind="valid",
|
|
)
|
|
|
|
def test_cp_split_and_rebuild_1d_keeps_batch_request_boundaries(self):
|
|
import torch
|
|
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
request_extend_lens=[4, 9],
|
|
request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]],
|
|
request_zigzag_indices=[[1, 2], [1, 2]],
|
|
)
|
|
)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local = cp_split_and_rebuild_1d(forward_batch, torch.arange(13))
|
|
|
|
self.assertEqual(local.tolist(), [8, 9, 10, 11, 12])
|
|
|
|
def test_cp_split_and_rebuild_1d_uses_zero_compute_padding_rows(self):
|
|
import torch
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[40320],
|
|
page_size=64,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
)
|
|
)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local = cp_split_and_rebuild_1d(forward_batch, torch.arange(65))
|
|
|
|
self.assertEqual(local.shape, (64,))
|
|
self.assertEqual(local[0].item(), 64)
|
|
self.assertEqual(local[1:].tolist(), [0] * 63)
|
|
|
|
def test_select_cp_local_valid_rows_filters_compute_padding_rows(self):
|
|
import torch
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[0],
|
|
page_size=64,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
)
|
|
)
|
|
local_compute_rows = torch.full((64, 2), -1.0)
|
|
local_compute_rows[0] = torch.tensor([50.0, 51.0])
|
|
|
|
selected = select_cp_local_valid_rows_for_cache_write(
|
|
forward_batch, local_compute_rows
|
|
)
|
|
|
|
self.assertEqual(selected.tolist(), [[50.0, 51.0]])
|
|
|
|
def test_cp_split_and_rebuild_position_is_batch_aware_and_compute_padded(self):
|
|
import torch
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[40320],
|
|
page_size=64,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
)
|
|
)
|
|
positions = torch.arange(40320, 40385, dtype=torch.int32)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local = cp_split_and_rebuild_position(forward_batch, positions)
|
|
|
|
self.assertEqual(local.shape, (64,))
|
|
self.assertEqual(local.tolist(), list(range(40384, 40448)))
|
|
|
|
def test_cp_split_and_rebuild_position_ignores_mlp_sync_static_padding_without_compute_padding(
|
|
self,
|
|
):
|
|
import torch
|
|
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[8],
|
|
prefix_lens=[0],
|
|
page_size=4,
|
|
cp_size=2,
|
|
cp_rank=1,
|
|
)
|
|
self.assertFalse(plan.compute_padding_enabled)
|
|
forward_batch = SimpleNamespace(
|
|
extend_num_tokens=9,
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
),
|
|
)
|
|
positions = torch.arange(9, dtype=torch.int32)
|
|
expected = split_tensor_by_cp_batch_plan(
|
|
positions[:8],
|
|
plan,
|
|
mode="position",
|
|
)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local = cp_split_and_rebuild_position(forward_batch, positions)
|
|
|
|
self.assertTrue(torch.equal(local, expected))
|
|
|
|
def test_cp_local_embedding_pad_len_uses_metadata_max_rank_len(self):
|
|
from types import SimpleNamespace
|
|
|
|
import torch
|
|
|
|
forward_batch = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4096] * 8)
|
|
)
|
|
|
|
self.assertEqual(
|
|
get_cp_local_embedding_padded_token_count(forward_batch, 4040), 4096
|
|
)
|
|
self.assertEqual(
|
|
get_cp_local_embedding_padded_token_count(forward_batch, 4096), 4096
|
|
)
|
|
self.assertEqual(
|
|
pad_cp_local_input_ids_for_embedding(
|
|
SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[6] * 8)
|
|
),
|
|
torch.tensor([11, 12, 13, 14]),
|
|
).tolist(),
|
|
[11, 12, 13, 14, 0, 0],
|
|
)
|
|
self.assertEqual(
|
|
pad_cp_local_input_ids_for_embedding(
|
|
SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4] * 8)
|
|
),
|
|
torch.tensor([11, 12, 13, 14]),
|
|
).tolist(),
|
|
[11, 12, 13, 14],
|
|
)
|
|
|
|
missing_metadata = SimpleNamespace(nsa_cp_metadata=None)
|
|
self.assertIsNone(
|
|
get_cp_local_embedding_padded_token_count(missing_metadata, 4040)
|
|
)
|
|
self.assertIsNone(
|
|
pad_cp_local_input_ids_for_embedding(
|
|
missing_metadata, torch.tensor([11, 12, 13, 14])
|
|
)
|
|
)
|
|
|
|
stale_metadata = SimpleNamespace(
|
|
nsa_cp_metadata=NSAContextParallelMetadata(max_rank_len=[4039] * 8)
|
|
)
|
|
self.assertIsNone(
|
|
get_cp_local_embedding_padded_token_count(stale_metadata, 4040)
|
|
)
|
|
|
|
def test_local_out_cache_loc_requires_compute_owner_pages(self):
|
|
import torch
|
|
from types import SimpleNamespace
|
|
|
|
page_size = 4
|
|
# Segment order for cp_size=4, cp_rank=1 is segment 1 then 6.
|
|
# The logical page ids below deliberately encode the same owners through
|
|
# (logical_page - 1) % cp_size:
|
|
# segment 1 -> page 2 owner 1
|
|
# segment 6 -> page 6 owner 1
|
|
segment_pages = [1, 2, 3, 4, 8, 7, 6, 5]
|
|
out_cache_loc = torch.cat(
|
|
[
|
|
torch.arange(page * page_size, (page + 1) * page_size)
|
|
for page in segment_pages
|
|
]
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=4,
|
|
cp_rank=1,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
split_list=[page_size] * 8,
|
|
zigzag_index=[1, 6],
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
),
|
|
out_cache_loc=out_cache_loc,
|
|
)
|
|
|
|
local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
|
|
|
self.assertIsNotNone(local_locs)
|
|
self.assertEqual(
|
|
local_locs.tolist(),
|
|
list(range(2 * page_size, 3 * page_size))
|
|
+ list(range(6 * page_size, 7 * page_size)),
|
|
)
|
|
|
|
def test_batch_local_out_cache_loc_keeps_request_boundaries(self):
|
|
import torch
|
|
from types import SimpleNamespace
|
|
|
|
page_size = 4
|
|
# cp_size=2/cp_rank=1 selects segment 1 and 2 for each request.
|
|
# req0 has no rank-1 local rows. req1 contributes segment 1 (page 2)
|
|
# and segment 2 (tail page 4). The synthetic logical page ids encode
|
|
# the owner-lane invariant through (page_id - 1) % cp_size.
|
|
out_cache_loc = torch.cat(
|
|
[
|
|
torch.arange(1 * page_size, 2 * page_size), # req0 seg0 owner 0
|
|
torch.arange(3 * page_size, 4 * page_size), # req1 seg0 owner 0
|
|
torch.arange(2 * page_size, 3 * page_size), # req1 seg1 owner 1
|
|
torch.tensor([4 * page_size]), # req1 seg2 owner 1 tail
|
|
]
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_num_tokens=8,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=2,
|
|
cp_rank=1,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
split_list=[4, 0, 0, 0],
|
|
zigzag_index=[1, 2],
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
request_extend_lens=[4, 9],
|
|
request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]],
|
|
request_zigzag_indices=[[1, 2], [1, 2]],
|
|
),
|
|
out_cache_loc=out_cache_loc,
|
|
)
|
|
|
|
local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
|
|
|
self.assertIsNotNone(local_locs)
|
|
self.assertEqual(
|
|
local_locs.tolist(),
|
|
list(range(2 * page_size, 3 * page_size)) + [4 * page_size],
|
|
)
|
|
|
|
def test_local_out_cache_loc_uses_valid_rows_under_compute_padding(self):
|
|
import torch
|
|
|
|
page_size = 64
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[0],
|
|
page_size=page_size,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
out_cache_loc = torch.cat(
|
|
[
|
|
torch.arange(page_size, 2 * page_size),
|
|
torch.tensor([2 * page_size]),
|
|
]
|
|
).to(torch.int64)
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
),
|
|
out_cache_loc=out_cache_loc,
|
|
)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
|
|
|
self.assertEqual(local_locs.tolist(), [2 * page_size])
|
|
|
|
def test_local_out_cache_loc_ignores_trailing_static_padding_locs(self):
|
|
import torch
|
|
|
|
page_size = 4
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[4],
|
|
prefix_lens=[0],
|
|
page_size=page_size,
|
|
cp_size=2,
|
|
cp_rank=0,
|
|
)
|
|
valid_locs = torch.arange(1 * page_size, 2 * page_size, dtype=torch.int64)
|
|
static_padding_locs = torch.arange(
|
|
99 * page_size, 100 * page_size, dtype=torch.int64
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_num_tokens=8,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=2,
|
|
cp_rank=0,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
),
|
|
out_cache_loc=torch.cat((valid_locs, static_padding_locs)),
|
|
)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
|
|
|
self.assertEqual(local_locs.tolist(), valid_locs.tolist())
|
|
|
|
def test_local_out_cache_loc_ignores_mlp_sync_static_padding_without_compute_padding(
|
|
self,
|
|
):
|
|
import torch
|
|
|
|
page_size = 4
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[8],
|
|
prefix_lens=[0],
|
|
page_size=page_size,
|
|
cp_size=2,
|
|
cp_rank=1,
|
|
)
|
|
self.assertFalse(plan.compute_padding_enabled)
|
|
static_padding_locs = torch.tensor([99 * page_size], dtype=torch.int64)
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
extend_num_tokens=9,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=2,
|
|
cp_rank=1,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
),
|
|
out_cache_loc=torch.cat(
|
|
(
|
|
torch.arange(1 * page_size, 3 * page_size, dtype=torch.int64),
|
|
static_padding_locs,
|
|
)
|
|
),
|
|
)
|
|
|
|
with patch(
|
|
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
|
|
return_value=False,
|
|
):
|
|
local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
|
|
|
self.assertEqual(
|
|
local_locs.tolist(),
|
|
[
|
|
2 * page_size,
|
|
2 * page_size + 1,
|
|
2 * page_size + 2,
|
|
2 * page_size + 3,
|
|
],
|
|
)
|
|
|
|
def test_local_out_cache_loc_rejects_unproven_trailing_padding_even_with_compute_padding(
|
|
self,
|
|
):
|
|
import torch
|
|
|
|
page_size = 4
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[4],
|
|
prefix_lens=[0],
|
|
page_size=page_size,
|
|
cp_size=2,
|
|
cp_rank=0,
|
|
)
|
|
self.assertTrue(plan.compute_padding_enabled)
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=2,
|
|
cp_rank=0,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
),
|
|
out_cache_loc=torch.cat(
|
|
(
|
|
torch.arange(1 * page_size, 2 * page_size, dtype=torch.int64),
|
|
torch.arange(99 * page_size, 100 * page_size, dtype=torch.int64),
|
|
)
|
|
),
|
|
)
|
|
|
|
with self.assertRaisesRegex(RuntimeError, "static_padded=None"):
|
|
get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
|
|
|
def test_batch_local_physical_out_cache_loc_reuses_layer_invariant_plan(self):
|
|
import torch
|
|
from types import SimpleNamespace
|
|
|
|
page_size = 4
|
|
out_cache_loc = torch.cat(
|
|
[
|
|
torch.arange(1 * page_size, 2 * page_size),
|
|
torch.arange(3 * page_size, 4 * page_size),
|
|
torch.arange(2 * page_size, 3 * page_size),
|
|
torch.tensor([4 * page_size]),
|
|
]
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=2,
|
|
cp_rank=1,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
split_list=[4, 0, 0, 0],
|
|
zigzag_index=[1, 2],
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
request_extend_lens=[4, 9],
|
|
request_split_lists=[[4, 0, 0, 0], [4, 4, 1, 0]],
|
|
request_zigzag_indices=[[1, 2], [1, 2]],
|
|
),
|
|
out_cache_loc=out_cache_loc,
|
|
)
|
|
|
|
physical_locs = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)
|
|
second_read = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)
|
|
|
|
self.assertIs(physical_locs, second_read)
|
|
self.assertEqual(
|
|
physical_locs.tolist(),
|
|
list(range(1 * page_size, 2 * page_size)) + [2 * page_size],
|
|
)
|
|
|
|
def test_local_physical_out_cache_loc_is_cached(self):
|
|
import torch
|
|
from types import SimpleNamespace
|
|
|
|
page_size = 4
|
|
segment_pages = [1, 2, 3, 4, 8, 7, 6, 5]
|
|
out_cache_loc = torch.cat(
|
|
[
|
|
torch.arange(page * page_size, (page + 1) * page_size)
|
|
for page in segment_pages
|
|
]
|
|
)
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=4,
|
|
cp_rank=1,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
split_list=[page_size] * 8,
|
|
zigzag_index=[1, 6],
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
),
|
|
out_cache_loc=out_cache_loc,
|
|
)
|
|
|
|
physical_locs = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)
|
|
second_read = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch)
|
|
|
|
self.assertIs(physical_locs, second_read)
|
|
self.assertEqual(
|
|
physical_locs.tolist(),
|
|
list(range(1 * page_size, 2 * page_size))
|
|
+ list(range(2 * page_size, 3 * page_size)),
|
|
)
|
|
|
|
def test_local_out_cache_loc_fails_fast_when_owner_mismatch(self):
|
|
import torch
|
|
from types import SimpleNamespace
|
|
|
|
page_size = 4
|
|
out_cache_loc = torch.arange(page_size * 8, page_size * 16)
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=4,
|
|
cp_rank=1,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
split_list=[page_size] * 8,
|
|
zigzag_index=[1, 6],
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
),
|
|
out_cache_loc=out_cache_loc,
|
|
)
|
|
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
"CP_SHARED_KV_FAIL_FAST.*local_loc_owner_mismatch",
|
|
):
|
|
get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
|
|
|
def test_local_out_cache_loc_fails_fast_every_invalid_event(self):
|
|
import torch
|
|
from types import SimpleNamespace
|
|
|
|
page_size = 4
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=4,
|
|
cp_rank=1,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
split_list=[page_size] * 8,
|
|
zigzag_index=[1, 6],
|
|
page_aligned=False,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
),
|
|
out_cache_loc=torch.arange(page_size * 8, page_size * 16),
|
|
)
|
|
|
|
with self.assertLogs(
|
|
"sglang.srt.layers.attention.nsa.utils", level="ERROR"
|
|
) as cm:
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
"CP_SHARED_KV_FAIL_FAST.*not_page_aligned",
|
|
):
|
|
get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
"CP_SHARED_KV_FAIL_FAST.*not_page_aligned",
|
|
):
|
|
get_cp_shared_kv_local_out_cache_loc(forward_batch)
|
|
|
|
self.assertEqual(len(cm.output), 2)
|
|
self.assertIn("[CP_SHARED_KV_FAIL_FAST][direct_write]", cm.output[0])
|
|
self.assertIn("metadata is not page-aligned", cm.output[0])
|
|
self.assertIn("[CP_SHARED_KV_FAIL_FAST][direct_write]", cm.output[1])
|
|
self.assertIn("metadata is not page-aligned", cm.output[1])
|
|
|
|
def test_indexer_direct_write_fails_fast_on_local_shape_mismatch(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa import nsa_indexer
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
indexer.nsa_enable_prefill_cp = True
|
|
forward_batch = SimpleNamespace(nsa_cp_metadata=object())
|
|
|
|
with (
|
|
patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True),
|
|
patch.object(
|
|
nsa_indexer,
|
|
"get_cp_shared_kv_local_out_cache_loc",
|
|
return_value=torch.tensor([1, 2], dtype=torch.int64),
|
|
),
|
|
):
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
"CP_SHARED_KV_FAIL_FAST.*index_local_shape_mismatch",
|
|
):
|
|
Indexer._store_cp_shared_local_index_k_cache(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=0,
|
|
local_key=torch.empty((1, 8)),
|
|
act_quant=None,
|
|
)
|
|
|
|
def test_indexer_direct_write_filters_compute_padding_rows(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa import nsa_indexer
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
page_size = 64
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[0],
|
|
page_size=page_size,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
out_cache_loc = torch.cat(
|
|
[
|
|
torch.arange(page_size, 2 * page_size),
|
|
torch.tensor([2 * page_size]),
|
|
]
|
|
).to(torch.int64)
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
),
|
|
out_cache_loc=out_cache_loc,
|
|
token_to_kv_pool=SimpleNamespace(page_size=page_size),
|
|
)
|
|
indexer = object.__new__(Indexer)
|
|
indexer.nsa_enable_prefill_cp = True
|
|
calls = []
|
|
|
|
def fake_store_index_k_cache(**kwargs):
|
|
calls.append(kwargs)
|
|
|
|
indexer._store_index_k_cache = fake_store_index_k_cache
|
|
local_key = torch.full((64, 2), -1.0)
|
|
local_key[0] = torch.tensor([9.0, 10.0])
|
|
|
|
with patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True):
|
|
stored = Indexer._store_cp_shared_local_index_k_cache(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=0,
|
|
local_key=local_key,
|
|
act_quant=None,
|
|
)
|
|
|
|
self.assertTrue(stored)
|
|
self.assertEqual(len(calls), 1)
|
|
self.assertEqual(calls[0]["key"].tolist(), [[9.0, 10.0]])
|
|
|
|
def test_mla_direct_write_fails_fast_on_local_shape_mismatch(self):
|
|
import torch
|
|
|
|
from sglang.srt.models.deepseek_common.attention_forward_methods import (
|
|
forward_mla,
|
|
)
|
|
|
|
mla = SimpleNamespace(attn_mqa=SimpleNamespace(layer_id=3))
|
|
forward_batch = SimpleNamespace()
|
|
|
|
with patch.object(
|
|
forward_mla,
|
|
"get_cp_shared_kv_local_out_cache_loc",
|
|
return_value=torch.tensor([1, 2], dtype=torch.int64),
|
|
):
|
|
with self.assertRaisesRegex(
|
|
RuntimeError,
|
|
"CP_SHARED_KV_FAIL_FAST.*mla_local_shape_mismatch",
|
|
):
|
|
forward_mla.DeepseekMLAForwardMixin._maybe_write_cp_shared_local_mla_kv(
|
|
mla,
|
|
forward_batch,
|
|
k_nope=torch.empty((1, 8)),
|
|
k_pe=torch.empty((2, 8)),
|
|
)
|
|
|
|
def test_mla_direct_write_filters_compute_padding_rows(self):
|
|
import torch
|
|
|
|
from sglang.srt.models.deepseek_common.attention_forward_methods import (
|
|
forward_mla,
|
|
)
|
|
|
|
page_size = 64
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[0],
|
|
page_size=page_size,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
out_cache_loc = torch.cat(
|
|
[
|
|
torch.arange(page_size, 2 * page_size),
|
|
torch.tensor([2 * page_size]),
|
|
]
|
|
).to(torch.int64)
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
),
|
|
out_cache_loc=out_cache_loc,
|
|
token_to_kv_pool=SimpleNamespace(page_size=page_size),
|
|
)
|
|
mla = SimpleNamespace(attn_mqa=SimpleNamespace(layer_id=0))
|
|
calls = []
|
|
|
|
def fake_tai_store(**kwargs):
|
|
calls.append(kwargs)
|
|
return True
|
|
|
|
k_nope = torch.full((64, 2), -1.0)
|
|
k_nope[0] = torch.tensor([1.0, 2.0])
|
|
k_pe = torch.full((64, 2), -1.0)
|
|
k_pe[0] = torch.tensor([3.0, 4.0])
|
|
|
|
with patch.object(forward_mla, "try_tai_fused_mla_store", fake_tai_store):
|
|
stored = (
|
|
forward_mla.DeepseekMLAForwardMixin._maybe_write_cp_shared_local_mla_kv(
|
|
mla,
|
|
forward_batch,
|
|
k_nope=k_nope,
|
|
k_pe=k_pe,
|
|
)
|
|
)
|
|
|
|
self.assertTrue(stored)
|
|
self.assertEqual(len(calls), 1)
|
|
self.assertEqual(calls[0]["k_nope"].tolist(), [[1.0, 2.0]])
|
|
self.assertEqual(calls[0]["k_rope"].tolist(), [[3.0, 4.0]])
|
|
self.assertEqual(calls[0]["logical_locs"].tolist(), [2 * page_size])
|
|
|
|
def test_index_partial_current_compose_accepts_local_valid_compute_padding_rows(
|
|
self,
|
|
):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa import nsa_indexer
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
page_size = 64
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[0],
|
|
page_size=page_size,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
|
|
class FakePool:
|
|
page_size = 64
|
|
index_head_dim = 2
|
|
|
|
def get_index_k_with_scale_buffer(self, layer_id):
|
|
return torch.zeros((4, 3), dtype=torch.float32)
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
),
|
|
cp_shared_kv_index_prefetcher=None,
|
|
token_to_kv_pool=FakePool(),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=1,
|
|
batch_plan=plan,
|
|
page_aligned=True,
|
|
page_size=page_size,
|
|
extend_prefix_len=0,
|
|
),
|
|
out_cache_loc=torch.cat(
|
|
[
|
|
torch.arange(page_size, 2 * page_size),
|
|
torch.tensor([2 * page_size]),
|
|
]
|
|
).to(torch.int64),
|
|
extend_prefix_lens_cpu=[page_size],
|
|
extend_seq_lens_cpu=[65],
|
|
)
|
|
logical_page_table = torch.tensor([[1, 2, 3]], dtype=torch.int32)
|
|
current_index_kv = (
|
|
torch.tensor([[7.0, 8.0]], dtype=torch.float32),
|
|
torch.tensor([[0.5]], dtype=torch.float32),
|
|
)
|
|
materialize_calls = []
|
|
expected_buffer = torch.ones((3, 3), dtype=torch.float32)
|
|
expected_pages = torch.tensor([[0, 1, 2]], dtype=torch.int32)
|
|
indexer = object.__new__(Indexer)
|
|
|
|
def fake_materialize(**kwargs):
|
|
materialize_calls.append(kwargs)
|
|
return expected_buffer, expected_pages
|
|
|
|
with patch.object(
|
|
nsa_indexer,
|
|
"get_or_build_shared_paged_buffer_slot_remap",
|
|
return_value=torch.tensor([0, 1, 2], dtype=torch.int64),
|
|
), patch.object(
|
|
nsa_indexer,
|
|
"materialize_prefix_and_reuse_current_index_page_slots",
|
|
side_effect=fake_materialize,
|
|
):
|
|
dense_buffer, dense_pages = indexer._maybe_materialize_shared_index_buffer(
|
|
forward_batch,
|
|
layer_id=0,
|
|
logical_page_table=logical_page_table,
|
|
current_index_kv=current_index_kv,
|
|
)
|
|
|
|
self.assertIs(dense_buffer, expected_buffer)
|
|
self.assertIs(dense_pages, expected_pages)
|
|
self.assertEqual(len(materialize_calls), 1)
|
|
self.assertIs(materialize_calls[0]["current_index_k"], current_index_kv[0])
|
|
self.assertEqual(materialize_calls[0]["current_locs"].tolist(), [2 * page_size])
|
|
|
|
def test_indexer_current_reuse_compute_padding_selects_local_key_not_gathered_key(
|
|
self,
|
|
):
|
|
import torch
|
|
import types
|
|
|
|
from sglang.srt.layers.attention.nsa import nsa_indexer
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
page_size = 64
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[0],
|
|
page_size=page_size,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
metadata_obj = _build_batch_metadata_from_plan(plan)
|
|
|
|
class Mode:
|
|
def is_extend_without_speculative(self):
|
|
return True
|
|
|
|
def is_decode_or_idle(self):
|
|
return False
|
|
|
|
def is_target_verify(self):
|
|
return False
|
|
|
|
def is_draft_extend(self, include_v2=False):
|
|
return False
|
|
|
|
def is_context_parallel_extend(self):
|
|
return True
|
|
|
|
class AttnBackend:
|
|
def get_indexer_metadata(self, layer_id, forward_batch):
|
|
return object()
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(
|
|
page_size=page_size,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
),
|
|
token_to_kv_pool=SimpleNamespace(page_size=page_size),
|
|
nsa_cp_metadata=metadata_obj,
|
|
out_cache_loc=torch.cat(
|
|
[
|
|
torch.arange(page_size, 2 * page_size),
|
|
torch.tensor([2 * page_size]),
|
|
]
|
|
).to(torch.int64),
|
|
extend_prefix_lens_cpu=[0],
|
|
extend_seq_lens_cpu=[65],
|
|
seq_lens_cpu=torch.tensor([65], dtype=torch.int64),
|
|
forward_mode=Mode(),
|
|
attn_backend=AttnBackend(),
|
|
hisparse_coordinator=None,
|
|
)
|
|
|
|
local_key = torch.full((64, 2), -1.0, dtype=torch.float32)
|
|
local_key[0] = torch.tensor([11.0, 12.0])
|
|
gathered_key = torch.full((64, 2), 99.0, dtype=torch.float32)
|
|
gathered_key[0] = torch.tensor([101.0, 102.0])
|
|
query = torch.zeros((64, 2), dtype=torch.float32)
|
|
act_quant_inputs = []
|
|
topk_current_index_kv = []
|
|
|
|
def fake_act_quant(tensor, block_size, scale_fmt):
|
|
act_quant_inputs.append(tensor.detach().clone())
|
|
return tensor.detach().clone(), torch.ones(
|
|
(int(tensor.shape[0]), 1), dtype=torch.float32
|
|
)
|
|
|
|
fake_triton_kernel = types.ModuleType(
|
|
"sglang.srt.layers.attention.nsa.triton_kernel"
|
|
)
|
|
fake_triton_kernel.act_quant = fake_act_quant
|
|
|
|
indexer = object.__new__(Indexer)
|
|
indexer.alt_stream = None
|
|
indexer.nsa_enable_prefill_cp = True
|
|
indexer.index_topk = 2
|
|
indexer.block_size = 64
|
|
indexer.scale_fmt = None
|
|
indexer._get_q_k_bf16 = (
|
|
lambda *args, **kwargs: (query, gathered_key, local_key)
|
|
)
|
|
indexer._store_cp_shared_local_index_k_cache = lambda **kwargs: True
|
|
indexer._can_reuse_current_index_kv = lambda forward_batch: True
|
|
indexer._get_logits_head_gate = (
|
|
lambda x_for_gate, q_scale: torch.zeros((64, 1), dtype=torch.float32)
|
|
)
|
|
|
|
def fake_topk(*args, **kwargs):
|
|
topk_current_index_kv.append(kwargs["current_index_kv"])
|
|
return torch.zeros((64, 2), dtype=torch.int32)
|
|
|
|
indexer._get_topk_in_seq_cp_pair = fake_topk
|
|
|
|
with (
|
|
patch.dict(
|
|
sys.modules,
|
|
{
|
|
"sglang.srt.layers.attention.nsa.triton_kernel": fake_triton_kernel
|
|
},
|
|
),
|
|
patch.object(nsa_indexer, "_is_cuda", True),
|
|
patch.object(nsa_indexer, "_is_hip", False),
|
|
patch.object(nsa_indexer, "_is_npu", False),
|
|
patch.object(
|
|
nsa_indexer,
|
|
"is_nsa_prefill_cp_in_seq_split",
|
|
return_value=True,
|
|
),
|
|
):
|
|
result = Indexer.forward_cuda(
|
|
indexer,
|
|
x=torch.zeros((64, 2), dtype=torch.float32),
|
|
q_lora=torch.zeros((64, 2), dtype=torch.float32),
|
|
positions=torch.arange(64, dtype=torch.int64),
|
|
forward_batch=forward_batch,
|
|
layer_id=0,
|
|
return_indices=True,
|
|
)
|
|
|
|
self.assertEqual(result.shape, (64, 2))
|
|
self.assertEqual(len(act_quant_inputs), 2)
|
|
self.assertEqual(act_quant_inputs[1].tolist(), [[11.0, 12.0]])
|
|
self.assertNotEqual(act_quant_inputs[1].tolist(), [[101.0, 102.0]])
|
|
self.assertEqual(len(topk_current_index_kv), 1)
|
|
self.assertEqual(topk_current_index_kv[0][0].tolist(), [[11.0, 12.0]])
|
|
|
|
def test_indexer_direct_write_does_not_log_missing_metadata_for_non_cp_batch(self):
|
|
import torch
|
|
from types import SimpleNamespace
|
|
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
indexer.nsa_enable_prefill_cp = True
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
nsa_cp_metadata=None,
|
|
)
|
|
|
|
with self.assertNoLogs(
|
|
"sglang.srt.layers.attention.nsa.utils", level="INFO"
|
|
):
|
|
stored = Indexer._store_cp_shared_local_index_k_cache(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=0,
|
|
local_key=torch.empty(0),
|
|
act_quant=None,
|
|
)
|
|
|
|
self.assertFalse(stored)
|
|
|
|
def test_indexer_in_seq_cp_pair_materializes_index_once_for_prev_next(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
logical_pages = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32)
|
|
materialized_index = torch.tensor([11], dtype=torch.int32)
|
|
dense_pages = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32)
|
|
materialize_calls = []
|
|
topk_calls = []
|
|
|
|
class Metadata:
|
|
def get_page_table_64(self):
|
|
return logical_pages
|
|
|
|
def get_page_table_1(self):
|
|
return torch.empty((1, 1000), dtype=torch.int32)
|
|
|
|
def fake_materialize(forward_batch, layer_id, logical_page_table):
|
|
materialize_calls.append((layer_id, logical_page_table))
|
|
return materialized_index, dense_pages
|
|
|
|
def fake_get_topk(
|
|
forward_batch,
|
|
layer_id,
|
|
q_fp8,
|
|
weights,
|
|
metadata,
|
|
kv_len,
|
|
actual_seq_q,
|
|
cp_index=None,
|
|
current_index_kv=None,
|
|
shared_index_buffer=None,
|
|
shared_block_tables=None,
|
|
actual_seq_q_tensor=None,
|
|
actual_seq_q_cu_tensor=None,
|
|
):
|
|
topk_calls.append(
|
|
{
|
|
"kv_len": kv_len,
|
|
"actual_seq_q": actual_seq_q,
|
|
"actual_seq_q_tensor": actual_seq_q_tensor,
|
|
"actual_seq_q_cu_tensor": actual_seq_q_cu_tensor,
|
|
"shared_index_buffer": shared_index_buffer,
|
|
"shared_block_tables": shared_block_tables,
|
|
"current_index_kv": current_index_kv,
|
|
}
|
|
)
|
|
return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32)
|
|
|
|
indexer._maybe_materialize_shared_index_buffer = fake_materialize
|
|
indexer._get_topk_ragged_with_cp = fake_get_topk
|
|
|
|
class Mode:
|
|
def is_extend_without_speculative(self):
|
|
return True
|
|
|
|
forward_batch = type(
|
|
"ForwardBatchStub",
|
|
(),
|
|
{
|
|
"forward_mode": Mode(),
|
|
"extend_prefix_lens_cpu": [0],
|
|
"extend_seq_lens_cpu": [5],
|
|
"seq_lens_cpu": torch.tensor([5], dtype=torch.int64),
|
|
"nsa_cp_metadata": NSAContextParallelMetadata(
|
|
kv_len_prev=5,
|
|
kv_len_next=9,
|
|
actual_seq_q_prev=3,
|
|
actual_seq_q_next=2,
|
|
actual_seq_q_prev_cu_tensor=torch.tensor([0, 3], dtype=torch.int32),
|
|
actual_seq_q_next_cu_tensor=torch.tensor([0, 2], dtype=torch.int32),
|
|
)
|
|
},
|
|
)()
|
|
q_fp8 = torch.arange(5 * 4, dtype=torch.float32).view(5, 4)
|
|
weights = torch.arange(5 * 2, dtype=torch.float32).view(5, 2)
|
|
|
|
result = Indexer._get_topk_in_seq_cp_pair(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
q_fp8=q_fp8,
|
|
weights=weights,
|
|
metadata=Metadata(),
|
|
current_index_kv=None,
|
|
)
|
|
|
|
self.assertEqual(len(materialize_calls), 1)
|
|
self.assertIs(materialize_calls[0][1], logical_pages)
|
|
self.assertEqual(len(topk_calls), 2)
|
|
self.assertIs(topk_calls[0]["shared_index_buffer"], materialized_index)
|
|
self.assertIs(topk_calls[1]["shared_index_buffer"], materialized_index)
|
|
self.assertIs(topk_calls[0]["shared_block_tables"], dense_pages)
|
|
self.assertIs(topk_calls[1]["shared_block_tables"], dense_pages)
|
|
self.assertIsNone(topk_calls[0]["current_index_kv"])
|
|
self.assertEqual(topk_calls[0]["kv_len"], 5)
|
|
self.assertEqual(topk_calls[1]["kv_len"], 9)
|
|
self.assertEqual(topk_calls[0]["actual_seq_q_cu_tensor"].tolist(), [0, 3])
|
|
self.assertEqual(topk_calls[1]["actual_seq_q_cu_tensor"].tolist(), [0, 2])
|
|
self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]])
|
|
|
|
def test_indexer_in_seq_cp_pair_batch_preserves_request_segment_order(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
indexer.index_topk = 2
|
|
logical_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32)
|
|
materialized_index = torch.tensor([11], dtype=torch.int32)
|
|
dense_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32)
|
|
materialize_calls = []
|
|
topk_calls = []
|
|
|
|
class Metadata:
|
|
def get_page_table_64(self):
|
|
return logical_pages
|
|
|
|
def get_page_table_1(self):
|
|
return torch.empty((2, 1000), dtype=torch.int32)
|
|
|
|
def fake_materialize(forward_batch, layer_id, logical_page_table):
|
|
materialize_calls.append((layer_id, logical_page_table))
|
|
return materialized_index, dense_pages
|
|
|
|
def fake_get_topk(
|
|
forward_batch,
|
|
layer_id,
|
|
q_fp8,
|
|
weights,
|
|
metadata,
|
|
kv_len,
|
|
actual_seq_q,
|
|
cp_index=None,
|
|
current_index_kv=None,
|
|
shared_index_buffer=None,
|
|
shared_block_tables=None,
|
|
actual_seq_q_tensor=None,
|
|
actual_seq_q_cu_tensor=None,
|
|
batch_idx=0,
|
|
):
|
|
topk_calls.append(
|
|
{
|
|
"batch_idx": batch_idx,
|
|
"kv_len": kv_len,
|
|
"actual_seq_q": actual_seq_q,
|
|
"cp_index": cp_index,
|
|
"q": q_fp8.flatten().tolist(),
|
|
"weights": weights.flatten().tolist(),
|
|
"actual_seq_q_tensor": actual_seq_q_tensor,
|
|
"actual_seq_q_cu_tensor": actual_seq_q_cu_tensor,
|
|
"shared_index_buffer": shared_index_buffer,
|
|
"shared_block_tables": shared_block_tables,
|
|
"current_index_kv": current_index_kv,
|
|
}
|
|
)
|
|
rows = int(q_fp8.shape[0])
|
|
return torch.arange(1, rows + 1, dtype=torch.int32).view(rows, 1).repeat(1, 2)
|
|
|
|
indexer._maybe_materialize_shared_index_buffer = fake_materialize
|
|
indexer._get_topk_ragged_with_cp = fake_get_topk
|
|
|
|
forward_batch = SimpleNamespace(
|
|
batch_size=2,
|
|
forward_mode=SimpleNamespace(
|
|
is_extend_without_speculative=lambda: True,
|
|
),
|
|
extend_prefix_lens_cpu=[0, 0],
|
|
extend_seq_lens_cpu=[1000, 1000],
|
|
seq_lens_cpu=torch.tensor([1000, 1000], dtype=torch.int64),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
kv_len_prev=100,
|
|
kv_len_next=200,
|
|
actual_seq_q_prev=2,
|
|
actual_seq_q_next=1,
|
|
actual_seq_q_prev_cu_tensor=torch.tensor([0, 2], dtype=torch.int32),
|
|
actual_seq_q_next_cu_tensor=torch.tensor([0, 1], dtype=torch.int32),
|
|
request_kv_len_prev=[2, 1],
|
|
request_kv_len_next=[3, 4],
|
|
request_actual_seq_q_prev=[2, 1],
|
|
request_actual_seq_q_next=[1, 3],
|
|
),
|
|
)
|
|
q_fp8 = torch.arange(7, dtype=torch.float32).view(7, 1)
|
|
weights = (torch.arange(7, dtype=torch.float32) + 100).view(7, 1)
|
|
|
|
result = Indexer._get_topk_in_seq_cp_pair(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
q_fp8=q_fp8,
|
|
weights=weights,
|
|
metadata=Metadata(),
|
|
current_index_kv=None,
|
|
)
|
|
|
|
self.assertEqual(len(materialize_calls), 1)
|
|
self.assertIs(materialize_calls[0][1], logical_pages)
|
|
self.assertEqual(len(topk_calls), 1)
|
|
self.assertEqual(topk_calls[0]["batch_idx"], 0)
|
|
self.assertEqual(topk_calls[0]["actual_seq_q"], 7)
|
|
self.assertEqual(
|
|
topk_calls[0]["cp_index"],
|
|
[(0, 0, 2), (0, 2, 3), (1, 0, 1), (1, 1, 4)],
|
|
)
|
|
self.assertEqual(topk_calls[0]["q"], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
|
|
self.assertEqual(
|
|
topk_calls[0]["weights"],
|
|
[100.0, 101.0, 102.0, 103.0, 104.0, 105.0, 106.0],
|
|
)
|
|
self.assertIsNone(topk_calls[0]["actual_seq_q_tensor"])
|
|
self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"])
|
|
self.assertTrue(all(call["shared_index_buffer"] is materialized_index for call in topk_calls))
|
|
self.assertTrue(all(call["shared_block_tables"] is dense_pages for call in topk_calls))
|
|
self.assertTrue(all(call["current_index_kv"] is None for call in topk_calls))
|
|
self.assertEqual(
|
|
result.tolist(),
|
|
[[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]],
|
|
)
|
|
|
|
def test_indexer_in_seq_cp_pair_compute_padding_outputs_dummy_safe_rows(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
page_size = 64
|
|
plan = build_batch_page_aligned_in_seq_split_plan(
|
|
extend_lens=[65],
|
|
prefix_lens=[0],
|
|
page_size=page_size,
|
|
cp_size=8,
|
|
cp_rank=1,
|
|
)
|
|
metadata_obj = _build_batch_metadata_from_plan(plan)
|
|
indexer = object.__new__(Indexer)
|
|
indexer.index_topk = 2
|
|
logical_pages = torch.tensor([[1, 2]], dtype=torch.int32)
|
|
materialized_index = torch.tensor([11], dtype=torch.int32)
|
|
dense_pages = torch.tensor([[1, 2]], dtype=torch.int32)
|
|
materialize_calls = []
|
|
topk_calls = []
|
|
|
|
class Metadata:
|
|
def get_page_table_64(self):
|
|
return logical_pages
|
|
|
|
def get_page_table_1(self):
|
|
return torch.empty((1, 65), dtype=torch.int32)
|
|
|
|
def fake_materialize(forward_batch, layer_id, logical_page_table):
|
|
materialize_calls.append((layer_id, logical_page_table))
|
|
return materialized_index, dense_pages
|
|
|
|
def fake_get_topk(
|
|
forward_batch,
|
|
layer_id,
|
|
q_fp8,
|
|
weights,
|
|
metadata,
|
|
kv_len,
|
|
actual_seq_q,
|
|
cp_index=None,
|
|
current_index_kv=None,
|
|
shared_index_buffer=None,
|
|
shared_block_tables=None,
|
|
actual_seq_q_tensor=None,
|
|
actual_seq_q_cu_tensor=None,
|
|
batch_idx=0,
|
|
):
|
|
topk_calls.append(
|
|
{
|
|
"actual_seq_q": actual_seq_q,
|
|
"cp_index": cp_index,
|
|
"q": q_fp8.flatten().tolist(),
|
|
"weights": weights.flatten().tolist(),
|
|
"shared_index_buffer": shared_index_buffer,
|
|
"shared_block_tables": shared_block_tables,
|
|
}
|
|
)
|
|
rows = int(q_fp8.shape[0])
|
|
return (
|
|
torch.arange(1, rows + 1, dtype=torch.int32)
|
|
.view(rows, 1)
|
|
.repeat(1, 2)
|
|
)
|
|
|
|
indexer._maybe_materialize_shared_index_buffer = fake_materialize
|
|
indexer._get_topk_ragged_with_cp = fake_get_topk
|
|
|
|
forward_batch = SimpleNamespace(
|
|
batch_size=1,
|
|
forward_mode=SimpleNamespace(
|
|
is_extend_without_speculative=lambda: True,
|
|
),
|
|
extend_prefix_lens_cpu=[0],
|
|
extend_seq_lens_cpu=[65],
|
|
seq_lens_cpu=torch.tensor([65], dtype=torch.int64),
|
|
nsa_cp_metadata=metadata_obj,
|
|
)
|
|
q_fp8 = torch.arange(64, dtype=torch.float32).view(64, 1)
|
|
weights = (torch.arange(64, dtype=torch.float32) + 100).view(64, 1)
|
|
|
|
result = Indexer._get_topk_in_seq_cp_pair(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
q_fp8=q_fp8,
|
|
weights=weights,
|
|
metadata=Metadata(),
|
|
current_index_kv=None,
|
|
)
|
|
|
|
self.assertEqual(len(materialize_calls), 1)
|
|
self.assertIs(materialize_calls[0][1], logical_pages)
|
|
self.assertEqual(len(topk_calls), 1)
|
|
self.assertEqual(topk_calls[0]["actual_seq_q"], 1)
|
|
self.assertEqual(topk_calls[0]["cp_index"], [(0, 64, 65)])
|
|
self.assertEqual(topk_calls[0]["q"], [0.0])
|
|
self.assertEqual(topk_calls[0]["weights"], [100.0])
|
|
self.assertIs(topk_calls[0]["shared_index_buffer"], materialized_index)
|
|
self.assertIs(topk_calls[0]["shared_block_tables"], dense_pages)
|
|
self.assertEqual(result.shape, (64, 2))
|
|
self.assertEqual(result[0].tolist(), [1, 1])
|
|
self.assertTrue(
|
|
torch.equal(result[1:], torch.full((63, 2), -1, dtype=torch.int32))
|
|
)
|
|
|
|
def test_indexer_in_seq_cp_pair_batch_materializes_partial_current_index_reuse_once(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
indexer.index_topk = 2
|
|
current_index_kv = (torch.arange(7), torch.arange(7))
|
|
logical_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32)
|
|
materialized_index = torch.tensor([11], dtype=torch.int32)
|
|
dense_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32)
|
|
materialize_calls = []
|
|
topk_calls = []
|
|
|
|
class Metadata:
|
|
def get_page_table_64(self):
|
|
return logical_pages
|
|
|
|
def get_page_table_1(self):
|
|
return torch.empty((2, 512), dtype=torch.int32)
|
|
|
|
def fake_materialize(
|
|
forward_batch,
|
|
layer_id,
|
|
logical_page_table,
|
|
current_index_kv=None,
|
|
):
|
|
materialize_calls.append(
|
|
{
|
|
"layer_id": layer_id,
|
|
"logical_page_table": logical_page_table,
|
|
"current_index_kv": current_index_kv,
|
|
}
|
|
)
|
|
return materialized_index, dense_pages
|
|
|
|
def fake_get_topk(
|
|
forward_batch,
|
|
layer_id,
|
|
q_fp8,
|
|
weights,
|
|
metadata,
|
|
kv_len,
|
|
actual_seq_q,
|
|
cp_index=None,
|
|
current_index_kv=None,
|
|
shared_index_buffer=None,
|
|
shared_block_tables=None,
|
|
actual_seq_q_tensor=None,
|
|
actual_seq_q_cu_tensor=None,
|
|
batch_idx=0,
|
|
):
|
|
topk_calls.append(
|
|
{
|
|
"batch_idx": batch_idx,
|
|
"actual_seq_q": actual_seq_q,
|
|
"cp_index": cp_index,
|
|
"q_rows": int(q_fp8.shape[0]),
|
|
"current_index_kv": current_index_kv,
|
|
"shared_index_buffer": shared_index_buffer,
|
|
"shared_block_tables": shared_block_tables,
|
|
"actual_seq_q_cu_tensor": actual_seq_q_cu_tensor,
|
|
}
|
|
)
|
|
rows = int(q_fp8.shape[0])
|
|
return torch.arange(1, rows + 1, dtype=torch.int32).view(rows, 1).repeat(1, 2)
|
|
|
|
indexer._maybe_materialize_shared_index_buffer = fake_materialize
|
|
indexer._get_topk_ragged_with_cp = fake_get_topk
|
|
|
|
forward_batch = SimpleNamespace(
|
|
batch_size=2,
|
|
forward_mode=SimpleNamespace(
|
|
is_extend_without_speculative=lambda: True,
|
|
),
|
|
extend_prefix_lens_cpu=[64, 64],
|
|
extend_seq_lens_cpu=[3, 4],
|
|
seq_lens_cpu=torch.tensor([67, 68], dtype=torch.int64),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
request_kv_len_prev=[2, 1],
|
|
request_kv_len_next=[3, 4],
|
|
request_actual_seq_q_prev=[2, 1],
|
|
request_actual_seq_q_next=[1, 3],
|
|
),
|
|
)
|
|
result = Indexer._get_topk_in_seq_cp_pair(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
q_fp8=torch.empty(7, 1),
|
|
weights=torch.empty(7, 1),
|
|
metadata=Metadata(),
|
|
current_index_kv=current_index_kv,
|
|
)
|
|
|
|
self.assertEqual(len(materialize_calls), 1)
|
|
self.assertIs(materialize_calls[0]["logical_page_table"], logical_pages)
|
|
self.assertIs(materialize_calls[0]["current_index_kv"], current_index_kv)
|
|
self.assertEqual(len(topk_calls), 1)
|
|
self.assertEqual(topk_calls[0]["actual_seq_q"], 7)
|
|
self.assertEqual(
|
|
topk_calls[0]["cp_index"],
|
|
[(0, 0, 2), (0, 2, 3), (1, 0, 1), (1, 1, 4)],
|
|
)
|
|
self.assertEqual(topk_calls[0]["q_rows"], 7)
|
|
self.assertTrue(
|
|
all(call["current_index_kv"] is None for call in topk_calls)
|
|
)
|
|
self.assertTrue(
|
|
all(call["shared_index_buffer"] is materialized_index for call in topk_calls)
|
|
)
|
|
self.assertTrue(
|
|
all(call["shared_block_tables"] is dense_pages for call in topk_calls)
|
|
)
|
|
self.assertEqual([call["batch_idx"] for call in topk_calls], [0])
|
|
self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"])
|
|
self.assertEqual(
|
|
result.tolist(),
|
|
[[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]],
|
|
)
|
|
|
|
def test_indexer_in_seq_cp_pair_batch_composes_current_only_index_reuse(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
indexer.index_topk = 2
|
|
current_index_kv = (torch.arange(7), torch.arange(7))
|
|
logical_pages = torch.tensor([[1, 2], [3, 4]], dtype=torch.int32)
|
|
materialized_index = torch.tensor([17], dtype=torch.int32)
|
|
dense_pages = torch.tensor([[10, 11], [12, 13]], dtype=torch.int32)
|
|
materialize_calls = []
|
|
topk_calls = []
|
|
|
|
class Mode:
|
|
def is_extend_without_speculative(self):
|
|
return True
|
|
|
|
class Metadata:
|
|
def get_page_table_64(self):
|
|
return logical_pages
|
|
|
|
def get_page_table_1(self):
|
|
return torch.empty((2, 512), dtype=torch.int32)
|
|
|
|
def fake_materialize(
|
|
forward_batch,
|
|
layer_id,
|
|
logical_page_table,
|
|
current_index_kv=None,
|
|
):
|
|
materialize_calls.append(
|
|
{
|
|
"layer_id": layer_id,
|
|
"logical_page_table": logical_page_table,
|
|
"current_index_kv": current_index_kv,
|
|
}
|
|
)
|
|
return materialized_index, dense_pages
|
|
|
|
def fake_get_topk(
|
|
forward_batch,
|
|
layer_id,
|
|
q_fp8,
|
|
weights,
|
|
metadata,
|
|
kv_len,
|
|
actual_seq_q,
|
|
cp_index=None,
|
|
current_index_kv=None,
|
|
shared_index_buffer=None,
|
|
shared_block_tables=None,
|
|
actual_seq_q_tensor=None,
|
|
actual_seq_q_cu_tensor=None,
|
|
batch_idx=0,
|
|
):
|
|
topk_calls.append(
|
|
{
|
|
"batch_idx": batch_idx,
|
|
"actual_seq_q": actual_seq_q,
|
|
"cp_index": cp_index,
|
|
"q_rows": int(q_fp8.shape[0]),
|
|
"current_index_kv": current_index_kv,
|
|
"shared_index_buffer": shared_index_buffer,
|
|
"shared_block_tables": shared_block_tables,
|
|
"actual_seq_q_cu_tensor": actual_seq_q_cu_tensor,
|
|
}
|
|
)
|
|
rows = int(q_fp8.shape[0])
|
|
return torch.arange(1, rows + 1, dtype=torch.int32).view(rows, 1).repeat(1, 2)
|
|
|
|
indexer._maybe_materialize_shared_index_buffer = fake_materialize
|
|
indexer._get_topk_ragged_with_cp = fake_get_topk
|
|
|
|
forward_batch = SimpleNamespace(
|
|
batch_size=2,
|
|
forward_mode=Mode(),
|
|
extend_prefix_lens_cpu=[0, 0],
|
|
extend_seq_lens_cpu=[3, 4],
|
|
seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64),
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
request_kv_len_prev=[3, 4],
|
|
request_kv_len_next=[3, 4],
|
|
request_actual_seq_q_prev=[2, 1],
|
|
request_actual_seq_q_next=[1, 3],
|
|
),
|
|
)
|
|
result = Indexer._get_topk_in_seq_cp_pair(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
q_fp8=torch.empty(7, 1),
|
|
weights=torch.empty(7, 1),
|
|
metadata=Metadata(),
|
|
current_index_kv=current_index_kv,
|
|
)
|
|
|
|
self.assertEqual(len(materialize_calls), 1)
|
|
self.assertIs(materialize_calls[0]["logical_page_table"], logical_pages)
|
|
self.assertIs(materialize_calls[0]["current_index_kv"], current_index_kv)
|
|
self.assertEqual(len(topk_calls), 1)
|
|
self.assertTrue(all(call["current_index_kv"] is None for call in topk_calls))
|
|
self.assertTrue(
|
|
all(call["shared_index_buffer"] is materialized_index for call in topk_calls)
|
|
)
|
|
self.assertTrue(
|
|
all(call["shared_block_tables"] is dense_pages for call in topk_calls)
|
|
)
|
|
self.assertEqual(topk_calls[0]["batch_idx"], 0)
|
|
self.assertEqual(topk_calls[0]["actual_seq_q"], 7)
|
|
self.assertEqual(
|
|
topk_calls[0]["cp_index"],
|
|
[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)],
|
|
)
|
|
self.assertEqual(topk_calls[0]["q_rows"], 7)
|
|
self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"])
|
|
self.assertEqual(
|
|
result.tolist(),
|
|
[[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]],
|
|
)
|
|
|
|
def test_indexer_shared_index_materialize_accepts_current_only_compose(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa import nsa_indexer
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
page_buffer = torch.zeros((8, 32), dtype=torch.uint8)
|
|
logical_pages = torch.tensor([[1, 2]], dtype=torch.int64)
|
|
current_index_kv = (
|
|
torch.zeros((2, 4), dtype=torch.uint8),
|
|
torch.zeros((2, 1), dtype=torch.float32),
|
|
)
|
|
compose_calls = []
|
|
|
|
class Pool:
|
|
page_size = 4
|
|
index_head_dim = 4
|
|
|
|
def get_index_k_with_scale_buffer(self, layer_id):
|
|
return page_buffer
|
|
|
|
class Mode:
|
|
def is_extend_without_speculative(self):
|
|
return True
|
|
|
|
forward_batch = SimpleNamespace(
|
|
uses_cp_shared_kv=True,
|
|
cp_shared_kv_layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0),
|
|
token_to_kv_pool=Pool(),
|
|
cp_shared_kv_index_prefetcher=None,
|
|
forward_mode=Mode(),
|
|
extend_prefix_lens_cpu=[0],
|
|
extend_seq_lens_cpu=[5],
|
|
seq_lens_cpu=torch.tensor([5], dtype=torch.int64),
|
|
cp_local_out_cache_loc=torch.tensor([4, 5], dtype=torch.int64),
|
|
)
|
|
|
|
def fake_compose(**kwargs):
|
|
compose_calls.append(kwargs)
|
|
return torch.empty((3, 32), dtype=torch.uint8), torch.tensor([[1, 2]])
|
|
|
|
with patch.object(
|
|
nsa_indexer,
|
|
"materialize_prefix_and_reuse_current_index_page_slots",
|
|
side_effect=fake_compose,
|
|
):
|
|
materialized, dense_pages = Indexer._maybe_materialize_shared_index_buffer(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
logical_page_table=logical_pages,
|
|
current_index_kv=current_index_kv,
|
|
)
|
|
|
|
self.assertEqual(len(compose_calls), 1)
|
|
self.assertEqual(compose_calls[0]["prefix_pages"], 0)
|
|
self.assertEqual(compose_calls[0]["current_slot_spans"], [(0, 2)])
|
|
self.assertIs(compose_calls[0]["current_index_k"], current_index_kv[0])
|
|
self.assertIs(compose_calls[0]["current_index_scale"], current_index_kv[1])
|
|
self.assertEqual(list(materialized.shape), [3, 32])
|
|
self.assertEqual(dense_pages.tolist(), [[1, 2]])
|
|
|
|
def test_indexer_ragged_cp_index_current_batch_does_not_materialize(self):
|
|
import contextlib
|
|
import torch
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from sglang.srt.layers.attention.nsa import nsa_indexer
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
indexer.index_topk = 2
|
|
indexer._with_real_sm_count = lambda: contextlib.nullcontext()
|
|
|
|
deep_gemm_calls = []
|
|
|
|
def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False):
|
|
deep_gemm_calls.append(
|
|
{
|
|
"q_rows": int(q_fp8.shape[0]),
|
|
"kv_rows": int(kv_fp8[0].shape[0]),
|
|
"weights_rows": int(weights.shape[0]),
|
|
"ks": ks.tolist(),
|
|
"ke": ke.tolist(),
|
|
}
|
|
)
|
|
return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32)
|
|
|
|
class Metadata:
|
|
def get_page_table_64(self):
|
|
raise AssertionError("current cp_index path must not materialize index pages")
|
|
|
|
def topk_transform(self, logits, topk, **kwargs):
|
|
return (
|
|
torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32)
|
|
.view(-1, 1)
|
|
.repeat(1, topk)
|
|
)
|
|
|
|
forward_batch = SimpleNamespace(
|
|
token_to_kv_pool=SimpleNamespace(page_size=64),
|
|
seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64),
|
|
extend_seq_lens_cpu=[3, 4],
|
|
)
|
|
q_fp8 = torch.empty((7, 1), dtype=torch.float32)
|
|
weights = torch.empty((7, 1, 1), dtype=torch.float32)
|
|
current_index_kv = (
|
|
torch.arange(7, dtype=torch.uint8).view(7, 1),
|
|
torch.arange(7, dtype=torch.float32).view(7, 1),
|
|
)
|
|
|
|
with patch.object(
|
|
nsa_indexer,
|
|
"deep_gemm",
|
|
SimpleNamespace(fp8_mqa_logits=fake_logits),
|
|
):
|
|
result = Indexer._get_topk_ragged_with_cp(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
q_fp8=q_fp8,
|
|
weights=weights,
|
|
metadata=Metadata(),
|
|
kv_len=0,
|
|
actual_seq_q=7,
|
|
cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)],
|
|
current_index_kv=current_index_kv,
|
|
)
|
|
|
|
self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]])
|
|
self.assertEqual(len(deep_gemm_calls), 1)
|
|
self.assertEqual(deep_gemm_calls[0]["q_rows"], 7)
|
|
self.assertEqual(deep_gemm_calls[0]["weights_rows"], 7)
|
|
self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14)
|
|
self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10])
|
|
self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14])
|
|
|
|
def test_indexer_ragged_cp_index_batch_uses_request_ragged_offsets(self):
|
|
import contextlib
|
|
import torch
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from sglang.srt.layers.attention.nsa import nsa_indexer
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
indexer.index_topk = 2
|
|
indexer._with_real_sm_count = lambda: contextlib.nullcontext()
|
|
|
|
topk_kwargs = []
|
|
|
|
def fake_prepare(**kwargs):
|
|
total_kv_len = int(kwargs["total_kv_len"])
|
|
return (
|
|
torch.zeros((total_kv_len, 1), dtype=torch.uint8),
|
|
torch.zeros((total_kv_len,), dtype=torch.float32),
|
|
torch.tensor([0, 0, 3, 6, 10, 10, 10], dtype=torch.int32),
|
|
torch.tensor([2, 3, 3, 4, 2, 3, 4], dtype=torch.int32),
|
|
)
|
|
|
|
def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False):
|
|
return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32)
|
|
|
|
class Metadata:
|
|
def get_page_table_64(self):
|
|
raise AssertionError("current cp_index path must not materialize index pages")
|
|
|
|
def topk_transform(self, logits, topk, **kwargs):
|
|
topk_kwargs.append(kwargs)
|
|
return torch.zeros((int(logits.shape[0]), topk), dtype=torch.int32)
|
|
|
|
forward_batch = SimpleNamespace(
|
|
token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1),
|
|
seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64),
|
|
extend_seq_lens_cpu=[3, 4],
|
|
)
|
|
current_index_kv = (
|
|
torch.arange(7, dtype=torch.uint8).view(7, 1),
|
|
torch.arange(7, dtype=torch.float32).view(7, 1),
|
|
)
|
|
|
|
with patch.object(
|
|
nsa_indexer,
|
|
"try_tai_prepare_cp_mqa_current_index_batch",
|
|
side_effect=fake_prepare,
|
|
create=True,
|
|
), patch.object(
|
|
nsa_indexer,
|
|
"deep_gemm",
|
|
SimpleNamespace(fp8_mqa_logits=fake_logits),
|
|
):
|
|
Indexer._get_topk_ragged_with_cp(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
q_fp8=torch.empty((7, 1), dtype=torch.float32),
|
|
weights=torch.empty((7, 1, 1), dtype=torch.float32),
|
|
metadata=Metadata(),
|
|
kv_len=0,
|
|
actual_seq_q=7,
|
|
cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)],
|
|
current_index_kv=current_index_kv,
|
|
)
|
|
|
|
self.assertEqual(len(topk_kwargs), 1)
|
|
offset = topk_kwargs[0].get("topk_indices_offset_override")
|
|
self.assertIsNotNone(offset)
|
|
# Batch cp_index compacts each CP segment's K into a temporary buffer,
|
|
# but flashmla_sparse consumes the normal ragged KV layout. The fused
|
|
# ragged topk offset therefore has to stay in request/KV coordinates,
|
|
# not compact segment coordinates or compact-q cu-seqlens.
|
|
self.assertEqual(offset.tolist(), [0, 0, 0, 3, 3, 3, 3])
|
|
|
|
def test_indexer_ragged_cp_index_shared_batch_uses_tai_prepare_once(self):
|
|
import contextlib
|
|
import torch
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from sglang.srt.layers.attention.nsa import index_buf_accessor, nsa_indexer
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
indexer.index_topk = 2
|
|
indexer._with_real_sm_count = lambda: contextlib.nullcontext()
|
|
|
|
prepare_calls = []
|
|
deep_gemm_calls = []
|
|
|
|
def fake_prepare(**kwargs):
|
|
prepare_calls.append(kwargs)
|
|
total_kv_len = int(kwargs["total_kv_len"])
|
|
total_q_count = int(kwargs["total_q_count"])
|
|
return (
|
|
torch.zeros((total_kv_len, 1), dtype=torch.uint8),
|
|
torch.zeros((total_kv_len,), dtype=torch.float32),
|
|
torch.tensor([0, 0, 3, 6, 10, 10, 10], dtype=torch.int32),
|
|
torch.tensor([2, 3, 3, 4, 2, 3, 4], dtype=torch.int32),
|
|
)
|
|
|
|
def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False):
|
|
deep_gemm_calls.append(
|
|
{
|
|
"q_rows": int(q_fp8.shape[0]),
|
|
"kv_rows": int(kv_fp8[0].shape[0]),
|
|
"weights_rows": int(weights.shape[0]),
|
|
"ks": ks.tolist(),
|
|
"ke": ke.tolist(),
|
|
}
|
|
)
|
|
return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32)
|
|
|
|
class Metadata:
|
|
def topk_transform(self, logits, topk, **kwargs):
|
|
return (
|
|
torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32)
|
|
.view(-1, 1)
|
|
.repeat(1, topk)
|
|
)
|
|
|
|
forward_batch = SimpleNamespace(
|
|
token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1),
|
|
seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64),
|
|
extend_seq_lens_cpu=[3, 4],
|
|
)
|
|
q_fp8 = torch.empty((7, 1), dtype=torch.float32)
|
|
weights = torch.empty((7, 1, 1), dtype=torch.float32)
|
|
shared_index_buffer = torch.zeros((8, 264), dtype=torch.uint8)
|
|
shared_block_tables = torch.arange(8, dtype=torch.int64).view(2, 4)
|
|
|
|
with patch.object(
|
|
nsa_indexer,
|
|
"try_tai_prepare_cp_mqa_index_batch",
|
|
side_effect=fake_prepare,
|
|
create=True,
|
|
), patch.object(
|
|
index_buf_accessor.GetK,
|
|
"execute",
|
|
side_effect=AssertionError("batched path must not call per-segment GetK"),
|
|
), patch.object(
|
|
index_buf_accessor.GetS,
|
|
"execute",
|
|
side_effect=AssertionError("batched path must not call per-segment GetS"),
|
|
), patch.object(
|
|
nsa_indexer,
|
|
"deep_gemm",
|
|
SimpleNamespace(fp8_mqa_logits=fake_logits),
|
|
):
|
|
result = Indexer._get_topk_ragged_with_cp(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
q_fp8=q_fp8,
|
|
weights=weights,
|
|
metadata=Metadata(),
|
|
kv_len=0,
|
|
actual_seq_q=7,
|
|
cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)],
|
|
shared_index_buffer=shared_index_buffer,
|
|
shared_block_tables=shared_block_tables,
|
|
)
|
|
|
|
self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]])
|
|
self.assertEqual(len(prepare_calls), 1)
|
|
call = prepare_calls[0]
|
|
self.assertIs(call["index_buffer"], shared_index_buffer)
|
|
self.assertIs(call["block_tables"], shared_block_tables)
|
|
self.assertEqual(call["batch_indices"].dtype, torch.int32)
|
|
self.assertEqual(call["batch_indices"].tolist(), [0, 0, 1, 1])
|
|
self.assertEqual(call["kv_lens"].dtype, torch.int32)
|
|
self.assertEqual(call["kv_lens"].tolist(), [3, 3, 4, 4])
|
|
self.assertEqual(call["q_starts"].tolist(), [1, 2, 3, 1])
|
|
self.assertEqual(call["q_lens"].tolist(), [2, 1, 1, 3])
|
|
self.assertEqual(call["k_bases"].tolist(), [0, 3, 6, 10])
|
|
self.assertEqual(call["q_bases"].tolist(), [0, 2, 3, 4])
|
|
self.assertEqual(call["total_kv_len"], 14)
|
|
self.assertEqual(call["total_q_count"], 7)
|
|
self.assertEqual(call["max_kv_len"], 4)
|
|
self.assertEqual(call["max_q_len"], 3)
|
|
self.assertEqual(len(deep_gemm_calls), 1)
|
|
self.assertEqual(deep_gemm_calls[0]["q_rows"], 7)
|
|
self.assertEqual(deep_gemm_calls[0]["weights_rows"], 7)
|
|
self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14)
|
|
self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10])
|
|
self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14])
|
|
|
|
def test_indexer_ragged_cp_index_current_batch_uses_tai_compact_once(self):
|
|
import contextlib
|
|
import torch
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from sglang.srt.layers.attention.nsa import nsa_indexer
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
indexer.index_topk = 2
|
|
indexer._with_real_sm_count = lambda: contextlib.nullcontext()
|
|
|
|
prepare_calls = []
|
|
deep_gemm_calls = []
|
|
|
|
def fake_prepare(**kwargs):
|
|
prepare_calls.append(kwargs)
|
|
total_kv_len = int(kwargs["total_kv_len"])
|
|
return (
|
|
torch.zeros((total_kv_len, 1), dtype=torch.uint8),
|
|
torch.zeros((total_kv_len,), dtype=torch.float32),
|
|
torch.tensor([0, 0, 3, 6, 10, 10, 10], dtype=torch.int32),
|
|
torch.tensor([2, 3, 3, 4, 2, 3, 4], dtype=torch.int32),
|
|
)
|
|
|
|
def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False):
|
|
deep_gemm_calls.append(
|
|
{
|
|
"q_rows": int(q_fp8.shape[0]),
|
|
"kv_rows": int(kv_fp8[0].shape[0]),
|
|
"ks": ks.tolist(),
|
|
"ke": ke.tolist(),
|
|
}
|
|
)
|
|
return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32)
|
|
|
|
class Metadata:
|
|
def get_page_table_64(self):
|
|
raise AssertionError("current cp_index path must not materialize index pages")
|
|
|
|
def topk_transform(self, logits, topk, **kwargs):
|
|
return (
|
|
torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32)
|
|
.view(-1, 1)
|
|
.repeat(1, topk)
|
|
)
|
|
|
|
forward_batch = SimpleNamespace(
|
|
token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1),
|
|
seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64),
|
|
extend_seq_lens_cpu=[3, 4],
|
|
)
|
|
current_index_kv = (
|
|
torch.arange(7, dtype=torch.uint8).view(7, 1),
|
|
torch.arange(7, dtype=torch.float32).view(7, 1),
|
|
)
|
|
|
|
with patch.object(
|
|
nsa_indexer,
|
|
"try_tai_prepare_cp_mqa_current_index_batch",
|
|
side_effect=fake_prepare,
|
|
create=True,
|
|
), patch.object(
|
|
nsa_indexer,
|
|
"deep_gemm",
|
|
SimpleNamespace(fp8_mqa_logits=fake_logits),
|
|
):
|
|
result = Indexer._get_topk_ragged_with_cp(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
q_fp8=torch.empty((7, 1), dtype=torch.float32),
|
|
weights=torch.empty((7, 1, 1), dtype=torch.float32),
|
|
metadata=Metadata(),
|
|
kv_len=0,
|
|
actual_seq_q=7,
|
|
cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)],
|
|
current_index_kv=current_index_kv,
|
|
)
|
|
|
|
self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]])
|
|
self.assertEqual(len(prepare_calls), 1)
|
|
call = prepare_calls[0]
|
|
self.assertIs(call["current_index_k"], current_index_kv[0])
|
|
self.assertIs(call["current_index_scale"], current_index_kv[1])
|
|
self.assertEqual(call["current_bases"].tolist(), [0, 0, 3, 3])
|
|
self.assertEqual(call["kv_lens"].tolist(), [3, 3, 4, 4])
|
|
self.assertEqual(call["q_starts"].tolist(), [1, 2, 3, 1])
|
|
self.assertEqual(call["q_lens"].tolist(), [2, 1, 1, 3])
|
|
self.assertEqual(call["k_bases"].tolist(), [0, 3, 6, 10])
|
|
self.assertEqual(call["q_bases"].tolist(), [0, 2, 3, 4])
|
|
self.assertEqual(call["total_kv_len"], 14)
|
|
self.assertEqual(call["total_q_count"], 7)
|
|
self.assertEqual(len(deep_gemm_calls), 1)
|
|
self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14)
|
|
self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10])
|
|
self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14])
|
|
|
|
def test_indexer_ragged_cp_index_current_batch_uses_cp_local_bases_and_uint8_k(
|
|
self,
|
|
):
|
|
import contextlib
|
|
import torch
|
|
from types import SimpleNamespace
|
|
from unittest.mock import patch
|
|
|
|
from sglang.srt.layers.attention.nsa import nsa_indexer
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
indexer.index_topk = 2
|
|
indexer._with_real_sm_count = lambda: contextlib.nullcontext()
|
|
|
|
prepare_calls = []
|
|
|
|
def fake_prepare(**kwargs):
|
|
prepare_calls.append(kwargs)
|
|
total_kv_len = int(kwargs["total_kv_len"])
|
|
return (
|
|
torch.zeros((total_kv_len, 1), dtype=torch.uint8),
|
|
torch.zeros((total_kv_len,), dtype=torch.float32),
|
|
torch.tensor([0, 0, 3, 3, 3], dtype=torch.int32),
|
|
torch.tensor([2, 3, 2, 3, 4], dtype=torch.int32),
|
|
)
|
|
|
|
class Metadata:
|
|
def get_page_table_64(self):
|
|
raise AssertionError("current cp_index path must not materialize index pages")
|
|
|
|
def topk_transform(self, logits, topk, **kwargs):
|
|
return (
|
|
torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32)
|
|
.view(-1, 1)
|
|
.repeat(1, topk)
|
|
)
|
|
|
|
forward_batch = SimpleNamespace(
|
|
token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1),
|
|
seq_lens_cpu=torch.tensor([5, 7], dtype=torch.int64),
|
|
extend_seq_lens_cpu=[5, 7],
|
|
nsa_cp_metadata=NSAContextParallelMetadata(
|
|
batch_size=2,
|
|
batch_plan=SimpleNamespace(
|
|
request_rank_local_offsets=[0, 2],
|
|
request_valid_rank_local_offsets=[0, 2],
|
|
),
|
|
),
|
|
)
|
|
current_index_k = (
|
|
torch.arange(6, dtype=torch.float32)
|
|
.to(torch.float8_e4m3fn)
|
|
.view(6, 1)
|
|
)
|
|
current_index_scale = torch.arange(6, dtype=torch.float32).view(6, 1)
|
|
current_index_kv = (current_index_k, current_index_scale)
|
|
|
|
with patch.object(
|
|
nsa_indexer,
|
|
"try_tai_prepare_cp_mqa_current_index_batch",
|
|
side_effect=fake_prepare,
|
|
create=True,
|
|
), patch.object(
|
|
nsa_indexer,
|
|
"deep_gemm",
|
|
SimpleNamespace(
|
|
fp8_mqa_logits=lambda q_fp8, kv_fp8, weights, ks, ke, clean_logits=False: torch.zeros(
|
|
(int(q_fp8.shape[0]), 8), dtype=torch.float32
|
|
)
|
|
),
|
|
):
|
|
result = Indexer._get_topk_ragged_with_cp(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
q_fp8=torch.empty((5, 1), dtype=torch.float32),
|
|
weights=torch.empty((5, 1, 1), dtype=torch.float32),
|
|
metadata=Metadata(),
|
|
kv_len=0,
|
|
actual_seq_q=5,
|
|
cp_index=[(0, 1, 3), (1, 1, 4)],
|
|
current_index_kv=current_index_kv,
|
|
)
|
|
|
|
self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5]])
|
|
self.assertEqual(len(prepare_calls), 1)
|
|
call = prepare_calls[0]
|
|
self.assertEqual(call["current_index_k"].dtype, torch.uint8)
|
|
self.assertEqual(call["current_bases"].tolist(), [0, 2])
|
|
self.assertEqual(call["kv_lens"].tolist(), [3, 4])
|
|
self.assertEqual(call["q_lens"].tolist(), [2, 3])
|
|
|
|
def test_eagle_capture_for_decode_clears_cp_local_hidden_marker(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
|
|
from sglang.srt.speculative.eagle_worker import EAGLEWorker
|
|
|
|
worker = object.__new__(EAGLEWorker)
|
|
worker.topk = 1
|
|
draft_input = EagleDraftInput(
|
|
hidden_states=torch.full((4, 2), -1.0),
|
|
cp_local_hidden_states=True,
|
|
)
|
|
draft_output_hidden = torch.tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
logits_output = LogitsProcessorOutput(
|
|
next_token_logits=torch.tensor([[0.1, 0.9], [0.7, 0.3]]),
|
|
hidden_states=draft_output_hidden,
|
|
)
|
|
|
|
worker.capture_for_decode(logits_output, draft_input)
|
|
|
|
self.assertIs(draft_input.hidden_states, draft_output_hidden)
|
|
self.assertFalse(draft_input.cp_local_hidden_states)
|
|
|
|
def test_indexer_in_seq_cp_pair_composes_current_only_index_reuse(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
|
|
|
|
indexer = object.__new__(Indexer)
|
|
current_index_kv = (torch.tensor([1]), torch.tensor([2]))
|
|
logical_pages = torch.tensor([[1, 2]], dtype=torch.int32)
|
|
materialized_index = torch.tensor([7], dtype=torch.int32)
|
|
dense_pages = torch.tensor([[3, 4]], dtype=torch.int32)
|
|
materialize_calls = []
|
|
topk_calls = []
|
|
|
|
class Metadata:
|
|
def get_page_table_64(self):
|
|
return logical_pages
|
|
|
|
def get_page_table_1(self):
|
|
return torch.empty((1, 512), dtype=torch.int32)
|
|
|
|
def fake_materialize(
|
|
forward_batch,
|
|
layer_id,
|
|
logical_page_table,
|
|
current_index_kv=None,
|
|
):
|
|
materialize_calls.append(
|
|
{
|
|
"layer_id": layer_id,
|
|
"logical_page_table": logical_page_table,
|
|
"current_index_kv": current_index_kv,
|
|
}
|
|
)
|
|
return materialized_index, dense_pages
|
|
|
|
def fake_get_topk(
|
|
forward_batch,
|
|
layer_id,
|
|
q_fp8,
|
|
weights,
|
|
metadata,
|
|
kv_len,
|
|
actual_seq_q,
|
|
cp_index=None,
|
|
current_index_kv=None,
|
|
shared_index_buffer=None,
|
|
shared_block_tables=None,
|
|
actual_seq_q_tensor=None,
|
|
actual_seq_q_cu_tensor=None,
|
|
):
|
|
topk_calls.append(
|
|
{
|
|
"current_index_kv": current_index_kv,
|
|
"actual_seq_q_tensor": actual_seq_q_tensor,
|
|
"actual_seq_q_cu_tensor": actual_seq_q_cu_tensor,
|
|
"shared_index_buffer": shared_index_buffer,
|
|
"shared_block_tables": shared_block_tables,
|
|
}
|
|
)
|
|
return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32)
|
|
|
|
indexer._maybe_materialize_shared_index_buffer = fake_materialize
|
|
indexer._get_topk_ragged_with_cp = fake_get_topk
|
|
|
|
class Mode:
|
|
def is_extend_without_speculative(self):
|
|
return True
|
|
|
|
forward_batch = type(
|
|
"ForwardBatchStub",
|
|
(),
|
|
{
|
|
"forward_mode": Mode(),
|
|
"extend_prefix_lens_cpu": [0],
|
|
"extend_seq_lens_cpu": [5],
|
|
"seq_lens_cpu": torch.tensor([5], dtype=torch.int64),
|
|
"nsa_cp_metadata": NSAContextParallelMetadata(
|
|
kv_len_prev=5,
|
|
kv_len_next=9,
|
|
actual_seq_q_prev=3,
|
|
actual_seq_q_next=2,
|
|
actual_seq_q_prev_cu_tensor=torch.tensor([0, 3], dtype=torch.int32),
|
|
actual_seq_q_next_cu_tensor=torch.tensor([0, 2], dtype=torch.int32),
|
|
)
|
|
},
|
|
)()
|
|
|
|
result = Indexer._get_topk_in_seq_cp_pair(
|
|
indexer,
|
|
forward_batch,
|
|
layer_id=7,
|
|
q_fp8=torch.empty(5, 4),
|
|
weights=torch.empty(5, 2),
|
|
metadata=Metadata(),
|
|
current_index_kv=current_index_kv,
|
|
)
|
|
|
|
self.assertEqual(len(materialize_calls), 1)
|
|
self.assertIs(materialize_calls[0]["logical_page_table"], logical_pages)
|
|
self.assertIs(materialize_calls[0]["current_index_kv"], current_index_kv)
|
|
self.assertEqual(len(topk_calls), 2)
|
|
self.assertIsNone(topk_calls[0]["current_index_kv"])
|
|
self.assertIsNone(topk_calls[1]["current_index_kv"])
|
|
self.assertIs(topk_calls[0]["shared_index_buffer"], materialized_index)
|
|
self.assertIs(topk_calls[1]["shared_block_tables"], dense_pages)
|
|
self.assertEqual(topk_calls[0]["actual_seq_q_cu_tensor"].tolist(), [0, 3])
|
|
self.assertEqual(topk_calls[1]["actual_seq_q_cu_tensor"].tolist(), [0, 2])
|
|
self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]])
|
|
|
|
def test_paged_topk_transform_uses_cu_override_without_scan_metadata_ops(self):
|
|
import torch
|
|
|
|
from sglang.srt.layers.attention.nsa_backend import (
|
|
NSAMetadata,
|
|
NSAIndexerMetadata,
|
|
TopkTransformMethod,
|
|
)
|
|
|
|
cu_override = torch.tensor([0, 4], dtype=torch.int32)
|
|
attn_metadata = NSAMetadata(
|
|
page_size=64,
|
|
cache_seqlens_int32=torch.tensor([4], dtype=torch.int32),
|
|
max_seq_len_q=4,
|
|
max_seq_len_k=8,
|
|
cu_seqlens_q=torch.tensor([0, 4], dtype=torch.int32),
|
|
cu_seqlens_k=torch.tensor([0, 8], dtype=torch.int32),
|
|
page_table_1=torch.arange(8, dtype=torch.int32).view(1, 8),
|
|
real_page_table=torch.arange(8, dtype=torch.int32).view(1, 8),
|
|
nsa_cache_seqlens_int32=torch.tensor([4], dtype=torch.int32),
|
|
nsa_cu_seqlens_q=torch.arange(2, dtype=torch.int32),
|
|
nsa_cu_seqlens_k=torch.tensor([0, 4], dtype=torch.int32),
|
|
nsa_extend_seq_lens_list=[4],
|
|
nsa_seqlens_expanded=torch.arange(1, 5, dtype=torch.int32),
|
|
topk_indices_offset=torch.zeros(4, dtype=torch.int32),
|
|
)
|
|
metadata = NSAIndexerMetadata(
|
|
attn_metadata=attn_metadata,
|
|
topk_transform_method=TopkTransformMethod.PAGED,
|
|
)
|
|
logits = torch.zeros(4, 8)
|
|
lengths = torch.arange(1, 5, dtype=torch.int32)
|
|
expected = torch.full((4, 2), 7, dtype=torch.int32)
|
|
|
|
def fake_fused(**kwargs):
|
|
self.assertIs(kwargs["cu_seqlens_q"], cu_override)
|
|
self.assertIs(kwargs["lengths"], lengths)
|
|
return expected
|
|
|
|
fake_sgl_kernel = SimpleNamespace(
|
|
fast_topk_transform_fused=fake_fused,
|
|
fast_topk_transform_ragged_fused=lambda **_: (_ for _ in ()).throw(
|
|
AssertionError("ragged path should not run")
|
|
),
|
|
fast_topk_v2=lambda *_, **__: (_ for _ in ()).throw(
|
|
AssertionError("unfused path should not run")
|
|
),
|
|
)
|
|
|
|
with (
|
|
patch.dict(sys.modules, {"sgl_kernel": fake_sgl_kernel}),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa_backend.envs.SGLANG_NSA_FUSE_TOPK.get",
|
|
return_value=True,
|
|
),
|
|
patch(
|
|
"sglang.srt.layers.attention.nsa_backend.compute_cu_seqlens",
|
|
side_effect=AssertionError("paged override should skip cumsum"),
|
|
),
|
|
patch(
|
|
"torch.repeat_interleave",
|
|
side_effect=AssertionError("paged topk should not build ragged offsets"),
|
|
),
|
|
):
|
|
actual = metadata.topk_transform(
|
|
logits,
|
|
topk=2,
|
|
cu_seqlens_q=torch.tensor([4], dtype=torch.int32),
|
|
ke_offset=lengths,
|
|
cu_seqlens_q_topk_override=cu_override,
|
|
)
|
|
|
|
self.assertIs(actual, expected)
|
|
|
|
def test_mla_partial_current_path_fails_fast_instead_of_compact_fallback(self):
|
|
backend_path = (
|
|
Path(__file__).resolve().parents[4]
|
|
/ "python/sglang/srt/layers/attention/nsa_backend.py"
|
|
)
|
|
source = backend_path.read_text()
|
|
tree = ast.parse(source)
|
|
compact_merge_calls = [
|
|
node
|
|
for node in ast.walk(tree)
|
|
if isinstance(node, ast.Call)
|
|
and isinstance(node.func, ast.Name)
|
|
and node.func.id == "merge_materialized_and_current_kv"
|
|
]
|
|
|
|
self.assertEqual(
|
|
compact_merge_calls,
|
|
[],
|
|
"MLA partial-current reuse must not fall back to compact "
|
|
"materialize/current merge when page-slot prefetch compose is unavailable.",
|
|
)
|
|
self.assertIn(
|
|
"[CP_SHARED_KV_FAIL_FAST][mla_partial_current_sync]",
|
|
source,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|