kvcache-ai/sglang
1
0
Fork 0
mirror of https://github.com/kvcache-ai/sglang.git synced 2026-07-10 17:27:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Deepseek V3.2] Clean up MTP (#13236)

Browse Source
This commit is contained in:
hlu1
2025-11-14 11:01:37 -08:00
committed by GitHub
parent 84e151ac7f
commit a7002e614b
3 changed files with 69 additions and 77 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -290,7 +290,10 @@ class Indexer(CustomOp):
)
blocksize = page_size
if forward_batch.forward_mode.is_target_verify():
if (
forward_batch.forward_mode.is_target_verify()
or forward_batch.forward_mode.is_draft_extend()
):
seqlens_32 = metadata.get_seqlens_expanded()
else:
seqlens_32 = metadata.get_seqlens_int32()
@@ -337,6 +340,8 @@ class Indexer(CustomOp):
if TYPE_CHECKING:
assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool)
assert forward_batch.forward_mode.is_extend_without_speculative()
page_size = forward_batch.token_to_kv_pool.page_size
assert page_size == 64, "only support page size 64"
assert len(weights.shape) == 3
@@ -649,6 +654,7 @@ class Indexer(CustomOp):
if (
forward_batch.forward_mode.is_decode_or_idle()
or forward_batch.forward_mode.is_target_verify()
or forward_batch.forward_mode.is_draft_extend()
):
topk_result = self._get_topk_paged(
forward_batch, layer_id, q_fp8, weights, metadata

134
python/sglang/srt/layers/attention/nsa_backend.py
View File

@@ -137,6 +137,9 @@ class NSAIndexerMetadata(BaseIndexerMetadata):
def get_seqlens_expanded(self) -> torch.Tensor:
return self.attn_metadata.nsa_seqlens_expanded
def get_cu_seqlens_k(self) -> torch.Tensor:
return self.attn_metadata.cu_seqlens_k
def topk_transform(
self,
logits: torch.Tensor,
@@ -304,8 +307,7 @@ class NativeSparseAttnBackend(AttentionBackend):
cu_seqlens_q = self.get_device_int32_arange(batch_size + 1)
seqlens_expanded = cache_seqlens_int32
elif forward_batch.forward_mode.is_target_verify():
max_seqlen_q = self.speculative_num_draft_tokens
nsa_max_seqlen_q = self.speculative_num_draft_tokens
max_seqlen_q = 1
cu_seqlens_q = torch.arange(
0,
batch_size * self.speculative_num_draft_tokens + 1,
@@ -338,6 +340,43 @@ class NativeSparseAttnBackend(AttentionBackend):
page_table = torch.repeat_interleave(
page_table, repeats=self.speculative_num_draft_tokens, dim=0
)
elif forward_batch.forward_mode.is_draft_extend():
assert (
forward_batch.extend_seq_lens_cpu is not None
and forward_batch.extend_seq_lens is not None
and forward_batch.extend_prefix_lens_cpu is not None
), "All of them must not be None"
extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu
assert forward_batch.extend_seq_lens is not None
max_seqlen_q = 1
cu_seqlens_q = torch.arange(
0,
forward_batch.extend_num_tokens + 1,
1,
dtype=torch.int32,
device=device,
)
seqlens_expanded = torch.cat(
[
torch.arange(
kv_len - qo_len + 1,
kv_len + 1,
dtype=torch.int32,
device=device,
)
for qo_len, kv_len in zip(
forward_batch.extend_seq_lens_cpu,
forward_batch.seq_lens_cpu.tolist(),
strict=True,
)
]
)
page_table = torch.repeat_interleave(
page_table, repeats=forward_batch.extend_seq_lens, dim=0
)
elif forward_batch.forward_mode.is_extend():
assert (
forward_batch.extend_seq_lens_cpu is not None
@@ -518,7 +557,7 @@ class NativeSparseAttnBackend(AttentionBackend):
)
else:
flashmla_metadata = None
elif forward_mode.is_target_verify():
elif forward_mode.is_target_verify() or forward_mode.is_draft_extend():
cache_seqlens_int32 = (seq_lens + self.speculative_num_draft_tokens).to(
torch.int32
)
@@ -576,64 +615,6 @@ class NativeSparseAttnBackend(AttentionBackend):
)
else:
flashmla_metadata = None
elif forward_mode.is_draft_extend():
cache_seqlens_int32 = (seq_lens + self.speculative_num_draft_tokens).to(
torch.int32
)
cu_seqlens_k = compute_cu_seqlens(cache_seqlens_int32)
page_table_1 = self.decode_cuda_graph_metadata["page_table"][:bs, :]
max_seqlen_k = page_table_1.shape[1]
extend_seq_lens_cpu = [self.speculative_num_draft_tokens] * bs
extend_seq_lens = torch.full(
(bs,),
self.speculative_num_draft_tokens,
device=self.device,
dtype=torch.int32,
)
max_seqlen_q = max(extend_seq_lens_cpu)
cu_seqlens_q = compute_cu_seqlens(extend_seq_lens.to(torch.int32))
seqlens_int32_cpu = [
self.speculative_num_draft_tokens + kv_len
for kv_len in seq_lens.tolist()
]
seqlens_expanded = torch.cat(
[
torch.arange(
kv_len - qo_len + 1,
kv_len + 1,
dtype=torch.int32,
device=self.device,
)
for qo_len, kv_len in zip(
extend_seq_lens_cpu,
seqlens_int32_cpu,
strict=True,
)
]
)
nsa_cache_seqlens_int32 = compute_nsa_seqlens(
seqlens_expanded, nsa_index_topk=self.nsa_index_topk
)
nsa_extend_seq_lens_list = [1] * bs
if NSA_DECODE_IMPL == "flashmla_kv":
flashmla_metadata = self.decode_cuda_graph_metadata[
"flashmla_metadata"
].slice(slice(0, bs * self.speculative_num_draft_tokens + 1))
# As the DeepGemm is not support for q_len = 3/4 in Indexer and every token has independent topk_indices,
# we made the Q shape [bs * speculative_num_draft_tokens, 1, head_nums, dim].
# So seq_len_q is 1 for flashmla_metadata in target_verify and draft_extend mode.
flashmla_metadata.copy_(
self._compute_flashmla_metadata(
cache_seqlens=nsa_cache_seqlens_int32,
seq_len_q=1,
)
)
else:
flashmla_metadata = None
nsa_cu_seqlens_k = compute_cu_seqlens(nsa_cache_seqlens_int32)
nsa_cu_seqlens_q = self.get_device_int32_arange(len(nsa_cu_seqlens_k))
@@ -747,14 +728,18 @@ class NativeSparseAttnBackend(AttentionBackend):
metadata.cu_seqlens_k[1:].copy_(
torch.cumsum(cache_seqlens, dim=0, dtype=torch.int32)
)
page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
metadata.page_table_1[:, :max_seqlen_k].copy_(page_indices)
extend_seq_lens_cpu = spec_info.accept_length[:bs].tolist()
seqlens_int32_cpu = [
self.speculative_num_draft_tokens + kv_len
for kv_len in seq_lens_cpu.tolist()
]
extend_seq_lens = spec_info.accept_length[:bs]
extend_seq_lens_cpu = extend_seq_lens.tolist()
page_indices = self.req_to_token[req_pool_indices, :max_seqlen_k]
page_indices = torch.repeat_interleave(
page_indices, repeats=extend_seq_lens, dim=0
)
metadata.page_table_1[: page_indices.shape[0], :max_seqlen_k].copy_(
page_indices
)
seqlens_expanded = torch.cat(
[
torch.arange(
@@ -765,21 +750,21 @@ class NativeSparseAttnBackend(AttentionBackend):
)
for qo_len, kv_len in zip(
extend_seq_lens_cpu,
seqlens_int32_cpu,
seq_lens_cpu.tolist(),
strict=True,
)
]
)
metadata.nsa_seqlens_expanded[: seqlens_expanded.size(0)].copy_(
metadata.nsa_seqlens_expanded[: seqlens_expanded.shape[0]].copy_(
seqlens_expanded
)
nsa_cache_seqlens = compute_nsa_seqlens(
seqlens_expanded, self.nsa_index_topk
)
metadata.nsa_cache_seqlens_int32[: seqlens_expanded.size(0)].copy_(
metadata.nsa_cache_seqlens_int32[: seqlens_expanded.shape[0]].copy_(
nsa_cache_seqlens
)
seqlens_expanded_size = seqlens_expanded.size(0)
seqlens_expanded_size = seqlens_expanded.shape[0]
assert (
metadata.nsa_cache_seqlens_int32 is not None
and metadata.nsa_cu_seqlens_k is not None
@@ -794,8 +779,9 @@ class NativeSparseAttnBackend(AttentionBackend):
assert self.real_page_size == metadata.page_size
if self.real_page_size > 1:
real_table = self._transform_table_1_to_real(page_indices)
new_len = real_table.shape[1]
metadata.real_page_table[:, :new_len].copy_(real_table)
new_rows = real_table.shape[0]
new_cols = real_table.shape[1]
metadata.real_page_table[:new_rows, :new_cols].copy_(real_table)
else:
assert metadata.real_page_table is metadata.page_table_1