diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx index 44bb8fb87..50802c1fb 100644 --- a/docs_new/docs/references/environment_variables.mdx +++ b/docs_new/docs/references/environment_variables.mdx @@ -521,16 +521,6 @@ SGLang supports various environment variables that can be used to configure its Tie-break mode for FlashInfer DSA topk when --dsa-topk-backend=flashinfer: unset disables explicit tie-breaking, small prefers the smaller candidate index for equal scores, and large prefers the larger candidate index for equal scores. Setting this variable makes FlashInfer use deterministic topk. unset - - SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA - Precompute metadata that can be shared among different draft steps when MTP is enabled. SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA is a deprecated alias. - true - - - SGLANG_USE_FUSED_METADATA_COPY - Control whether to use fused metadata copy kernel for cuda graph replay - true - SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD When the maximum kv len in current prefill batch exceeds this value, the sparse mla kernel will be applied, else it falls back to dense MHA implementation. Default to the index topk of model (2048 for DeepSeek V3.2). SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD is a deprecated alias. diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index c97cb1199..3e91b6499 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py index d75fe9611..4b2be7dee 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py +++ b/python/sglang/srt/layers/attention/dsa/dsa_backend_mtp_precompute.py @@ -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, ) diff --git a/python/sglang/srt/layers/attention/dsa/dsa_mtp_verification.py b/python/sglang/srt/layers/attention/dsa/dsa_mtp_verification.py deleted file mode 100644 index bf6c73b5e..000000000 --- a/python/sglang/srt/layers/attention/dsa/dsa_mtp_verification.py +++ /dev/null @@ -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"], - ) diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 1b8f70e7e..e5b6b2337 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -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: diff --git a/python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py b/python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py deleted file mode 100644 index ee111ce43..000000000 --- a/python/sglang/srt/layers/attention/nsa/nsa_mtp_verification.py +++ /dev/null @@ -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 diff --git a/test/manual/test_dsa_alias_cli_registry_env.py b/test/manual/test_dsa_alias_cli_registry_env.py index c9221cd2b..aa02eeb37 100644 --- a/test/manual/test_dsa_alias_cli_registry_env.py +++ b/test/manual/test_dsa_alias_cli_registry_env.py @@ -202,8 +202,6 @@ class TestEnvVarAliases(unittest.TestCase): "SGLANG_NSA_FUSE_TOPK", "SGLANG_DSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD", "SGLANG_NSA_PREFILL_DENSE_ATTN_KV_LEN_THRESHOLD", - "SGLANG_DSA_ENABLE_MTP_PRECOMPUTE_METADATA", - "SGLANG_NSA_ENABLE_MTP_PRECOMPUTE_METADATA", ]: os.environ.pop(key, None) # Re-import to reset descriptor state