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

[NPU] Support DeepSeek-V4 DSpark and refactor DSV4 cache management (#33676)

Browse Source 
Co-authored-by: JiaruiChang5268 <jc5268@columbia.edu>
Co-authored-by: Kelon <kelonlu@163.com>
Co-authored-by: unknown <z8ruev42yk@gmail.com>
Co-authored-by: Talantan1102 <545811257@qq.com>
Co-authored-by: Talantan1102 <44429302+Talantan1102@users.noreply.github.com>
This commit is contained in:
Zed
2026-08-17 16:27:44 +08:00
committed by GitHub
co-authored by JiaruiChang5268 Kelon unknown Talantan1102 Talantan1102
parent e03c53fc13
commit b83d507cd7
37 changed files with 2010 additions and 2018 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/unit/npu/attention/test_npu_ascend_dsv4_backend.py
-73
View File
@@ -59,7 +59,6 @@ from sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend import (
DeepseekV4AscendMultiStepDraftBackend,
_apply_hadamard,
_get_kv_indices,
_overlap_transform,
_walsh_hadamard_matrix,
)
@@ -182,78 +181,6 @@ class TestApplyHadamard(unittest.TestCase):
self.assertTrue(torch.equal(out, expected))
class TestOverlapTransform(unittest.TestCase):
def test_shape(self):
# (n_chunks, ratio, 2*d) -> (n_chunks, 2*ratio, d)
n_chunks, r, d = 3, 2, 4
tensor = torch.randn(n_chunks, r, 2 * d)
out = _overlap_transform(tensor, value=0.0, head_dim=d)
self.assertEqual(out.shape, (n_chunks, 2 * r, d))
def test_first_chunk_left_half_filled_with_value(self):
n_chunks, r, d = 3, 2, 4
tensor = torch.randn(n_chunks, r, 2 * d)
fill = float("-inf")
out = _overlap_transform(tensor, value=fill, head_dim=d)
self.assertTrue(torch.equal(out[0, :r], torch.full((r, d), fill)))
def test_first_chunk_left_half_filled_with_zero(self):
n_chunks, r, d = 2, 2, 4
tensor = torch.randn(n_chunks, r, 2 * d)
out = _overlap_transform(tensor, value=0.0, head_dim=d)
self.assertTrue(torch.equal(out[0, :r], torch.zeros(r, d)))
def test_right_half_mirrors_tensor_second_half(self):
n_chunks, r, d = 3, 2, 4
tensor = torch.randn(n_chunks, r, 2 * d)
out = _overlap_transform(tensor, value=0.0, head_dim=d)
self.assertTrue(torch.equal(out[:, r:], tensor[..., d:]))
def test_previous_chunk_left_half(self):
n_chunks, r, d = 3, 2, 4
tensor = torch.randn(n_chunks, r, 2 * d)
out = _overlap_transform(tensor, value=0.0, head_dim=d)
self.assertTrue(torch.equal(out[1:, :r], tensor[:-1, :, :d]))
def test_single_chunk(self):
n_chunks, r, d = 1, 2, 4
tensor = torch.randn(n_chunks, r, 2 * d)
fill = 7.0
out = _overlap_transform(tensor, value=fill, head_dim=d)
self.assertEqual(out.shape, (1, 2 * r, d))
self.assertTrue(torch.equal(out[0, :r], torch.full((r, d), fill)))
self.assertTrue(torch.equal(out[0, r:], tensor[0, :, d:]))
def test_full_element_mapping(self):
n_chunks, r, d = 2, 2, 3
tensor = torch.arange(n_chunks * r * 2 * d, dtype=torch.float32).reshape(
n_chunks, r, 2 * d
)
fill = -1.0
out = _overlap_transform(tensor, value=fill, head_dim=d)
for c in range(n_chunks):
for row in range(2 * r):
for col in range(d):
if c == 0 and row < r:
expected = fill
elif row >= r:
expected = tensor[c, row - r, d + col].item()
else:
expected = tensor[c - 1, row, col].item()
self.assertEqual(
out[c, row, col].item(),
expected,
f"mismatch at (c={c}, row={row}, col={col})",
)
def test_preserves_input_dtype(self):
n_chunks, r, d = 2, 2, 4
tensor = torch.randn(n_chunks, r, 2 * d, dtype=torch.bfloat16)
out = _overlap_transform(tensor, value=0.0, head_dim=d)
self.assertEqual(out.dtype, torch.bfloat16)
class TestGetKvIndices(unittest.TestCase):
_PATCH_TARGET = (
"sglang.srt.hardware_backend.npu.attention.ascend_dsv4_backend.get_attn_backend"
test/registered/unit/spec/test_decode_bookkeeping_ownership.py
+9
View File
@@ -41,6 +41,10 @@ _EVICT_METHOD = "maybe_evict_swa"
# Any added/removed/recounted site fails until reviewed here.
_SB = "managers/schedule_batch.py"
_EAGLE_DECODE = ("speculative/eagle_utils.py", "eagle_prepare_for_decode")
_DFLASH_DECODE = (
"speculative/dflash_info_v2.py",
"DFlashDraftInputV2.prepare_for_decode",
)
_RESOLVE = (
"managers/scheduler_components/batch_result_processor.py",
"SchedulerBatchResultProcessor._resolve_spec_v2_tokens",
@@ -62,6 +66,11 @@ _OWNER_SITES = {
# inside the owned-kv alloc_for_spec_decode function (op42).
(*_EAGLE_DECODE, "decode_batch_idx"): 1,
(*_EAGLE_DECODE, "evict"): 1,
# DFlash uses its stateful scheduler-side preparation instead of
# eagle_prepare_for_decode. spec_prepare_for_decode dispatches to exactly
# one of these two owners for each speculative decode iteration.
(*_DFLASH_DECODE, "decode_batch_idx"): 1,
(*_DFLASH_DECODE, "evict"): 1,
(
"mem_cache/allocation.py",
"alloc_for_spec_decode",