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

Remove retired DSA env paths (#29912)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-07-05 22:58:02 -07:00
committed by GitHub
parent 6f22790943
commit 5eb1b6a7ba
7 changed files with 136 additions and 584 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/environ.py
-5
View File
@@ -627,9 +627,6 @@ class Envs:
)
SGLANG_DSA_TOPK_FLASHINFER_DETERMINISTIC = EnvBool(False)
SGLANG_DSA_TOPK_FLASHINFER_TIE_BREAK = EnvStr(None)
SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA = EnvBoolWithAlias(
True, deprecated_name="SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA"
)
SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD = EnvIntWithAlias(
2048, deprecated_name="SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD"
)
@@ -638,8 +635,6 @@ class Envs:
)
SGLANG_DSA_MQA_LOGITS_FREE_MEM_FRACTION = EnvFloat(0.2)
SGLANG_ENABLE_PCG_DSV2_DUAL_STREAM = EnvBool(False)
SGLANG_USE_FUSED_METADATA_COPY = EnvBool(True)
SGLANG_DSA_USE_FUSED_METADATA_GENERATION = EnvBool(True)
SGLANG_DSA_TOPK_BROADCAST = EnvBool(False)
SGLANG_DISABLE_DSA_INDEXER_FUSION = EnvBool(False)
python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py
+2 -6
View File
@@ -11,7 +11,6 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.environ import envs
from sglang.srt.layers.attention.dsa.utils import compute_dsa_seqlens
from sglang.srt.layers.attention.utils import seqlens_expand_triton
from sglang.srt.utils import is_cuda, is_hip
@@ -21,9 +20,6 @@ if TYPE_CHECKING:
_is_cuda = is_cuda()
_is_hip = is_hip()
_USE_FUSED_METADATA_GENERATION = (
envs.SGLANG_DSA_USE_FUSED_METADATA_GENERATION.get() and not _is_hip
)
@dataclass
@@ -125,7 +121,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
"""Precompute metadata for normal decode mode."""
max_len = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1]
if _USE_FUSED_METADATA_GENERATION and _is_cuda:
if _is_cuda and not _is_hip:
from sglang.srt.layers.attention.triton_ops.dsa_metadata import (
fused_dsa_decode_metadata,
)
@@ -244,7 +240,7 @@ class DeepseekSparseAttnBackendMTPPrecomputeMixin:
max_seqlen_k = self.decode_cuda_graph_metadata[bs].page_table_1.shape[1]
seqlens_expanded_size = bs * self.speculative_num_draft_tokens
if _USE_FUSED_METADATA_GENERATION and _is_cuda:
if _is_cuda and not _is_hip:
from sglang.srt.layers.attention.triton_ops.dsa_metadata import (
fused_dsa_target_verify_metadata,
)
python/sglang/srt/layers/attention/dsa/dsa_mtp_verification.py
-391
View File
@@ -1,391 +0,0 @@
"""
Verification utilities for DSA backend fused metadata copy operations.
This module contains verification code to ensure that fused metadata copy kernels
produce the same results as individual copy operations.
"""
import torch
def verify_single_backend_fused_metadata_copy(
metadata,
precomputed,
forward_mode,
bs,
flashmla_num_splits_src=None,
flashmla_metadata_src=None,
flashmla_num_splits_dst=None,
flashmla_metadata_dst=None,
):
"""
Verify that the fused metadata copy kernel produces the same results as individual copies.
Args:
metadata: The DSA metadata object containing destination tensors
precomputed: The precomputed metadata containing source tensors
forward_mode: The forward mode (decode, target_verify, or draft_extend)
bs: Batch size
flashmla_num_splits_src: Source FlashMLA num_splits tensor (optional)
flashmla_metadata_src: Source FlashMLA metadata tensor (optional)
flashmla_num_splits_dst: Destination FlashMLA num_splits tensor (optional)
flashmla_metadata_dst: Destination FlashMLA metadata tensor (optional)
Raises:
RuntimeError: If verification fails (tensors don't match)
"""
# Clone destination tensors to preserve fused kernel results
fused_cache_seqlens = metadata.cache_seqlens_int32.clone()
fused_cu_seqlens_k = metadata.cu_seqlens_k.clone()
fused_page_table_1 = metadata.page_table_1.clone()
fused_dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32.clone()
fused_dsa_seqlens_expanded = metadata.dsa_seqlens_expanded.clone()
fused_dsa_cu_seqlens_k = metadata.dsa_cu_seqlens_k.clone()
fused_real_page_table = (
metadata.real_page_table.clone()
if precomputed.real_page_table is not None
else None
)
fused_flashmla_num_splits = None
fused_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
fused_flashmla_num_splits = flashmla_num_splits_dst.clone()
fused_flashmla_metadata = flashmla_metadata_dst.clone()
# Create reference tensors (zeroed out)
ref_cache_seqlens = torch.zeros_like(metadata.cache_seqlens_int32)
ref_cu_seqlens_k = torch.zeros_like(metadata.cu_seqlens_k)
ref_page_table_1 = torch.zeros_like(metadata.page_table_1)
ref_dsa_cache_seqlens = torch.zeros_like(metadata.dsa_cache_seqlens_int32)
ref_dsa_seqlens_expanded = torch.zeros_like(metadata.dsa_seqlens_expanded)
ref_dsa_cu_seqlens_k = torch.zeros_like(metadata.dsa_cu_seqlens_k)
ref_real_page_table = (
torch.zeros_like(metadata.real_page_table)
if precomputed.real_page_table is not None
else None
)
ref_flashmla_num_splits = None
ref_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits = torch.zeros_like(flashmla_num_splits_dst)
ref_flashmla_metadata = torch.zeros_like(flashmla_metadata_dst)
# Run individual copy operations (reference implementation)
ref_cache_seqlens.copy_(precomputed.cache_seqlens)
ref_cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:])
if forward_mode.is_decode_or_idle():
# Decode mode
ref_page_table_1[:, : precomputed.max_len].copy_(precomputed.page_indices)
ref_dsa_cache_seqlens.copy_(precomputed.dsa_cache_seqlens)
elif forward_mode.is_target_verify():
# Target verify mode
ref_page_table_1[:, : precomputed.max_seqlen_k].copy_(precomputed.page_indices)
ref_dsa_seqlens_expanded.copy_(precomputed.seqlens_expanded)
ref_dsa_cache_seqlens.copy_(precomputed.dsa_cache_seqlens)
# Copy DSA cu_seqlens
size = precomputed.seqlens_expanded_size
ref_dsa_cu_seqlens_k[1 : 1 + size].copy_(precomputed.dsa_cu_seqlens_k[1 : 1 + size])
# Copy real page table
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
ref_real_page_table[:rows, :cols].copy_(precomputed.real_page_table)
# Copy FlashMLA metadata
if precomputed.flashmla_metadata is not None:
size = precomputed.seqlens_expanded_size
ref_flashmla_num_splits[: size + 1].copy_(flashmla_num_splits_src[: size + 1])
ref_flashmla_metadata.copy_(flashmla_metadata_src)
# Compare results and crash if inconsistent
def check_tensor_equal(name, fused, ref):
if not torch.equal(fused, ref):
max_diff = (fused.float() - ref.float()).abs().max().item()
mismatched_elements = (fused != ref).sum().item()
total_elements = fused.numel()
raise RuntimeError(
f"FUSED METADATA COPY VERIFICATION FAILED!\n"
f"Tensor: {name}\n"
f"Max difference: {max_diff}\n"
f"Mismatched elements: {mismatched_elements}/{total_elements}\n"
f"Fused shape: {fused.shape}, Ref shape: {ref.shape}\n"
f"Forward mode: {forward_mode}, bs={bs}\n"
f"The fused kernel produces different results than individual copies.\n"
f"This indicates a bug in the fused metadata copy kernel."
)
# Verify all tensors (only compare the slices that were actually updated)
check_tensor_equal("cache_seqlens", fused_cache_seqlens, ref_cache_seqlens)
check_tensor_equal("cu_seqlens_k", fused_cu_seqlens_k, ref_cu_seqlens_k)
# Compare page_table_1 only for the region that was updated
if forward_mode.is_decode_or_idle():
check_tensor_equal(
"page_table_1",
fused_page_table_1[:, : precomputed.max_len],
ref_page_table_1[:, : precomputed.max_len],
)
elif forward_mode.is_target_verify():
check_tensor_equal(
"page_table_1",
fused_page_table_1[:, : precomputed.max_seqlen_k],
ref_page_table_1[:, : precomputed.max_seqlen_k],
)
# Compare dsa_cache_seqlens only for the region that was updated
if forward_mode.is_decode_or_idle():
check_tensor_equal(
"dsa_cache_seqlens",
fused_dsa_cache_seqlens,
ref_dsa_cache_seqlens,
)
else: # TARGET_VERIFY
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"dsa_cache_seqlens",
fused_dsa_cache_seqlens[:size],
ref_dsa_cache_seqlens[:size],
)
# Compare dsa_seqlens_expanded only for TARGET_VERIFY
if forward_mode.is_target_verify():
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"dsa_seqlens_expanded",
fused_dsa_seqlens_expanded[:size],
ref_dsa_seqlens_expanded[:size],
)
# Compare dsa_cu_seqlens_k only for the region that was updated
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"dsa_cu_seqlens_k",
fused_dsa_cu_seqlens_k[: 1 + size],
ref_dsa_cu_seqlens_k[: 1 + size],
)
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
check_tensor_equal(
"real_page_table",
fused_real_page_table[:rows, :cols],
ref_real_page_table[:rows, :cols],
)
if precomputed.flashmla_metadata is not None:
size = precomputed.seqlens_expanded_size
check_tensor_equal(
"flashmla_num_splits",
fused_flashmla_num_splits[: size + 1],
ref_flashmla_num_splits[: size + 1],
)
check_tensor_equal(
"flashmla_metadata",
fused_flashmla_metadata,
ref_flashmla_metadata,
)
def verify_multi_backend_fused_metadata_copy(
metadata0,
metadata1,
metadata2,
precomputed,
bs,
flashmla_num_splits_src=None,
flashmla_metadata_src=None,
):
"""
Verify that the multi-backend fused metadata copy kernel produces the same results
as individual copies for all three backends.
Args:
metadata0: The DSA metadata object for backend 0
metadata1: The DSA metadata object for backend 1
metadata2: The DSA metadata object for backend 2
precomputed: The precomputed metadata containing source tensors
bs: Batch size
flashmla_num_splits_src: Source FlashMLA num_splits tensor (optional)
flashmla_metadata_src: Source FlashMLA metadata tensor (optional)
Raises:
RuntimeError: If verification fails (tensors don't match)
"""
# Clone destination tensors to preserve fused kernel results
fused_results = []
for idx, metadata in enumerate([metadata0, metadata1, metadata2]):
fused_cache_seqlens = metadata.cache_seqlens_int32.clone()
fused_cu_seqlens_k = metadata.cu_seqlens_k.clone()
fused_page_table_1 = metadata.page_table_1.clone()
fused_dsa_cache_seqlens = metadata.dsa_cache_seqlens_int32.clone()
fused_dsa_cu_seqlens_k = metadata.dsa_cu_seqlens_k.clone()
fused_real_page_table = (
metadata.real_page_table.clone()
if precomputed.real_page_table is not None
else None
)
fused_flashmla_num_splits = None
fused_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
fused_flashmla_num_splits = metadata.flashmla_metadata.num_splits.clone()
fused_flashmla_metadata = (
metadata.flashmla_metadata.flashmla_metadata.clone()
)
fused_results.append(
{
"cache_seqlens": fused_cache_seqlens,
"cu_seqlens_k": fused_cu_seqlens_k,
"page_table_1": fused_page_table_1,
"dsa_cache_seqlens": fused_dsa_cache_seqlens,
"dsa_cu_seqlens_k": fused_dsa_cu_seqlens_k,
"real_page_table": fused_real_page_table,
"flashmla_num_splits": fused_flashmla_num_splits,
"flashmla_metadata": fused_flashmla_metadata,
}
)
# Run individual copy operations for each backend (reference implementation)
ref_results = []
for idx in range(3):
metadata = [metadata0, metadata1, metadata2][idx]
# Create reference tensors (zeroed out)
ref_cache_seqlens = torch.zeros_like(metadata.cache_seqlens_int32)
ref_cu_seqlens_k = torch.zeros_like(metadata.cu_seqlens_k)
ref_page_table_1 = torch.zeros_like(metadata.page_table_1)
ref_dsa_cache_seqlens = torch.zeros_like(metadata.dsa_cache_seqlens_int32)
ref_dsa_cu_seqlens_k = torch.zeros_like(metadata.dsa_cu_seqlens_k)
ref_real_page_table = (
torch.zeros_like(metadata.real_page_table)
if precomputed.real_page_table is not None
else None
)
ref_flashmla_num_splits = None
ref_flashmla_metadata = None
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits = torch.zeros_like(
metadata.flashmla_metadata.num_splits
)
ref_flashmla_metadata = torch.zeros_like(
metadata.flashmla_metadata.flashmla_metadata
)
# Copy operations (decode mode)
ref_cache_seqlens.copy_(precomputed.cache_seqlens)
ref_cu_seqlens_k[1:].copy_(precomputed.cu_seqlens_k[1:])
ref_page_table_1[:, : precomputed.max_len].copy_(precomputed.page_indices)
ref_dsa_cache_seqlens.copy_(precomputed.dsa_cache_seqlens)
# Copy DSA cu_seqlens
size = precomputed.seqlens_expanded_size
ref_dsa_cu_seqlens_k[1 : 1 + size].copy_(
precomputed.dsa_cu_seqlens_k[1 : 1 + size]
)
# Copy real page table
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
ref_real_page_table[:rows, :cols].copy_(precomputed.real_page_table)
# Copy FlashMLA metadata
if precomputed.flashmla_metadata is not None:
ref_flashmla_num_splits[: size + 1].copy_(
flashmla_num_splits_src[: size + 1]
)
ref_flashmla_metadata.copy_(flashmla_metadata_src)
ref_results.append(
{
"cache_seqlens": ref_cache_seqlens,
"cu_seqlens_k": ref_cu_seqlens_k,
"page_table_1": ref_page_table_1,
"dsa_cache_seqlens": ref_dsa_cache_seqlens,
"dsa_cu_seqlens_k": ref_dsa_cu_seqlens_k,
"real_page_table": ref_real_page_table,
"flashmla_num_splits": ref_flashmla_num_splits,
"flashmla_metadata": ref_flashmla_metadata,
}
)
# Compare results for all 3 backends
def check_tensor_equal(backend_idx, name, fused, ref):
if not torch.equal(fused, ref):
max_diff = (fused.float() - ref.float()).abs().max().item()
mismatched_elements = (fused != ref).sum().item()
total_elements = fused.numel()
raise RuntimeError(
f"MULTI-BACKEND FUSED METADATA COPY VERIFICATION FAILED!\n"
f"Backend: {backend_idx}\n"
f"Tensor: {name}\n"
f"Max difference: {max_diff}\n"
f"Mismatched elements: {mismatched_elements}/{total_elements}\n"
f"Fused shape: {fused.shape}, Ref shape: {ref.shape}\n"
f"Batch size: {bs}\n"
f"The multi-backend fused kernel produces different results than individual copies.\n"
f"This indicates a bug in the fused metadata copy kernel."
)
# Verify all tensors for all 3 backends (multi-backend is DECODE mode only)
for idx in range(3):
fused = fused_results[idx]
ref = ref_results[idx]
check_tensor_equal(
idx,
"cache_seqlens",
fused["cache_seqlens"],
ref["cache_seqlens"],
)
check_tensor_equal(
idx,
"cu_seqlens_k",
fused["cu_seqlens_k"],
ref["cu_seqlens_k"],
)
# Multi-backend is DECODE mode only, so compare only [:, :max_len]
check_tensor_equal(
idx,
"page_table_1",
fused["page_table_1"][:, : precomputed.max_len],
ref["page_table_1"][:, : precomputed.max_len],
)
check_tensor_equal(
idx,
"dsa_cache_seqlens",
fused["dsa_cache_seqlens"],
ref["dsa_cache_seqlens"],
)
# DECODE mode uses bs for dsa_cu_seqlens_k size
check_tensor_equal(
idx,
"dsa_cu_seqlens_k",
fused["dsa_cu_seqlens_k"][: bs + 1],
ref["dsa_cu_seqlens_k"][: bs + 1],
)
if precomputed.real_page_table is not None:
rows, cols = precomputed.real_page_table.shape
check_tensor_equal(
idx,
"real_page_table",
fused["real_page_table"][:rows, :cols],
ref["real_page_table"][:rows, :cols],
)
if precomputed.flashmla_metadata is not None:
# DECODE mode uses bs + 1 for flashmla_num_splits
check_tensor_equal(
idx,
"flashmla_num_splits",
fused["flashmla_num_splits"][: bs + 1],
ref["flashmla_num_splits"][: bs + 1],
)
check_tensor_equal(
idx,
"flashmla_metadata",
fused["flashmla_metadata"],
ref["flashmla_metadata"],
)
python/sglang/srt/layers/attention/dsa_backend.py
+134 -160
View File
@@ -138,13 +138,6 @@ def _to_2d_context_lens(seqlens_32: torch.Tensor, batch_size: int) -> torch.Tens
# Reuse this workspace buffer across all DSA backend instances
global_workspace_buffer = None
# Control whether to use fused metadata copy kernel for cuda graph replay (default: enabled)
# Set SGLANG_USE_FUSED_METADATA_COPY=0 or false to disable
_USE_FUSED_METADATA_COPY = envs.SGLANG_USE_FUSED_METADATA_COPY.get() and not _is_hip
_USE_FUSED_METADATA_GENERATION = (
envs.SGLANG_DSA_USE_FUSED_METADATA_GENERATION.get() and not _is_hip
)
@dataclass(frozen=True)
class DSAFlashMLAMetadata:
@@ -1274,7 +1267,7 @@ class DeepseekSparseAttnBackend(
# Normal Decode
max_len = metadata.page_table_1.shape[1]
if _USE_FUSED_METADATA_GENERATION and is_cuda():
if is_cuda() and not _is_hip:
from sglang.srt.layers.attention.triton_ops.dsa_metadata import (
fused_dsa_decode_metadata,
)
@@ -1316,7 +1309,7 @@ class DeepseekSparseAttnBackend(
elif forward_mode.is_target_verify():
max_seqlen_k = metadata.page_table_1.shape[1]
if _USE_FUSED_METADATA_GENERATION and is_cuda():
if is_cuda() and not _is_hip:
from sglang.srt.layers.attention.triton_ops.dsa_metadata import (
fused_dsa_target_verify_metadata,
)
@@ -1414,7 +1407,7 @@ class DeepseekSparseAttnBackend(
device=self.device,
)
if _USE_FUSED_METADATA_GENERATION and is_cuda():
if is_cuda() and not _is_hip:
from sglang.srt.layers.attention.triton_ops.dsa_metadata import (
fused_dsa_draft_extend_metadata,
)
@@ -1558,7 +1551,7 @@ class DeepseekSparseAttnBackend(
fused_kernel_succeeded = False
# Use fused CUDA kernel for all copy operations
if _USE_FUSED_METADATA_COPY:
if not _is_hip:
try:
from sglang.jit_kernel.fused_metadata_copy import (
fused_metadata_copy_cuda,
@@ -1631,7 +1624,8 @@ class DeepseekSparseAttnBackend(
f"Warning: Fused metadata copy kernel failed with error: {e}, falling back to individual copies."
)
# Fallback to individual copy operations if fused kernel disabled or failed
# Fallback to individual copy operations if the fused kernel is unavailable
# or fails at runtime.
if not fused_kernel_succeeded:
# Copy basic seqlens
metadata.cache_seqlens_int32.copy_(precomputed.cache_seqlens)
@@ -2822,153 +2816,134 @@ class DeepseekSparseAttnMultiStepBackend:
return
bs = forward_batch.batch_size
if envs.SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA.get():
# Precompute metadata once (shared across all backends)
precomputed = self.attn_backends[0]._precompute_replay_metadata(
bs=bs,
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
seq_lens_cpu=forward_batch.seq_lens_cpu,
forward_mode=ForwardMode.DECODE,
)
# Precompute metadata once (shared across all backends)
precomputed = self.attn_backends[0]._precompute_replay_metadata(
bs=bs,
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
seq_lens_cpu=forward_batch.seq_lens_cpu,
forward_mode=ForwardMode.DECODE,
)
# Use multi-backend fused copy when we have 3 or more backends
# This is 3x faster than calling the single-backend copy 3 times
if self.speculative_num_steps > 3:
try:
from sglang.jit_kernel.fused_metadata_copy import (
fused_metadata_copy_multi_cuda,
# Use multi-backend fused copy when we have 3 or more backends
# This is 3x faster than calling the single-backend copy 3 times
if self.speculative_num_steps > 3:
try:
from sglang.jit_kernel.fused_metadata_copy import (
fused_metadata_copy_multi_cuda,
)
metadata0 = self.attn_backends[0].decode_cuda_graph_metadata[bs]
metadata1 = self.attn_backends[1].decode_cuda_graph_metadata[bs]
metadata2 = self.attn_backends[2].decode_cuda_graph_metadata[bs]
# Set dsa_prefill_impl for first 3 backends (required by the method)
for i in range(3):
self.attn_backends[i].set_dsa_prefill_impl(forward_batch=None)
# Prepare FlashMLA tensors if needed
flashmla_num_splits_src = None
flashmla_metadata_src = None
flashmla_num_splits_dst0 = None
flashmla_num_splits_dst1 = None
flashmla_num_splits_dst2 = None
flashmla_metadata_dst0 = None
flashmla_metadata_dst1 = None
flashmla_metadata_dst2 = None
if precomputed.flashmla_metadata is not None:
flashmla_num_splits_src = precomputed.flashmla_metadata.num_splits
flashmla_metadata_src = (
precomputed.flashmla_metadata.flashmla_metadata
)
flashmla_num_splits_dst0 = metadata0.flashmla_metadata.num_splits
flashmla_num_splits_dst1 = metadata1.flashmla_metadata.num_splits
flashmla_num_splits_dst2 = metadata2.flashmla_metadata.num_splits
flashmla_metadata_dst0 = (
metadata0.flashmla_metadata.flashmla_metadata
)
flashmla_metadata_dst1 = (
metadata1.flashmla_metadata.flashmla_metadata
)
flashmla_metadata_dst2 = (
metadata2.flashmla_metadata.flashmla_metadata
)
metadata0 = self.attn_backends[0].decode_cuda_graph_metadata[bs]
metadata1 = self.attn_backends[1].decode_cuda_graph_metadata[bs]
metadata2 = self.attn_backends[2].decode_cuda_graph_metadata[bs]
# Call the multi-backend fused kernel for first 3 backends
fused_metadata_copy_multi_cuda(
# Source tensors
precomputed.cache_seqlens,
precomputed.cu_seqlens_k,
precomputed.page_indices,
precomputed.dsa_cache_seqlens,
precomputed.dsa_cu_seqlens_k,
precomputed.real_page_table,
flashmla_num_splits_src,
flashmla_metadata_src,
# Destination tensors for backend 0
metadata0.cache_seqlens_int32,
metadata0.cu_seqlens_k,
metadata0.page_table_1,
metadata0.dsa_cache_seqlens_int32,
metadata0.dsa_cu_seqlens_k,
(
metadata0.real_page_table
if precomputed.real_page_table is not None
else None
),
flashmla_num_splits_dst0,
flashmla_metadata_dst0,
# Destination tensors for backend 1
metadata1.cache_seqlens_int32,
metadata1.cu_seqlens_k,
metadata1.page_table_1,
metadata1.dsa_cache_seqlens_int32,
metadata1.dsa_cu_seqlens_k,
(
metadata1.real_page_table
if precomputed.real_page_table is not None
else None
),
flashmla_num_splits_dst1,
flashmla_metadata_dst1,
# Destination tensors for backend 2
metadata2.cache_seqlens_int32,
metadata2.cu_seqlens_k,
metadata2.page_table_1,
metadata2.dsa_cache_seqlens_int32,
metadata2.dsa_cu_seqlens_k,
(
metadata2.real_page_table
if precomputed.real_page_table is not None
else None
),
flashmla_num_splits_dst2,
flashmla_metadata_dst2,
# Parameters
bs,
precomputed.max_len,
precomputed.seqlens_expanded_size,
)
# Set dsa_prefill_impl for first 3 backends (required by the method)
for i in range(3):
self.attn_backends[i].set_dsa_prefill_impl(forward_batch=None)
# Prepare FlashMLA tensors if needed
flashmla_num_splits_src = None
flashmla_metadata_src = None
flashmla_num_splits_dst0 = None
flashmla_num_splits_dst1 = None
flashmla_num_splits_dst2 = None
flashmla_metadata_dst0 = None
flashmla_metadata_dst1 = None
flashmla_metadata_dst2 = None
if precomputed.flashmla_metadata is not None:
flashmla_num_splits_src = (
precomputed.flashmla_metadata.num_splits
)
flashmla_metadata_src = (
precomputed.flashmla_metadata.flashmla_metadata
)
flashmla_num_splits_dst0 = (
metadata0.flashmla_metadata.num_splits
)
flashmla_num_splits_dst1 = (
metadata1.flashmla_metadata.num_splits
)
flashmla_num_splits_dst2 = (
metadata2.flashmla_metadata.num_splits
)
flashmla_metadata_dst0 = (
metadata0.flashmla_metadata.flashmla_metadata
)
flashmla_metadata_dst1 = (
metadata1.flashmla_metadata.flashmla_metadata
)
flashmla_metadata_dst2 = (
metadata2.flashmla_metadata.flashmla_metadata
)
# Call the multi-backend fused kernel for first 3 backends
fused_metadata_copy_multi_cuda(
# Source tensors
precomputed.cache_seqlens,
precomputed.cu_seqlens_k,
precomputed.page_indices,
precomputed.dsa_cache_seqlens,
precomputed.dsa_cu_seqlens_k,
precomputed.real_page_table,
flashmla_num_splits_src,
flashmla_metadata_src,
# Destination tensors for backend 0
metadata0.cache_seqlens_int32,
metadata0.cu_seqlens_k,
metadata0.page_table_1,
metadata0.dsa_cache_seqlens_int32,
metadata0.dsa_cu_seqlens_k,
(
metadata0.real_page_table
if precomputed.real_page_table is not None
else None
),
flashmla_num_splits_dst0,
flashmla_metadata_dst0,
# Destination tensors for backend 1
metadata1.cache_seqlens_int32,
metadata1.cu_seqlens_k,
metadata1.page_table_1,
metadata1.dsa_cache_seqlens_int32,
metadata1.dsa_cu_seqlens_k,
(
metadata1.real_page_table
if precomputed.real_page_table is not None
else None
),
flashmla_num_splits_dst1,
flashmla_metadata_dst1,
# Destination tensors for backend 2
metadata2.cache_seqlens_int32,
metadata2.cu_seqlens_k,
metadata2.page_table_1,
metadata2.dsa_cache_seqlens_int32,
metadata2.dsa_cu_seqlens_k,
(
metadata2.real_page_table
if precomputed.real_page_table is not None
else None
),
flashmla_num_splits_dst2,
flashmla_metadata_dst2,
# Parameters
bs,
precomputed.max_len,
precomputed.seqlens_expanded_size,
# Copy remaining backends one by one (if > 3 backends)
for i in range(3, self.speculative_num_steps - 1):
self.attn_backends[
i
].init_forward_metadata_replay_cuda_graph_from_precomputed(
bs=bs,
precomputed=precomputed,
forward_mode=ForwardMode.DECODE,
)
except (ImportError, Exception) as e:
# Fallback to loop if multi-backend kernel not available or fails
if isinstance(e, ImportError):
print(
"Warning: Multi-backend fused metadata copy kernel not available, falling back to loop."
)
else:
print(
f"Warning: Multi-backend fused metadata copy kernel failed with error: {e}, falling back to loop."
)
# Copy remaining backends one by one (if > 3 backends)
for i in range(3, self.speculative_num_steps - 1):
self.attn_backends[
i
].init_forward_metadata_replay_cuda_graph_from_precomputed(
bs=bs,
precomputed=precomputed,
forward_mode=ForwardMode.DECODE,
)
except (ImportError, Exception) as e:
# Fallback to loop if multi-backend kernel not available or fails
if isinstance(e, ImportError):
print(
"Warning: Multi-backend fused metadata copy kernel not available, falling back to loop."
)
else:
print(
f"Warning: Multi-backend fused metadata copy kernel failed with error: {e}, falling back to loop."
)
for i in range(self.speculative_num_steps - 1):
self.attn_backends[
i
].init_forward_metadata_replay_cuda_graph_from_precomputed(
bs=bs,
precomputed=precomputed,
forward_mode=ForwardMode.DECODE,
)
else:
# Less than 3 backends: copy to each backend individually
for i in range(self.speculative_num_steps - 1):
self.attn_backends[
i
@@ -2978,15 +2953,14 @@ class DeepseekSparseAttnMultiStepBackend:
forward_mode=ForwardMode.DECODE,
)
else:
# Less than 3 backends: copy to each backend individually
for i in range(self.speculative_num_steps - 1):
self.attn_backends[i]._apply_cuda_graph_metadata(
self.attn_backends[
i
].init_forward_metadata_replay_cuda_graph_from_precomputed(
bs=bs,
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
seq_lens_cpu=forward_batch.seq_lens_cpu,
precomputed=precomputed,
forward_mode=ForwardMode.DECODE,
spec_info=forward_batch.spec_info,
out_cache_loc=None,
)
def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch) -> None:
python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py
-10
View File
@@ -1,10 +0,0 @@
# [Deprecated] Re-export shim for backward compatibility. Use dsa.dsa_mtp_verification instead.
import warnings
warnings.warn(
"sglang.srt.layers.attention.nsa.nsa_mtp_verification is deprecated; "
"use sglang.srt.layers.attention.dsa.dsa_mtp_verification instead.",
DeprecationWarning,
stacklevel=2,
)
from sglang.srt.layers.attention.dsa.dsa_mtp_verification import * # noqa: F401, F403