diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index d229a4bc0..ecbc8a3c2 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -89,6 +89,9 @@ class ForwardMetadata: seq_lens: Optional[torch.Tensor] = None actual_seq_lengths_q: Optional[torch.Tensor] = None actual_seq_lengths_q_pa: Optional[torch.Tensor] = None + # CPU mirror of actual_seq_lengths_q_pa for the host metadata op + # (torch.ops.npu.sparse_attn_sharedkv_metadata_host reads CPU int32 inputs). + actual_seq_lengths_q_pa_cpu: Optional[torch.Tensor] = None actual_seq_lengths_kv: Optional[torch.Tensor] = None # swa attention mask for graph mode decode diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py index bc9c816f7..a662d6cba 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_dsv4_backend.py @@ -398,7 +398,7 @@ class CompressorAscendBackendMixin: allow_build=False, ) - cmp_kv = torch.ops.custom.compressor( + cmp_kv = torch.ops.npu.compressor( x, compressor._fused_wkv_w, compressor._fused_wgate_w, @@ -989,6 +989,14 @@ class DeepseekV4AscendAttnBackend( dtype=torch.int32, device=device, ) + # q_pa is constant per graph shape (never rewritten at replay), so the + # CPU mirror for the host metadata op is built once here. + metadata.actual_seq_lengths_q_pa_cpu = torch.arange( + 0, + bs * tokens_per_req + tokens_per_req, + tokens_per_req, + dtype=torch.int32, + ) # init >=1 so the captured kernel records valid attention work; replay overwrites in-place metadata.actual_seq_lengths_kv = torch.ones( @@ -1189,6 +1197,10 @@ class DeepseekV4AscendAttnBackend( fm.seq_lens_cpu_int, ctx.live_seq_lens_cpu.int(), ) + else: + # CPU mirror of the kv buffer written below, from its CPU source — the + # host metadata op reads this instead of a D2H sync. + fm.seq_lens_cpu_int = ctx.final_seq_lens_cpu[: ctx.bs].int().clamp(min=1) fm.actual_seq_lengths_kv.copy_(attn_seq_lens.clamp(min=1)) def _refresh_graph_compress_page_tables_direct(self, ctx) -> None: @@ -1432,6 +1444,13 @@ class DeepseekV4AscendAttnBackend( [torch.zeros(1, dtype=torch.int32, device=device), actual_q], dim=0, ) + fm.actual_seq_lengths_q_pa_cpu = torch.cat( + [ + torch.zeros(1, dtype=torch.int32), + torch.cumsum(seq_lens_cpu, dim=0).int(), + ], + dim=0, + ) elif forward_batch.forward_mode.is_decode(): B = forward_batch.batch_size fm.actual_seq_lengths_q = torch.arange( @@ -1440,6 +1459,7 @@ class DeepseekV4AscendAttnBackend( fm.actual_seq_lengths_q_pa = torch.arange( 0, B + 1, dtype=torch.int32, device=device ) + fm.actual_seq_lengths_q_pa_cpu = torch.arange(0, B + 1, dtype=torch.int32) elif ( forward_batch.forward_mode.is_target_verify() or forward_batch.forward_mode.is_draft_extend_v2() @@ -1458,9 +1478,17 @@ class DeepseekV4AscendAttnBackend( [torch.zeros(1, dtype=torch.int32, device=device), actual_q], dim=0, ) + fm.actual_seq_lengths_q_pa_cpu = torch.cat( + [ + torch.zeros(1, dtype=torch.int32), + torch.arange(n_draft, B * n_draft + 1, n_draft, dtype=torch.int32), + ], + dim=0, + ) else: fm.actual_seq_lengths_q = None fm.actual_seq_lengths_q_pa = None + fm.actual_seq_lengths_q_pa_cpu = None fm.swa_page_table = ( fm.block_tables_swa if fm.block_tables_swa is not None else fm.block_tables @@ -1514,8 +1542,6 @@ class DeepseekV4AscendAttnBackend( ) -> dict: fm = self.forward_metadata common = { - "cu_seqlens_q": actual_seq_lengths_q_pa, - "seqused_kv": actual_seq_lengths_kv, "cmp_ratio": 1, "ori_mask_mode": 4, "cmp_mask_mode": 3, @@ -1534,28 +1560,27 @@ class DeepseekV4AscendAttnBackend( "has_ori_kv": True, "has_cmp_kv": False, } + # The host metadata op reads CPU int32 mirrors — never a D2H sync of the + # device tensors (that would drain the stream and stall overlapped prep). c1a_kwargs = base_kwargs | common if self._is_dspark_draft_worker: - seq_lens_cpu = getattr(fm, "seq_lens_cpu_int", None) - max_seqlen_kv = ( - int(seq_lens_cpu[:bs].max().item()) - if seq_lens_cpu is not None and bs > 0 - else int(actual_seq_lengths_kv[:bs].max().item()) - ) - c1a_kwargs.update( - cu_seqlens_ori_kv=actual_seq_lengths_q_pa, - max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_kv, - ) - c1a_metadata = torch.ops._C_ascend.npu_sparse_attn_sharedkv_metadata( - device=str(actual_seq_lengths_kv.device), - **c1a_kwargs, - ) + cu_q_cpu = fm.actual_seq_lengths_q_pa_cpu + if cu_q_cpu is not None and cu_q_cpu.numel() > bs + 1: + cu_q_cpu = cu_q_cpu[: bs + 1] + host_inputs = {"seqused_kv": fm.seq_lens_cpu_int[:bs].int()} + if cu_q_cpu is not None: + host_inputs["cu_seqlens_q"] = cu_q_cpu + c1a_kwargs = c1a_kwargs | host_inputs + metadata_op = torch.ops.npu.sparse_attn_sharedkv_metadata_host else: - c1a_metadata = torch.ops.custom.npu_sparse_attn_sharedkv_metadata( - **c1a_kwargs, - ) - kernel_metadata = {"c1a_metadata": c1a_metadata} + # The device-side op requires tensor args for backend dispatch; pass + # the device mirrors just like the pre-refactor call did. + c1a_kwargs = c1a_kwargs | { + "cu_seqlens_q": actual_seq_lengths_q_pa, + "seqused_kv": actual_seq_lengths_kv, + } + metadata_op = torch.ops.custom.npu_sparse_attn_sharedkv_metadata + kernel_metadata = {"c1a_metadata": metadata_op(**c1a_kwargs)} if self._dsv4_has_c4: c4a_overrides = { @@ -1665,10 +1690,7 @@ class DeepseekV4AscendAttnBackend( if ori_sparse_indices is not None: attn_kwargs["ori_sparse_indices"] = ori_sparse_indices q_arg = attn_kwargs.pop("q") - if self._is_dspark_draft_worker: - out, _ = torch.ops._C_ascend.npu_sparse_attn_sharedkv(q_arg, **attn_kwargs) - else: - out, _ = torch.ops.custom.npu_sparse_attn_sharedkv(q_arg, **attn_kwargs) + out, _ = torch.ops.npu.sparse_attn_sharedkv(q_arg, **attn_kwargs) return out def _forward_compressed( @@ -1739,7 +1761,7 @@ class DeepseekV4AscendAttnBackend( else: attn_kwargs["cmp_sparse_indices"] = None q_arg = attn_kwargs.pop("q") - out, _ = torch.ops.custom.npu_sparse_attn_sharedkv(q_arg, **attn_kwargs) + out, _ = torch.ops.npu.sparse_attn_sharedkv(q_arg, **attn_kwargs) return out def get_swa_out_cache_loc(self, forward_batch: ForwardBatch) -> torch.Tensor: diff --git a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py index 228e3cd9a..eddda0412 100644 --- a/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py +++ b/python/sglang/srt/hardware_backend/npu/dsv4/dsv4_memory_pool.py @@ -97,7 +97,7 @@ class NPUCompressStatePool(CompressStatePool): 128, ), f"NPUCompressStatePool only supports ratio in (4, 128); got {ratio}" assert dtype == torch.float32, ( - "Atlas A3 custom.compressor requires FP32 state_cache, " + "Atlas A3 npu.compressor requires FP32 state_cache, " f"but NPUCompressStatePool got {dtype}." ) assert ring_size > 0, f"ring_size must be positive, got {ring_size}" @@ -404,7 +404,7 @@ class DSV4NPUTokenToKVPool(DeepSeekV4TokenToKVPool): def get_state_cache(self, layer_id: int, from_indexer: bool) -> torch.Tensor: """FP32 ``[block_num, ring_size, 2*coff*D]`` view of this layer's kv+score buffer — the fused compressor op - (``torch.ops.custom.compressor``)'s ``state_cache`` argument.""" + (``torch.ops.npu.compressor``)'s ``state_cache`` argument.""" return self._get_state_pool(layer_id, from_indexer).state_cache_3d # ------------------------------------------------------------------ diff --git a/python/sglang/srt/hardware_backend/npu/extra_ops_loader.py b/python/sglang/srt/hardware_backend/npu/extra_ops_loader.py index b7ac5e528..8638ceadf 100644 --- a/python/sglang/srt/hardware_backend/npu/extra_ops_loader.py +++ b/python/sglang/srt/hardware_backend/npu/extra_ops_loader.py @@ -25,7 +25,36 @@ class OpLibSpec: class TorchOpLoader: - """Load a standalone .so into ``torch.ops`` and validate its operators.""" + """ + Loader for PyTorch custom operators from shared libraries. + + This class handles the registration and initialization of custom PyTorch + operators (Ops) from dynamically linked shared object (.so) files. It supports + environment-based library path discovery, dependency pre-loading, and + operator existence validation. + + Usage: + 1. Create an OpLibSpec with operator metadata + 2. Instantiate TorchOpLoader with the spec + 3. Call initialize() to load and register the operators + + Example: + >>> spec = OpLibSpec( + ... name="My custom ops", + ... so_env="MY_LIB_SO_PATH", + ... namespace="_C_my_lib", + ... required_ops=("op1", "op2"), + ... pre_load_imports=("torch", "other_dep"), + ... ) + >>> loader = TorchOpLoader(spec) + >>> lib_path = loader.initialize() + >>> # Ops are now registered under namespace: _C_my_lib.op1() + + The loader will raise appropriate exceptions if: + - The shared library cannot be found (via SO_PATH env var or default paths) + - Pre-load imports fail + - Required operators are missing after loading + """ def __init__(self, spec: OpLibSpec) -> None: self._spec = spec @@ -105,18 +134,3 @@ class TorchOpLoader: self._loaded_library = library_path logger.info("Registered %s operators from %s", self._spec.name, library_path) return library_path - - -def initialize_dspark_sparse_attn_ops() -> Optional[Path]: - """Register the DSpark sparse-attention ops before backend execution.""" - spec = OpLibSpec( - name="DSpark sparse-attention", - so_env="SGLANG_DSPARK_EXTRA_OPS_SO", - namespace="_C_ascend", - required_ops=( - "npu_sparse_attn_sharedkv_metadata", - "npu_sparse_attn_sharedkv", - ), - pre_load_imports=("torch_npu",), - ) - return TorchOpLoader(spec).initialize() diff --git a/python/sglang/srt/mem_cache/allocation.py b/python/sglang/srt/mem_cache/allocation.py index 2b8105e3a..fe20022aa 100644 --- a/python/sglang/srt/mem_cache/allocation.py +++ b/python/sglang/srt/mem_cache/allocation.py @@ -114,16 +114,15 @@ def get_last_loc( "torch_native", ) and decode_backend not in ("ascend", "torch_native") - if _is_hip and uses_triton_dispatch: - # HIP-only: the legacy get_last_loc_triton kernel emits a + if (_is_hip or _is_npu) and uses_triton_dispatch: + # HIP and NPU DSV4: the legacy get_last_loc_triton kernel emits a # mixed-width int32->int64 store that Triton mis-compiles on HIP, # producing out-of-range last_loc values under EAGLE + # page_size>1 (e.g. with aiter unified attention or the triton - # attention backend). The bug is in the Triton HIP codegen, not - # in any particular attention backend, so route every HIP path - # that would otherwise use get_last_loc_triton through the - # int32-safe variant. Non-HIP hardware keeps the original - # dispatcher below. + # attention backend), and can fault in the equivalent NPU DSV4 + # allocation path. Route those paths through the variant whose + # in-kernel result remains int32 and is promoted only after launch. + # Other hardware/backends keep the original dispatcher below. return get_last_loc_triton_safe( req_to_token, req_pool_indices_tensor, prefix_lens_tensor ) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py index 8e4d69a52..adcd6c42f 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -349,12 +349,6 @@ class DSparkWorkerV2(BaseSpecWorker): def init_attention_backends(self): with self._draft_context(): - if _is_npu: - from sglang.srt.hardware_backend.npu.extra_ops_loader import ( - initialize_dspark_sparse_attn_ops, - ) - - initialize_dspark_sparse_attn_ops() self._draft_worker.init_attention_backends() self._need_mamba_verify_commit = mambaish_config( self.model_runner.model_config diff --git a/test/registered/unit/npu/attention/test_npu_ascend_backend.py b/test/registered/unit/npu/attention/test_npu_ascend_backend.py index ddc6e9093..c110ea919 100644 --- a/test/registered/unit/npu/attention/test_npu_ascend_backend.py +++ b/test/registered/unit/npu/attention/test_npu_ascend_backend.py @@ -153,6 +153,7 @@ class TestForwardMetadata(unittest.TestCase): "seq_lens", "actual_seq_lengths_q", "actual_seq_lengths_q_pa", + "actual_seq_lengths_q_pa_cpu", "actual_seq_lengths_kv", "swa_mask", "prefix_lens",