diff --git a/python/sglang/kernels/ops/elementwise/add3.py b/python/sglang/kernels/ops/elementwise/add3.py index 97590b8ce..911a28dd3 100644 --- a/python/sglang/kernels/ops/elementwise/add3.py +++ b/python/sglang/kernels/ops/elementwise/add3.py @@ -12,6 +12,7 @@ from sglang.kernels.jit.utils import ( load_jit, make_cpp_args, ) +from sglang.srt.utils import is_npu if TYPE_CHECKING: from tvm_ffi.module import Module @@ -38,7 +39,8 @@ def covered(a: torch.Tensor, b: torch.Tensor, c: torch.Tensor) -> bool: """Same-shape contiguous CUDA bf16 tensors, numel a multiple of the widest vector (16 elements).""" return ( - a.dtype == b.dtype == c.dtype == torch.bfloat16 + not is_npu() + and a.dtype == b.dtype == c.dtype == torch.bfloat16 and a.shape == b.shape == c.shape and a.is_contiguous() and b.is_contiguous() diff --git a/python/sglang/kernels/ops/kimi_k3/__init__.py b/python/sglang/kernels/ops/kimi_k3/__init__.py index 7e8ce029a..c95fe21de 100644 --- a/python/sglang/kernels/ops/kimi_k3/__init__.py +++ b/python/sglang/kernels/ops/kimi_k3/__init__.py @@ -2,9 +2,13 @@ from __future__ import annotations from typing import TYPE_CHECKING, Optional +from sglang.srt.utils import is_npu + if TYPE_CHECKING: import torch +_is_npu = is_npu() + _K3_N_GEMM_DISPATCH_MAP = { (144, 7168): 16, (896, 7168): 8, @@ -65,12 +69,13 @@ def kimi_k3_tiny_gemm( m, k = x.shape n, _ = w.shape - if max_num_tokens := _K3_N_GEMM_DISPATCH_MAP.get((n, k)): - if 0 < m <= max_num_tokens: - return tiny_n_gemm_bf16(x, w) - if max_num_tokens := _K3_K_GEMM_DISPATCH_MAP.get((n, k)): - if 0 < m <= max_num_tokens: - return tiny_k_gemm_bf16(x, w) + if not _is_npu: + if max_num_tokens := _K3_N_GEMM_DISPATCH_MAP.get((n, k)): + if 0 < m <= max_num_tokens: + return tiny_n_gemm_bf16(x, w) + if max_num_tokens := _K3_K_GEMM_DISPATCH_MAP.get((n, k)): + if 0 < m <= max_num_tokens: + return tiny_k_gemm_bf16(x, w) return torch.nn.functional.linear(x, w) diff --git a/python/sglang/kernels/ops/kimi_k3/mla_output_gate.py b/python/sglang/kernels/ops/kimi_k3/mla_output_gate.py index 60a64f899..721f82b1b 100644 --- a/python/sglang/kernels/ops/kimi_k3/mla_output_gate.py +++ b/python/sglang/kernels/ops/kimi_k3/mla_output_gate.py @@ -12,11 +12,13 @@ from sglang.kernels.jit.utils import ( load_jit, make_cpp_args, ) +from sglang.srt.utils import is_npu if TYPE_CHECKING: from tvm_ffi.module import Module _THREADS: int = 256 +_is_npu = is_npu() @cache_once @@ -33,7 +35,8 @@ def _jit_mla_output_gate_module() -> Module: def covered(x: torch.Tensor, gate: torch.Tensor) -> bool: return ( - x.dtype == torch.bfloat16 + not _is_npu + and x.dtype == torch.bfloat16 and gate.dtype == torch.bfloat16 and x.shape == gate.shape and x.is_contiguous() diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py b/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py index af6ffc86d..38a9e94b0 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py @@ -9,6 +9,7 @@ from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_f from sglang.kernels.ops.speculative.dspark.dispatch import inputs_on_cuda from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout +from sglang.srt.utils import is_npu class RaggedVerifyWindow(msgspec.Struct, frozen=True): @@ -797,7 +798,7 @@ def build_commit_inject_layout_triton( class BuildOutTokens: @classmethod def execute(cls, *args, **kwargs) -> torch.Tensor: - if inputs_on_cuda(*args, **kwargs): + if not is_npu() and inputs_on_cuda(*args, **kwargs): return cls.triton(*args, **kwargs) return cls.torch(*args, **kwargs) diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py b/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py index 2c40bbda3..aebf21f91 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py +++ b/python/sglang/multimodal_gen/runtime/layers/attention/backends/ascend_fa.py @@ -52,6 +52,12 @@ class AscendFABackend(AttentionBackend): def get_builder_cls() -> type["AttentionMetadataBuilder"]: return AscendFAMetadataBuilder + @classmethod + def supports_ring_rotation(cls) -> bool: + """Whether this backend can serve as the ring-attention kernel; the + per-hop online-softmax merge needs the kernel's softmax LSE.""" + return True + class AscendFAImpl(AttentionImpl): diff --git a/python/sglang/srt/arg_groups/speculative_hook.py b/python/sglang/srt/arg_groups/speculative_hook.py index 1b1f65e3c..cc6e30bd1 100644 --- a/python/sglang/srt/arg_groups/speculative_hook.py +++ b/python/sglang/srt/arg_groups/speculative_hook.py @@ -276,14 +276,17 @@ def _target_checkpoint_bundles_dspark_draft(server_args: ServerArgs) -> bool: def _handle_dspark(server_args: ServerArgs) -> None: - if not server_args.device.startswith("cuda"): - raise ValueError("DSpark speculative decoding only supports CUDA device.") + _is_npu = server_args.device.startswith("npu") + if not server_args.device.startswith("cuda") and not _is_npu: + raise ValueError( + "DSpark speculative decoding only supports CUDA and NPU devices." + ) # dp_size==1 with dp_attention is a degenerate flag under DSV4 CP; skip DP-only checks. if server_args.enable_dp_attention and server_args.dp_size > 1: if not server_args.enable_dp_lm_head: raise ValueError("DSpark with dp attention requires --enable-dp-lm-head.") - if server_args.moe_a2a_backend != "none": + if not _is_npu and server_args.moe_a2a_backend != "none": raise ValueError( "DSpark with dp attention only supports the built-in TP MoE " f"(moe_a2a_backend='none'), got {server_args.moe_a2a_backend!r}." @@ -294,7 +297,8 @@ def _handle_dspark(server_args: ServerArgs) -> None: f"(attn_cp_size={server_args.attn_cp_size})." ) if ( - server_args.speculative_moe_a2a_backend is not None + not _is_npu + and server_args.speculative_moe_a2a_backend is not None and server_args.speculative_moe_a2a_backend != server_args.moe_a2a_backend ): raise ValueError( diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index cc0ab28b7..5390e8583 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -348,6 +348,9 @@ class Envs: SGLANG_DSPARK_FP32_LM_HEAD = EnvBool(False) SGLANG_DSPARK_FAST_SAMPLING = EnvBool(True) SGLANG_DSPARK_FOLDED_SAMPLING = EnvInt(DsparkFoldedSampling.AUTO) + SGLANG_DSPARK_FOLDED_PROPOSAL = EnvBool(True) + SGLANG_DSPARK_STACKED_CTX_KV = EnvBool(True) + SGLANG_DSPARK_EMBED_IN_GRAPH = EnvBool(True) SGLANG_DSPARK_OPT_MARKOV_W2_BF16 = EnvBool(True) SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD = EnvBool(True) SGLANG_DSPARK_ENABLE_MULTI_STREAM = EnvBool(True) @@ -670,6 +673,13 @@ class Envs: SGLANG_EXPERIMENTAL_LORA_OPTI = EnvBool(False) # Enable int4x2 weights loading SGLANG_NPU_W4A4_NEW_PACKING = EnvBool(False) + # Keep K3 shared experts and dense MLPs sharded over attention TP. + SGLANG_K3_SHARED_EXPERTS_ATTN_TP = EnvBool(False) + SGLANG_K3_DENSE_MLP_ATTN_TP = EnvBool(False) + # Use the graph-safe Triton-Ascend kernel for masked speculative KV commits. + SGLANG_NPU_USE_TRITON_PREFIX_KV_CACHE_STORE = EnvBoolWithAlias( + False, deprecated_name="SGLANG_NPU_USE_TRITON_KV_CACHE_STORE" + ) # Quantize x to int8 in the dispatch operator DEEP_NORMAL_MODE_USE_INT8_QUANT = EnvBool(False) # This argument is deprecated SGLANG_NPU_FUSED_MOE_MODE = EnvInt(1) 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 89e9de06d..5d6ae2a74 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -28,7 +28,11 @@ from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_flags, get_spec from sglang.srt.speculative.spec_info import SpecInput, SpecInputType -from sglang.srt.utils import get_bool_env_var, get_current_device_stream_fast +from sglang.srt.utils import ( + get_bool_env_var, + get_current_device_stream_fast, + next_power_of_2, +) if TYPE_CHECKING: from sglang.srt.layers.radix_attention import RadixAttention @@ -337,6 +341,15 @@ class AscendAttnBackend(AttentionBackend): self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False") self.enable_torch_compile = get_flags().capture.enable_torch_compile self.speculative_num_draft_tokens = get_spec().speculative_num_draft_tokens + if ( + self.speculative_num_draft_tokens is not None + and model_runner.is_draft_worker + ): + self.speculative_num_draft_tokens = ( + model_runner.spec_algorithm.get_num_tokens_per_req_for_target_verify( + int(self.speculative_num_draft_tokens), is_draft_worker=True + ) + ) self.ascend_attn_mask_builder = AscendAttnMaskBuilder( model_runner, self.device, self.use_fia, self.use_mla ) @@ -430,14 +443,23 @@ class AscendAttnBackend(AttentionBackend): def init_forward_metadata(self, forward_batch: ForwardBatch): """Init the metadata for a forward pass.""" self.forward_metadata = ForwardMetadata() - seq_lens_max = forward_batch.seq_lens.max() if forward_batch.forward_mode.is_target_verify(): - seq_lens_max += self.speculative_num_draft_tokens + # Overlap scheduling can publish the CPU sequence length one step + # ahead of the device tensor. FIA consumes seq_lens_cpu below, so + # derive the block-table width from the same source. Otherwise a + # page-aligned request can expose KV_S=N while asking FIA for N+1. + seq_lens_max = ( + forward_batch.seq_lens_cpu.max().item() + + self.speculative_num_draft_tokens + ) elif ( forward_batch.forward_mode.is_decode_or_idle() and forward_batch.spec_info is not None ): + seq_lens_max = forward_batch.seq_lens.max() seq_lens_max += self.speculative_step_id + 1 + else: + seq_lens_max = forward_batch.seq_lens.max() self.forward_metadata.block_tables = ( self.req_to_token_pool.req_to_token[ forward_batch.req_pool_indices, :seq_lens_max @@ -2781,13 +2803,32 @@ class AscendAttnBackend(AttentionBackend): layer.tp_q_head_num, self.qk_rope_head_dim, ) + if (layer.tp_q_head_num & (layer.tp_q_head_num - 1)) != 0: + power_of_2_head = next_power_of_2(layer.tp_q_head_num) + padding_head = power_of_2_head - layer.tp_q_head_num + q_padding_tensor = torch.zeros( + [num_tokens, q.shape[1], padding_head, q.shape[-1]], + dtype=q.dtype, + device=q.device, + ) + q = torch.cat((q, q_padding_tensor), dim=-2) + q_rope_padding_tensor = torch.zeros( + [num_tokens, q_rope.shape[1], padding_head, q_rope.shape[-1]], + dtype=q_rope.dtype, + device=q_rope.device, + ) + q_rope = torch.cat((q_rope, q_rope_padding_tensor), dim=-2) + tp_q_head_num = power_of_2_head + else: + tp_q_head_num = layer.tp_q_head_num + attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score( q, kv_c, kv_c, query_rope=q_rope, key_rope=k_pe, - num_heads=layer.tp_q_head_num, + num_heads=tp_q_head_num, num_key_value_heads=layer.tp_k_head_num, input_layout="BSND", atten_mask=None, @@ -2799,6 +2840,7 @@ class AscendAttnBackend(AttentionBackend): block_size=self.page_size, actual_seq_lengths_kv=self.forward_metadata.seq_lens_cpu_int, ) + attn_output = attn_output[:, :, : layer.tp_q_head_num, :] else: assert ( self.graph_mode == False diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_kda_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_kda_backend.py new file mode 100644 index 000000000..8292b2f18 --- /dev/null +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_kda_backend.py @@ -0,0 +1,634 @@ +import math +from typing import Optional + +import torch +from sgl_kernel_npu.fla.kda_chunk_delta_h import ( + chunk_gated_delta_rule_fwd_h_npu, +) +from sgl_kernel_npu.fla.kda_gate import fused_kda_gate_npu +from sgl_kernel_npu.fla.kda_prefill import ( + chunk_gla_fwd_o_gk_npu, + recompute_w_u_fwd_npu, +) +from sgl_kernel_npu.fla.kda_target_verify import kda_target_verify_npu +from sgl_kernel_npu.fla.solve_tril import solve_tril_npu +from sgl_kernel_npu.fla.utils import prepare_chunk_indices +from sgl_kernel_npu.mamba.causal_conv1d import ( + causal_conv1d_fn_npu, + causal_conv1d_update_npu, +) +from sgl_kernel_npu.mamba.causal_conv1d_verify import ( + causal_conv1d_linear_verify_npu, +) + +from sglang.kernels.ops.attention.fla.cumsum import chunk_local_cumsum +from sglang.kernels.ops.attention.fla.kda import chunk_kda_scaled_dot_kkt_fwd +from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd +from sglang.srt.layers.attention.linear.kda_backend import ( + KDAAttnBackend, + ragged_verify_dense_scatter_indices, +) +from sglang.srt.layers.radix_linear_attention import RadixLinearAttention +from sglang.srt.model_executor.forward_batch_info import ForwardBatch + +_LOG2_E = math.log2(math.e) + + +class _AscendKDAExtendKernel: + """Ascend-only KDA prefill decomposition backed by sgl-kernel-npu.""" + + def extend( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + ssm_states: torch.Tensor, + cache_indices: torch.Tensor, + query_start_loc: torch.Tensor, + return_intermediate_states: bool = False, + **kwargs, + ): + chunk_size = 64 + q = l2norm_fwd(q.contiguous()) + k = l2norm_fwd(k.contiguous()) + v = v.contiguous() + beta = beta.contiguous() + chunk_indices = prepare_chunk_indices(query_start_loc, chunk_size) + g = chunk_local_cumsum( + g.contiguous(), + chunk_size=chunk_size, + scale=_LOG2_E, + cu_seqlens=query_start_loc, + chunk_indices=chunk_indices, + ) + + triangular, query_key = chunk_kda_scaled_dot_kkt_fwd( + q=q, + k=k, + gk=g, + beta=beta, + scale=k.shape[-1] ** -0.5, + cu_seqlens=query_start_loc, + output_dtype=torch.float32, + ) + triangular = solve_tril_npu( + A=triangular, + cu_seqlens=query_start_loc, + output_dtype=k.dtype, + ) + w, u, gated_k = recompute_w_u_fwd_npu( + k=k, + v=v, + beta=beta, + A=triangular, + gk=g, + cu_seqlens=query_start_loc, + chunk_indices=chunk_indices, + ) + del triangular + chunk_states, new_values = chunk_gated_delta_rule_fwd_h_npu( + k=gated_k, + w=w, + u=u, + gk=g, + initial_state=ssm_states, + initial_state_indices=cache_indices, + cu_seqlens=query_start_loc, + chunk_indices=chunk_indices, + use_exp2=True, + ) + del w, u, gated_k + out = chunk_gla_fwd_o_gk_npu( + q=q, + v=new_values, + g=g, + A=query_key, + h=chunk_states, + out=v, + scale=k.shape[-1] ** -0.5, + cu_seqlens=query_start_loc, + chunk_size=chunk_size, + chunk_indices=chunk_indices, + ) + del query_key, new_values + if return_intermediate_states: + return out, chunk_states.transpose(-1, -2).contiguous() + return out + + +class AscendKDAAttnBackend(KDAAttnBackend): + """Ascend implementation of Kimi Delta Attention. + + The model, scheduler, metadata, and non-operator control flow stay in the + shared KDA backend. This class contains only the layout and operator + differences required by Ascend. + """ + + supports_speculative_conv_state_snapshots: bool = True + + def __init__(self, model_runner): + super().__init__(model_runner) + # The NPU pool is allocated directly as [pool, channels, window]. + self.conv_states_shape = ( + model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape + ) + self.kernel_dispatcher.extend_kernel = _AscendKDAExtendKernel() + + def forward_decode( + self, + layer: RadixLinearAttention, + forward_batch: ForwardBatch, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + **kwargs, + ): + """Run KDA decode against the native channel-first Ascend cache.""" + assert isinstance(mixed_qkv, torch.Tensor) + layer_cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id) + conv_states = layer_cache.conv[0] + ssm_states = layer_cache.temporal + query_start_loc = self.forward_metadata.query_start_loc + cache_indices = self.forward_metadata.mamba_cache_indices + + qkv = causal_conv1d_update_npu( + mixed_qkv, + conv_states, + layer.conv_weights, + layer.bias, + activation="silu", + conv_state_indices=cache_indices, + ) + + if self.kernel_dispatcher.supports_packed_decode: + assert qkv.shape[0] == cache_indices.shape[0], ( + "KDA packed decode requires one token per sequence (T=1): " + f"got {qkv.shape[0]} tokens for {cache_indices.shape[0]} requests." + ) + core_attn_out = self.kernel_dispatcher.packed_decode( + mixed_qkv=qkv, + a=a, + b=b, + A_log=layer.A_log, + dt_bias=layer.dt_bias, + scale=layer.head_k_dim**-0.5, + ssm_states=ssm_states, + cache_indices=cache_indices, + num_v_heads=layer.num_v_heads, + head_v_dim=layer.head_v_dim, + lower_bound=layer.lower_bound, + replayssm_d=layer_cache.replayssm_d, + replayssm_k=layer_cache.replayssm_k, + replayssm_g=layer_cache.replayssm_g, + replayssm_write_pos=getattr( + self.forward_metadata, "replayssm_write_pos", None + ), + replayssm_force_flush=getattr( + self.forward_metadata, "replayssm_force_flush", None + ), + ) + else: + q, k, v = qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1) + q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0) + k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0) + v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0) + core_attn_out = self.kernel_dispatcher.decode( + q=q, + k=k, + v=v, + a=a, + b=b, + A_log=layer.A_log, + dt_bias=layer.dt_bias, + ssm_states=ssm_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + lower_bound=layer.lower_bound, + ) + + self._track_mamba_state_decode( + forward_batch, + conv_states, + ssm_states, + cache_indices, + layer.layer_id, + ) + return core_attn_out + + def forward_extend( + self, + layer: RadixLinearAttention, + forward_batch: ForwardBatch, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + **kwargs, + ): + """Run Ascend prefill without changing the shared KDA backend.""" + assert isinstance(mixed_qkv, torch.Tensor) + if forward_batch.forward_mode.is_target_verify(): + return self._forward_target_verify(layer, forward_batch, mixed_qkv, a, b) + + query_start_loc = self.forward_metadata.query_start_loc + cache_indices = self.forward_metadata.mamba_cache_indices + cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id) + conv_states = cache.conv[0] + ssm_states = cache.temporal + + if forward_batch.extend_prefix_lens is None: + raise RuntimeError( + "extend_prefix_lens cannot be None in non-TARGET_VERIFY mode." + ) + has_initial_state = forward_batch.extend_prefix_lens > 0 + + if self.forward_metadata.has_mamba_track_mask: + conv_states[self.forward_metadata.conv_states_mask_indices] = mixed_qkv[ + self.forward_metadata.track_conv_indices + ].transpose(-1, -2) + + splits = [layer.q_dim, layer.k_dim, layer.v_dim] + q, k, v = mixed_qkv.transpose(0, 1).split(splits, dim=0) + q_conv_weight, k_conv_weight, v_conv_weight = layer.conv_weights.split( + splits, dim=0 + ) + q_conv_state, k_conv_state, v_conv_state = conv_states.split(splits, dim=-2) + if layer.bias is not None: + q_bias, k_bias, v_bias = layer.bias.split(splits, dim=0) + else: + q_bias, k_bias, v_bias = None, None, None + + conv_kwargs = dict( + has_initial_state=has_initial_state, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + seq_lens_cpu=forward_batch.extend_seq_lens_cpu, + ) + q = self._causal_conv1d_extend( + q, q_conv_weight, q_bias, q_conv_state, **conv_kwargs + ) + k = self._causal_conv1d_extend( + k, k_conv_weight, k_bias, k_conv_state, **conv_kwargs + ) + v = self._causal_conv1d_extend( + v, v_conv_weight, v_bias, v_conv_state, **conv_kwargs + ) + + q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0) + k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0) + v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0) + g, beta, extend_A_log, extend_dt_bias = self._prepare_extend_gate_inputs( + layer, a, b + ) + track_ssm = self.forward_metadata.has_mamba_track_mask + core_attn_out = self.kernel_dispatcher.extend( + q=q, + k=k, + v=v, + g=g, + beta=beta, + ssm_states=ssm_states, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + A_log=extend_A_log, + dt_bias=extend_dt_bias, + lower_bound=layer.lower_bound, + extend_seq_lens_cpu=forward_batch.extend_seq_lens_cpu, + is_spec_decode=forward_batch.forward_mode.is_draft_extend_v2(), + return_intermediate_states=track_ssm, + track_ssm_h_src=( + self.forward_metadata.track_ssm_h_src if track_ssm else None + ), + ) + if track_ssm: + core_attn_out, h = core_attn_out + self._track_mamba_state_extend( + forward_batch, h, ssm_states, self.forward_metadata + ) + return core_attn_out + + def _causal_conv1d_extend( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: Optional[torch.Tensor], + state: torch.Tensor, + *, + has_initial_state: torch.Tensor, + cache_indices: torch.Tensor, + query_start_loc: torch.Tensor, + seq_lens_cpu: torch.Tensor, + ) -> torch.Tensor: + # The Ascend varlen kernel pads in the weight dtype. K3 keeps its + # weights in FP32 and its persistent convolution cache in BF16, so use + # a compact FP32 working set for the active rows and cast it back. + local_indices = torch.arange( + cache_indices.shape[0], + device=cache_indices.device, + dtype=cache_indices.dtype, + ) + state_work = state.index_select(0, cache_indices.to(torch.int64)) + state_work = state_work.to(weight.dtype).contiguous() + out = causal_conv1d_fn_npu( + x.to(weight.dtype), + weight, + bias, + activation="silu", + conv_states=state_work, + has_initial_state=has_initial_state, + cache_indices=local_indices, + query_start_loc=query_start_loc, + seq_lens_cpu=seq_lens_cpu, + ) + state.index_copy_(0, cache_indices.to(torch.int64), state_work.to(state.dtype)) + return out.to(x.dtype).transpose(0, 1) + + def _prepare_extend_gate_inputs( + self, + layer: RadixLinearAttention, + g: torch.Tensor, + beta: torch.Tensor, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + Optional[torch.Tensor], + Optional[torch.Tensor], + ]: + """Apply the Ascend prefill gate contract. + + The checkpoint was validated with FP32 gate activation before + ``chunk_kda``. Keeping this platform override here leaves the shared + GPU model/backend paths unchanged. + """ + preactivated_g = fused_kda_gate_npu( + g.flatten(-2), + layer.A_log, + layer.head_k_dim, + gate_bias=layer.dt_bias, + lower_bound=layer.lower_bound, + ) + return preactivated_g, beta, None, None + + def _forward_target_verify( + self, + layer: RadixLinearAttention, + forward_batch: ForwardBatch, + mixed_qkv: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + ) -> torch.Tensor: + """Run fixed-width DSpark verify with Ascend-native state snapshots.""" + metadata = self.forward_metadata + seq_len = mixed_qkv.shape[0] + query_start_loc = metadata.query_start_loc + cache_indices = metadata.mamba_cache_indices + + cache = self.req_to_token_pool.mamba2_layer_cache(layer.layer_id) + intermediate_state = cache.intermediate_ssm + if intermediate_state is None: + raise RuntimeError( + "Ascend KDA target verify requires speculative Mamba scratch." + ) + + draft_token_num = forward_batch.spec_info.draft_token_num + batch_size = query_start_loc.shape[0] - 1 + num_dense_tokens = batch_size * draft_token_num + ragged_layout = forward_batch.spec_info.ragged_verify_layout + if ragged_layout is None and seq_len == num_dense_tokens: + dense_token_indices = None + dense_qkv = mixed_qkv.view(batch_size, draft_token_num, -1) + dense_a = a + dense_b = b + else: + dense_token_indices = ragged_verify_dense_scatter_indices( + query_start_loc=query_start_loc, + seq_len=seq_len, + draft_token_num=draft_token_num, + ) + dense_qkv = self._scatter_tokens_to_dense( + mixed_qkv, dense_token_indices, num_dense_tokens + ).view(batch_size, draft_token_num, -1) + dense_a = self._scatter_gate_to_dense( + a, dense_token_indices, num_dense_tokens + ) + dense_b = self._scatter_gate_to_dense( + b, dense_token_indices, num_dense_tokens + ) + + intermediate_indices = self.verify_intermediate_state_indices[:batch_size] + processed_qkv = causal_conv1d_linear_verify_npu( + dense_qkv.transpose(1, 2).contiguous(), + cache.conv[0], + layer.conv_weights, + layer.bias, + cache_indices[:batch_size], + cache.intermediate_conv_window[0], + intermediate_indices, + activation="silu", + update_persistent_state=False, + ) + processed_qkv = processed_qkv.transpose(1, 2).reshape(num_dense_tokens, -1) + q, k, v = processed_qkv.split([layer.q_dim, layer.k_dim, layer.v_dim], dim=-1) + q = q.unflatten(-1, (-1, layer.head_q_dim)).unsqueeze(0) + k = k.unflatten(-1, (-1, layer.head_k_dim)).unsqueeze(0) + v = v.unflatten(-1, (-1, layer.head_v_dim)).unsqueeze(0) + + # Activate the forget gate and beta in FP32 before entering the + # recurrent kernel to match the checkpoint's verify contract. + # This stays in the Ascend backend so shared/GPU model code is unchanged. + preactivated_a = fused_kda_gate_npu( + dense_a.flatten(-2), + layer.A_log, + layer.head_k_dim, + gate_bias=layer.dt_bias, + lower_bound=layer.lower_bound, + ) + preactivated_b = dense_b.float().sigmoid() + out = kda_target_verify_npu( + A_log=layer.A_log, + dt_bias=layer.dt_bias, + q=q, + k=k, + v=v, + a=preactivated_a, + b=preactivated_b, + initial_state_source=cache.temporal, + initial_state_indices=cache_indices[:batch_size], + intermediate_states_buffer=intermediate_state, + intermediate_state_indices=intermediate_indices, + cache_steps=draft_token_num, + lower_bound=None, + gates_are_preactivated=True, + ) + if dense_token_indices is None: + return out + padded_out = out.new_zeros(1, num_dense_tokens + 1, *out.shape[2:]) + padded_out[:, :num_dense_tokens] = out + return padded_out[:, dense_token_indices] + + @staticmethod + def _scatter_tokens_to_dense( + value: torch.Tensor, + dense_token_indices: torch.Tensor, + num_dense_tokens: int, + ) -> torch.Tensor: + dense = value.new_zeros((num_dense_tokens + 1, *value.shape[1:])) + dense.index_copy_(0, dense_token_indices, value) + return dense[:num_dense_tokens] + + @classmethod + def _scatter_gate_to_dense( + cls, + value: torch.Tensor, + dense_token_indices: torch.Tensor, + num_dense_tokens: int, + ) -> torch.Tensor: + has_leading_singleton = value.ndim >= 2 and value.shape[0] == 1 + token_value = value.squeeze(0) if has_leading_singleton else value + dense = cls._scatter_tokens_to_dense( + token_value, dense_token_indices, num_dense_tokens + ) + return dense.unsqueeze(0) if has_leading_singleton else dense + + +class AscendKDAHybridLinearAttnBackend: + """KDA-specific hybrid backend with strided destination state mover. + + ``AscendHybridLinearAttnBackend`` uses ``move_intermediate_cache`` which + assumes a contiguous destination layout. KDA's temporal SSM state on NPU + is transposed (-1, -2) and requires the strided variant + ``move_intermediate_cache_kda`` to preserve correct (V, K) indexing. + + This class overrides only ``update_mamba_state_after_mtp_verify`` to + substitute the KDA-aware mover; the rest of the hybrid behaviour is + inherited unchanged. + """ + + def __new__(cls, *args, **kwargs): + # Delay importing AscendHybridLinearAttnBackend to avoid circular deps. + from sglang.srt.hardware_backend.npu.attention.ascend_hybrid_linear_attn_backend import ( + AscendHybridLinearAttnBackend as _Base, + ) + + # Dynamically create a subclass of _Base with our override. + class _AscendKDAHybrid(_Base): + def update_mamba_state_after_mtp_verify( + self, + last_correct_step_indices, + mamba_track_indices, + mamba_steps_to_track, + model, + req_pool_indices=None, + ): + from sgl_kernel_npu.mamba.mamba_state_update_triton import ( + conv_state_rollback, + move_intermediate_cache_kda, + ) + from sgl_kernel_npu.mamba.speculative_state_scatter import ( + speculative_state_scatter_npu, + ) + + del req_pool_indices + request_number = last_correct_step_indices.shape[0] + + state_indices_tensor = ( + self.linear_attn_backend.forward_metadata.mamba_cache_indices[ + :request_number + ] + ) + + mamba_caches = ( + self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers() + ) + + conv_states = mamba_caches.conv[0] + ssm_states = mamba_caches.temporal + intermediate_state_cache = mamba_caches.intermediate_ssm + dst_indices_tensor = state_indices_tensor.to(torch.int32) + src_indices_tensor = torch.arange( + dst_indices_tensor.shape[0], + device=dst_indices_tensor.device, + dtype=torch.int32, + ) + last_steps = last_correct_step_indices.to(torch.int32) + + move_intermediate_cache_kda( + ssm_states, + intermediate_state_cache, + dst_indices_tensor, + src_indices_tensor, + last_steps, + h_block_size=1, + ) + draft_token_num = intermediate_state_cache.shape[2] + has_conv_snapshots = getattr( + self.linear_attn_backend, + "supports_speculative_conv_state_snapshots", + False, + ) + if has_conv_snapshots: + intermediate_conv_window_cache = ( + mamba_caches.intermediate_conv_window[0] + ) + speculative_state_scatter_npu( + conv_states, + intermediate_conv_window_cache, + dst_indices_tensor, + src_indices_tensor, + last_steps, + ) + if mamba_track_indices is not None: + assert mamba_steps_to_track is not None + mamba_track_indices = mamba_track_indices.to(torch.int32) + mamba_steps_to_track = mamba_steps_to_track.to(torch.int32) + + move_intermediate_cache_kda( + ssm_states, + intermediate_state_cache, + mamba_track_indices, + src_indices_tensor, + mamba_steps_to_track, + h_block_size=1, + ) + + if has_conv_snapshots: + speculative_state_scatter_npu( + conv_states, + intermediate_conv_window_cache, + mamba_track_indices, + src_indices_tensor, + mamba_steps_to_track, + ) + else: + track_mask = mamba_steps_to_track >= 0 + track_indices = mamba_track_indices[track_mask] + if track_indices.numel() > 0: + conv_states[:, track_indices] = conv_states[ + :, dst_indices_tensor[track_mask] + ] + + if not has_conv_snapshots: + if dst_indices_tensor.numel() > 0: + conv_state_rollback( + conv_states, + dst_indices_tensor, + last_steps, + draft_token_num, + ) + + if ( + mamba_track_indices is not None + and mamba_track_indices.numel() > 0 + ): + conv_state_rollback( + conv_states, + mamba_track_indices, + mamba_steps_to_track, + draft_token_num, + ) + + return + + return _AscendKDAHybrid(*args, **kwargs) diff --git a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py index a685406f2..bfda4ebf5 100644 --- a/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py +++ b/python/sglang/srt/hardware_backend/npu/memory_pool_npu.py @@ -3,6 +3,7 @@ from typing import TYPE_CHECKING, Optional import torch from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE +from sglang.srt.environ import envs from sglang.srt.mem_cache.memory_pool import ( MHATokenToKVPool, MLATokenToKVPool, @@ -20,19 +21,26 @@ if is_npu(): def _init_npu_conv_state( - conv_state_in, conv_state_shape, speculative_num_draft_tokens: Optional[int] = None + conv_state_in, + conv_state_shape, + speculative_num_draft_tokens: Optional[int] = None, + is_kda: bool = False, ): extra_conv_len = 0 if speculative_num_draft_tokens is not None: extra_conv_len = speculative_num_draft_tokens - 1 - # conv_state shape (layers, pool_size, conv_wind + draft_step, dim) for conv1d ascendc ops require dim as last dim + # Mamba shapes are (channels, window), while KDA shapes are + # (window, channels). NPU kernels consume KDA state as + # [layers, pool, channels, window] and other Mamba state as + # [layers, pool, window, channels]. KDA keeps the base window fixed; + # speculative per-step windows live in the intermediate cache. conv_state = [ torch.zeros( size=( conv_state_in.shape[0], conv_state_in.shape[1], - conv_shape[1] + extra_conv_len, + conv_shape[1] if is_kda else conv_shape[1] + extra_conv_len, conv_shape[0], ), dtype=conv_state_in.dtype, @@ -66,6 +74,9 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): **kwargs, ): self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False") + self.use_triton_prefix_kv_cache_store = ( + envs.SGLANG_NPU_USE_TRITON_PREFIX_KV_CACHE_STORE.get() + ) super().__init__( size=size, page_size=page_size, @@ -201,16 +212,36 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): if self.use_fia: k_buffer_layer = self.k_buffer[layer_id - self.start_layer] v_buffer_layer = self.v_buffer[layer_id - self.start_layer] + num_rows = loc.numel() + expected_k_numel = num_rows * self.head_num * self.head_dim + expected_v_numel = num_rows * self.head_num * self.v_head_dim + if ( + cache_k.numel() != expected_k_numel + or cache_v.numel() != expected_v_numel + ): + raise ValueError( + "NPU FIA KV scatter row mismatch: " + f"loc_rows={num_rows}, cache_k_shape={tuple(cache_k.shape)}, " + f"cache_v_shape={tuple(cache_v.shape)}, " + f"head_num={self.head_num}, head_dim={self.head_dim}, " + f"v_head_dim={self.v_head_dim}." + ) + # aclnnScatterNdUpdate on the deployed CANN rejects the otherwise + # valid 4-D [slot, 1, head, dim] update during tiling. Flatten only + # the singleton FIA layout axis and scatter through an equivalent + # 3-D view; the underlying KV storage and attention layout stay + # unchanged. + loc_indices = loc.contiguous().view(-1, 1) torch_npu.npu_scatter_nd_update_( - k_buffer_layer, - loc.view(-1, 1), - cache_k.view(-1, 1, self.head_num, self.head_dim), + k_buffer_layer.view(-1, self.head_num, self.head_dim), + loc_indices, + cache_k.contiguous().view(num_rows, self.head_num, self.head_dim), ) torch_npu.npu_scatter_nd_update_( - v_buffer_layer, - loc.view(-1, 1), - cache_v.view(-1, 1, self.head_num, self.v_head_dim), + v_buffer_layer.view(-1, self.head_num, self.v_head_dim), + loc_indices, + cache_v.contiguous().view(num_rows, self.head_num, self.v_head_dim), ) else: loc = loc.to(torch.int32) @@ -226,6 +257,80 @@ class NPUMHATokenToKVPool(MHATokenToKVPool): slot_indices=loc, ) + def set_kv_buffer_prefix_valid( + self, + layer: "RadixAttention", + loc_2d: torch.Tensor, + commit_lens: torch.Tensor, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + k_scale: Optional[float] = None, + v_scale: Optional[float] = None, + layer_id_override: Optional[int] = None, + ): + if not self.use_triton_prefix_kv_cache_store: + return super().set_kv_buffer_prefix_valid( + layer, + loc_2d, + commit_lens, + cache_k, + cache_v, + k_scale, + v_scale, + layer_id_override, + ) + + if layer_id_override is not None: + layer_id = layer_id_override + else: + layer_id = layer.layer_id + if loc_2d.ndim != 2: + raise ValueError(f"loc_2d must be rank-2, got {tuple(loc_2d.shape)}") + + num_rows = loc_2d.numel() + if ( + cache_k.numel() != num_rows * self.head_num * self.head_dim + or cache_v.numel() != num_rows * self.head_num * self.v_head_dim + ): + raise ValueError( + "dense NPU KV rows must match loc_2d size: " + f"cache_k={tuple(cache_k.shape)}, cache_v={tuple(cache_v.shape)}, " + f"loc_2d={tuple(loc_2d.shape)}" + ) + + if cache_k.dtype != self.dtype: + if k_scale is not None: + cache_k.div_(k_scale) + if v_scale is not None: + cache_v.div_(v_scale) + cache_k = cache_k.to(self.dtype) + cache_v = cache_v.to(self.dtype) + if self.store_dtype != self.dtype: + cache_k = cache_k.contiguous().view(self.store_dtype) + cache_v = cache_v.contiguous().view(self.store_dtype) + + k_buffer_layer = self.k_buffer[layer_id - self.start_layer] + v_buffer_layer = self.v_buffer[layer_id - self.start_layer] + if loc_2d.device != k_buffer_layer.device: + loc_2d = loc_2d.to(device=k_buffer_layer.device, non_blocking=True) + if commit_lens.device != k_buffer_layer.device: + commit_lens = commit_lens.to( + device=k_buffer_layer.device, non_blocking=True + ) + self._debug_prefix_valid_backend = "npu_triton" + from sgl_kernel_npu.mem_cache.kv_cache_store import ( + store_kv_cache_prefix_valid_npu_triton, + ) + + store_kv_cache_prefix_valid_npu_triton( + k_buffer_layer.view(-1, self.head_num, self.head_dim), + v_buffer_layer.view(-1, self.head_num, self.v_head_dim), + cache_k.reshape(num_rows, self.head_num, self.head_dim), + cache_v.reshape(num_rows, self.head_num, self.v_head_dim), + loc_2d, + commit_lens, + ) + def _chunk_copy_npu_to_cpu(self, buf_of_layers, indices): chunk_size = self.cpu_offloading_chunk_size out = [] @@ -294,10 +399,10 @@ class NPUMLATokenToKVPool(MLATokenToKVPool): dtype: torch.dtype, kv_lora_rank: int, qk_rope_head_dim: int, - index_head_dim: Optional[int], layer_num: int, device: str, enable_memory_saver: bool, + index_head_dim: Optional[int] = None, start_layer: Optional[int] = None, end_layer: Optional[int] = None, ): diff --git a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py index b3f60a561..55949768f 100644 --- a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py +++ b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py @@ -203,8 +203,10 @@ def forward_mla_prepare_npu( k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1) k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1) else: - if qkv_latent.shape[0] < 65536 and not dsa_use_prefill_cp( - forward_batch + if ( + qkv_latent.shape[0] < 65536 + and not dsa_use_prefill_cp(forward_batch) + and not getattr(m, "_disable_npu_fused_split_qk_norm", False) ): q, k_nope, k_pe = fused_split_qk_norm( qkv_latent, @@ -216,6 +218,8 @@ def forward_mla_prepare_npu( eps=m.q_a_layernorm.variance_epsilon, ) else: + # The fused split+RMSNorm kernel is not numerically equivalent + # on Ascend. Keep the unfused path for models that opt out. q, latent_cache = qkv_latent.split( [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1, @@ -245,7 +249,8 @@ def forward_mla_prepare_npu( q_nope_out = q_nope_out.transpose(0, 1) - q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe) + if m.rotary_emb is not None: + q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe) if dsa_use_prefill_cp(forward_batch): # support allgather+rerrange @@ -297,14 +302,16 @@ def forward_mla_core_npu( attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank) - attn_bmm_output = torch.empty( - (attn_output.shape[0], m.num_local_heads, m.v_head_dim), - dtype=attn_output.dtype, - device=attn_output.device, - ) - attn_output = attn_output.contiguous() - torch.ops.npu.batch_matmul_transpose(attn_output, m.w_vc, attn_bmm_output) + # torch.ops.npu.batch_matmul_transpose is not numerically equivalent for + # Kimi-K3, so use the numerically validated torch_npu implementation. + attn_bmm_output = torch_npu.npu_transpose_batchmatmul( + attn_output, + m.w_vc, + perm_x1=(1, 0, 2), + perm_x2=(0, 1, 2), + perm_y=(1, 0, 2), + ) attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim) output, _ = m.o_proj(attn_bmm_output) @@ -380,8 +387,10 @@ def forward_dsa_prepare_npu( if q_event is not None: torch.npu.current_stream().wait_event(q_event) else: - if fused_qkv_a_proj_out.shape[0] < 65535 and not dsa_use_prefill_cp( - forward_batch + if ( + fused_qkv_a_proj_out.shape[0] < 65535 + and not dsa_use_prefill_cp(forward_batch) + and not getattr(m, "_disable_npu_fused_split_qk_norm", False) ): q_lora, k_nope, k_pe = fused_split_qk_norm( fused_qkv_a_proj_out, @@ -393,6 +402,8 @@ def forward_dsa_prepare_npu( eps=m.q_a_layernorm.variance_epsilon, ) else: + # Keep the numerically validated unfused path for models that + # explicitly opt out of the fused split and RMSNorm kernel. q, latent_cache = fused_qkv_a_proj_out.split( [m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1 ) diff --git a/python/sglang/srt/hardware_backend/npu/moe/activation.py b/python/sglang/srt/hardware_backend/npu/moe/activation.py index 981191c9f..bfad1c65d 100644 --- a/python/sglang/srt/hardware_backend/npu/moe/activation.py +++ b/python/sglang/srt/hardware_backend/npu/moe/activation.py @@ -85,6 +85,39 @@ class NPUSwigluDeepEPKernel(BaseActivation): return hidden_states, None +class NPUSitu(BaseActivation): + """SiTU activation and optional INT8 requantization for grouped rows.""" + + def __init__( + self, + *, + need_quant: bool, + beta: float = 4.0, + linear_beta: Optional[float] = 25.0, + ): + from sgl_kernel_npu.activation.situ import situ + + self.situ = situ + self.need_quant = need_quant + self.beta = float(beta) + self.linear_beta = None if linear_beta is None else float(linear_beta) + + def _apply_activation( + self, + hidden_states: torch.Tensor, + group_list: torch.Tensor, + group_list_type: int, + ): + return self.situ( + hidden_states, + group_list, + group_list_type, + need_quant=self.need_quant, + beta=self.beta, + linear_beta=self.linear_beta, + ) + + class NPUGeluAndMul(BaseActivation): def __init__(self): self._gelu = GeluAndMul() diff --git a/python/sglang/srt/layers/activation.py b/python/sglang/srt/layers/activation.py index 61ef0724d..fce55bbb0 100644 --- a/python/sglang/srt/layers/activation.py +++ b/python/sglang/srt/layers/activation.py @@ -208,6 +208,11 @@ class SituAndMul(BaseFusedOp): return situ_and_mul(x, None, self.beta, self.linear_beta) + def forward_npu(self, x: torch.Tensor) -> torch.Tensor: + from sgl_kernel_npu.activation.situ import situ_and_mul + + return situ_and_mul(x) + def forward_cpu(self, x: torch.Tensor) -> torch.Tensor: return self.forward_native(x) diff --git a/python/sglang/srt/layers/attention/attention_registry.py b/python/sglang/srt/layers/attention/attention_registry.py index 603ea0c99..84237b424 100644 --- a/python/sglang/srt/layers/attention/attention_registry.py +++ b/python/sglang/srt/layers/attention/attention_registry.py @@ -453,7 +453,16 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac else: linear_attn_backend = Mamba2AttnBackend(runner) elif kimi_linear_config(runner.model_config) is not None: - linear_attn_backend = KDAAttnBackend(runner) + if _is_npu: + from sglang.srt.hardware_backend.npu.attention.ascend_kda_backend import ( + AscendKDAAttnBackend, + AscendKDAHybridLinearAttnBackend, + ) + + linear_attn_backend = AscendKDAAttnBackend(runner) + hybrid_backend_cls = AscendKDAHybridLinearAttnBackend + else: + linear_attn_backend = KDAAttnBackend(runner) elif hybrid_lightning_config(runner.model_config) is not None: linear_attn_backend = LightningAttentionBackend(runner) else: diff --git a/python/sglang/srt/layers/attn_residual.py b/python/sglang/srt/layers/attn_residual.py index 491a7b4f8..2e3cc94e4 100644 --- a/python/sglang/srt/layers/attn_residual.py +++ b/python/sglang/srt/layers/attn_residual.py @@ -24,7 +24,7 @@ import triton.language as tl from sglang.srt.layers.layernorm import RMSNorm from sglang.srt.layers.linear import ReplicatedLinear -from sglang.srt.utils import is_hip +from sglang.srt.utils import is_hip, is_npu _BLOCK_H: int = 1024 # H = 7168 = 7 x 1024 _MAX_ROWS: int = 16 # next_pow2(8 + 1), K3 has <= 8 snapshots @@ -37,6 +37,8 @@ def _use_fast(hidden_size: int) -> bool: """The TMA kernel needs SM100+ (tcgen05, cp.async.bulk) and its H=7168 template instantiation; everything else takes the triton pipeline.""" global _FAST_SUPPORTED + if is_npu(): + return False if _FAST_SUPPORTED is None: major, _ = torch.cuda.get_device_capability() _FAST_SUPPORTED = major >= 10 @@ -208,7 +210,19 @@ def _mix_fused( ) -> torch.Tensor: """Triton score + combine pair: returns the pre-norm mixture.""" T, H = prefix_sum.shape + if T == 0: + return prefix_sum cw = get_cw(score_proj, score_norm) + if is_npu(): + from sgl_kernel_npu.kimi_k3.attn_residual import mix_fused + + return mix_fused( + prefix_sum, + bank, + nvb, + cw, + score_norm.variance_epsilon, + ) n_h_blocks = H // _BLOCK_H # Step 1: score each row (2D grid, full row-parallelism) @@ -394,6 +408,8 @@ def _aggregate( into bank[:, nvb, :]); the triton path keeps the standalone .write() copy — the caller (AttnResidual.forward) owns that fallback. """ + if prefix_sum.shape[0] == 0: + return prefix_sum if _use_fast(prefix_sum.shape[1]): return _aggregate_fast( prefix_sum, diff --git a/python/sglang/srt/layers/moe/moe_runner/ascend.py b/python/sglang/srt/layers/moe/moe_runner/ascend.py index 6dfee6448..c72c00eee 100644 --- a/python/sglang/srt/layers/moe/moe_runner/ascend.py +++ b/python/sglang/srt/layers/moe/moe_runner/ascend.py @@ -10,6 +10,7 @@ import torch from sglang.srt.hardware_backend.npu.moe.activation import ( AllGatherActivationWrapper, NPUGeluAndMul, + NPUSitu, NPUSwiglu, NPUSwigluDeepEPKernel, NPUSwigluOAI, @@ -101,7 +102,16 @@ class AscendRunnerCore(MoeRunnerCore): is_quant_kernel = isinstance( kernel, (NPUW4A8Int8MoEMethod, NPUW8A8Int8MoEMethod) ) - self.activation = NPUSwigluDeepEPKernel(need_quant=is_quant_kernel) + if config.activation == "situ": + self.activation = NPUSitu( + need_quant=is_quant_kernel, + beta=( + config.gemm1_alpha if config.gemm1_alpha is not None else 4.0 + ), + linear_beta=config.gemm1_clamp_limit, + ) + else: + self.activation = NPUSwigluDeepEPKernel(need_quant=is_quant_kernel) else: # Non‑DeepEP (ascend_tp) path # 1. Choose the base activation according to the quant method @@ -169,8 +179,11 @@ class AscendRunnerCore(MoeRunnerCore): ) # --- Activation --- - # The DeepEP kernel expects extra dispatch metadata - if isinstance(self.activation, NPUSwigluDeepEPKernel): + # Grouped-row activations require dispatch metadata. + if isinstance( + self.activation, + (NPUSwigluDeepEPKernel, NPUSitu), + ): hidden_states, pertoken_scale = self.activation._apply_activation( hidden_states, group_list=expert_tokens, diff --git a/python/sglang/srt/layers/quantization/modelslim/modelslim.py b/python/sglang/srt/layers/quantization/modelslim/modelslim.py index 227e20010..036c7392d 100644 --- a/python/sglang/srt/layers/quantization/modelslim/modelslim.py +++ b/python/sglang/srt/layers/quantization/modelslim/modelslim.py @@ -90,6 +90,8 @@ class ModelSlimConfig(QuantizationConfig): Config class for ModelSlim Quantization, a NPU-specific quantization type. """ + supports_kimi_k3_quantized_latent_projections = True + def __init__(self, quant_config: Dict[str, Any] = {}): super().__init__() keys = [k for k in quant_config if isinstance(k, str)] @@ -127,6 +129,33 @@ class ModelSlimConfig(QuantizationConfig): def update_packed_modules_mapping(self, mapping: Dict[str, List[str]]) -> None: self.packed_modules_mapping.update(mapping) + def _quant_prefix_candidates(self, prefix: str) -> List[str]: + """Return checkpoint-name variants without copying the large config. + + Kimi-K3's upstream model uses ``mlp`` internally while its ModelSlim + checkpoint retains the Hugging Face ``block_sparse_moe`` hierarchy. + Some multimodal checkpoints also keep the outer ``language_model`` + prefix. Resolve those layout-only differences at the quantization + boundary. + """ + candidates = [prefix] + if ".mlp." in prefix: + candidates.append(prefix.replace(".mlp.", ".block_sparse_moe.")) + + for candidate in list(candidates): + if candidate.startswith("language_model."): + candidates.append(candidate.removeprefix("language_model.")) + else: + candidates.append(f"language_model.{candidate}") + + return list(dict.fromkeys(candidates)) + + def _resolve_quant_prefix(self, prefix: str) -> str: + for candidate in self._quant_prefix_candidates(prefix): + if candidate + ".weight" in self.quant_description: + return candidate + return prefix + def get_linear_method(self) -> ModelSlimLinearMethod: return ModelSlimLinearMethod(self) @@ -176,10 +205,7 @@ class ModelSlimConfig(QuantizationConfig): prefix_in_quant_config = prefix.replace( proj_name, packed_modules_mapping_subset[proj_name][0] ) - # Verify the remapped prefix exists in quant_description. - # If not (e.g. json uses fused name as-is), fall back to original. - if prefix_in_quant_config + ".weight" not in self.quant_description: - prefix_in_quant_config = prefix + prefix_in_quant_config = self._resolve_quant_prefix(prefix_in_quant_config) if self.is_layer_skipped( prefix, packed_modules_mapping_subset ) or self.is_layer_skipped(prefix, self.packed_modules_mapping): @@ -217,6 +243,7 @@ class ModelSlimConfig(QuantizationConfig): ("W4A4_MXFP4", ModelSlimMXFP4Scheme), ] + prefix = self._resolve_quant_prefix(prefix) quant_schemes = [self.quant_description.get(prefix + ".weight", "")] for scheme_name, scheme_class in linear_quant_schemes: @@ -252,46 +279,54 @@ class ModelSlimConfig(QuantizationConfig): w13_scheme_name = None w2_scheme_name = None - for gate_name, up_name, down_name in naming_conventions: - w13_keys = [ - f"{prefix}.0.{gate_name}.weight", - f"{prefix}.0.{up_name}.weight", - ] - w2_key = f"{prefix}.0.{down_name}.weight" - w13_entries = { - key: self.quant_description[key] - for key in w13_keys - if key in self.quant_description - } - if w13_entries and w2_key in self.quant_description: - w13_names = list(w13_entries.values()) - # For w13, both projections must agree on the scheme - unique_w13 = set(w13_names) - if len(unique_w13) > 1: - raise ValueError( - f"Mismatched ModelSlim quantization for W13 in layer {prefix}: " - f"{w13_entries}" - ) - w13_scheme_name = w13_names[0] - w2_scheme_name = self.quant_description[w2_key] + resolved_prefix = prefix + for candidate in self._quant_prefix_candidates(prefix): + for gate_name, up_name, down_name in naming_conventions: + w13_keys = [ + f"{candidate}.0.{gate_name}.weight", + f"{candidate}.0.{up_name}.weight", + ] + w2_key = f"{candidate}.0.{down_name}.weight" + w13_entries = { + key: self.quant_description[key] + for key in w13_keys + if key in self.quant_description + } + if w13_entries and w2_key in self.quant_description: + w13_names = list(w13_entries.values()) + # For w13, both projections must agree on the scheme + unique_w13 = set(w13_names) + if len(unique_w13) > 1: + raise ValueError( + "Mismatched ModelSlim quantization for W13 in layer " + f"{prefix}: {w13_entries}" + ) + w13_scheme_name = w13_names[0] + w2_scheme_name = self.quant_description[w2_key] + resolved_prefix = candidate + break + if w13_scheme_name is not None: break if w13_scheme_name is None: # Build a helpful error message listing all attempted key patterns all_attempted = [] - for gate_name, up_name, down_name in naming_conventions: - w13_keys = [ - f"{prefix}.0.{gate_name}.weight", - f"{prefix}.0.{up_name}.weight", - ] - w2_key = f"{prefix}.0.{down_name}.weight" - w13_found = any(k in self.quant_description for k in w13_keys) - w2_found = w2_key in self.quant_description - status = ( - f"({gate_name}/{up_name}={'found' if w13_found else 'missing'}, " - f"{down_name}={'found' if w2_found else 'missing'})" - ) - all_attempted.append(status) + for candidate in self._quant_prefix_candidates(prefix): + for gate_name, up_name, down_name in naming_conventions: + w13_keys = [ + f"{candidate}.0.{gate_name}.weight", + f"{candidate}.0.{up_name}.weight", + ] + w2_key = f"{candidate}.0.{down_name}.weight" + w13_found = any(k in self.quant_description for k in w13_keys) + w2_found = w2_key in self.quant_description + status = ( + f"{candidate} " + f"({gate_name}/{up_name}=" + f"{'found' if w13_found else 'missing'}, " + f"{down_name}={'found' if w2_found else 'missing'})" + ) + all_attempted.append(status) raise ValueError( f"Missing ModelSlim MoE quantization description for layer {prefix}: " + "; ".join(all_attempted) @@ -306,7 +341,9 @@ class ModelSlimConfig(QuantizationConfig): def instantiate(name, weight_group): cls = scheme_map.get(name) if cls is None: - logger.warning(f"Unsupported scheme '{name}' for layer {prefix}") + logger.warning( + f"Unsupported scheme '{name}' for layer {resolved_prefix}" + ) return None return cls(self, weight_group) @@ -335,6 +372,7 @@ class ModelSlimConfig(QuantizationConfig): is_skipped = None for shard_prefix in shard_prefixes: + shard_prefix = self._resolve_quant_prefix(shard_prefix) is_shard_skipped = ( self.quant_description.get(shard_prefix + ".weight", "") == "FLOAT" ) @@ -348,6 +386,7 @@ class ModelSlimConfig(QuantizationConfig): "to have the same precision." ) else: + prefix = self._resolve_quant_prefix(prefix) is_skipped = self.quant_description.get(prefix + ".weight", "") == "FLOAT" assert is_skipped is not None diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index cecd7b940..261b2018e 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -558,7 +558,10 @@ class MambaPool: ) conv_state = _init_npu_conv_state( - conv_state[0], conv_state_shape, speculative_num_draft_tokens + conv_state[0], + conv_state_shape, + speculative_num_draft_tokens, + is_kda=cache_params.is_kda, ) if _is_cpu and _cpu_has_amx_support: @@ -759,6 +762,14 @@ class MambaPool: # Original dense layout (NPU/CPU, or EAGLE tree verify): one # [dim, K-1] window per draft token. # Shape: [num_layers, size+1, draft_tokens, dim, K-1] + dense_conv_shapes = [ + ( + (conv_shape[1], conv_shape[0]) + if _is_npu and cache_params.is_kda + else conv_shape + ) + for conv_shape in conv_state_shape + ] intermediate_conv_window_cache = [ torch.zeros( size=( @@ -771,7 +782,7 @@ class MambaPool: dtype=conv_dtype, device="cuda", ) - for conv_shape in conv_state_shape + for conv_shape in dense_conv_shapes ] self._intermediate_conv_window_phys = intermediate_conv_window_cache self.mamba_cache = self.SpeculativeState( diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 1ceb52f91..643ff7c5c 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -263,7 +263,17 @@ class EagerRunner(BaseRunner): if cp_v2_active: prepare_cp_forward(forward_batch) - if forward_batch.needs_forward_metadata_init() or cp_v2_active: + # Target verify can arrive with ``forward_metadata_ready`` set by an + # upstream/speculative planning step. That mark does not initialize + # the final target hybrid backend, and unlike a graph replay eager has + # no static metadata load to fill the gap. Re-plan target verify from + # the final batch every time; eager metadata is intentionally derived + # directly from the live ``spec_info`` tensors. + if ( + forward_batch.needs_forward_metadata_init() + or cp_v2_active + or forward_batch.forward_mode.is_target_verify() + ): if model_runner.ps.attn_dcp_size > 1 and hasattr( model_runner.model, "prepare_context_parallel_metadata_for_dcp" ): diff --git a/python/sglang/srt/models/dspark.py b/python/sglang/srt/models/dspark.py index 0eb78a98f..4116f527b 100644 --- a/python/sglang/srt/models/dspark.py +++ b/python/sglang/srt/models/dspark.py @@ -8,6 +8,7 @@ import torch.nn.functional as F from torch import nn from sglang.srt.distributed.communication_op import tensor_model_parallel_all_gather +from sglang.srt.environ import envs from sglang.srt.model_loader.weight_utils import default_weight_loader from sglang.srt.models.dflash import DFlashDraftModel from sglang.srt.speculative.dflash_utils import can_dflash_slice_qkv_weight @@ -554,6 +555,8 @@ class DSparkDraftMixin: Cached; None (per-layer fallback) when a QKV weight cannot be sliced (quantized) or layers disagree on norm epsilon / bias presence. """ + if not envs.SGLANG_DSPARK_STACKED_CTX_KV.get(): + return None cached = getattr(self, "_stacked_ctx_kv_cache", False) if cached is not False: return cached diff --git a/python/sglang/srt/models/kimi_k3.py b/python/sglang/srt/models/kimi_k3.py index 2c82e986a..9630872a9 100644 --- a/python/sglang/srt/models/kimi_k3.py +++ b/python/sglang/srt/models/kimi_k3.py @@ -38,6 +38,8 @@ from sglang.srt.layers.activation import SiluAndMul, SituAndMul from sglang.srt.layers.attn_residual import AttnResidual, aggregate_stream, get_cw from sglang.srt.layers.dcp.planner import prepare_decode_context_parallel_metadata from sglang.srt.layers.dp_attention import ( + attn_tp_all_gather_into_tensor, + attn_tp_reduce_scatter_tensor, dp_gather_replicate, dp_scatter, get_global_dp_buffer, @@ -117,7 +119,7 @@ from sglang.srt.runtime_context import ( get_parallel, get_server_args, ) -from sglang.srt.utils import is_blackwell_supported, is_hip, make_layers +from sglang.srt.utils import is_blackwell_supported, is_hip, is_npu, make_layers from sglang.srt.utils.common import ( BumpAllocator, add_prefix, @@ -129,7 +131,10 @@ from sglang.srt.utils.common import ( logger = logging.getLogger(__name__) _is_hip = is_hip() +_is_npu = is_npu() _aiter_k3_opt = get_bool_env_var("SGLANG_AITER_K3_OPT") +_k3_shared_experts_attn_tp = envs.SGLANG_K3_SHARED_EXPERTS_ATTN_TP.get() +_k3_dense_mlp_attn_tp = envs.SGLANG_K3_DENSE_MLP_ATTN_TP.get() def _cdiv(a: int, b: int) -> int: @@ -207,8 +212,6 @@ def _merge_weights_as_views( return merged, sizes -# DP attention helpers. -# # K3 cannot use LayerCommunicator: the attn-res aggregation kernels replace # input_layernorm / post_attention_layernorm, which the communicator expects # to own. Instead the MLP/MoE modules gather/scatter around their own body: @@ -217,8 +220,6 @@ def _merge_weights_as_views( # semantics (its internal all-reduces are unchanged and required — the latent # reduce must happen in latent space before the norm), and the delayed # prefix_sum add stays local, applied after the scatter back. - - def _dp_local_buffer_group(): """Symmetric-memory group for the local DP buffer (mirrors CommunicateSummableTensorPairFn._scatter_hidden_states).""" @@ -267,6 +268,20 @@ class KimiK3MLP(nn.Module): tp_size: Optional[int] = None, ) -> None: super().__init__() + # The Ascend path shards the dense MLP inside each attention-TP + # replica. The GPU K3 refactor instead gathers all DP rows and shards + # this one dense layer over the full TP group. Keep the GPU default, + # but allow the NPU launcher to retain the proven attention-TP layout + # without a device-type branch in shared model code. + self._dense_attn_tp = ( + _k3_dense_mlp_attn_tp + and is_dp_attention_enabled() + and tp_rank is None + and tp_size is None + ) + if self._dense_attn_tp: + tp_rank = get_parallel().attn_tp_rank + tp_size = get_parallel().attn_tp_size _tp_kwargs = ( dict(tp_rank=tp_rank, tp_size=tp_size) if tp_size is not None else {} ) @@ -284,6 +299,7 @@ class KimiK3MLP(nn.Module): bias=False, quant_config=quant_config, reduce_results=reduce_results, + use_dp_attention_reduce=self._dense_attn_tp, prefix=f"{prefix}.down_proj", **_tp_kwargs, ) @@ -308,7 +324,9 @@ class KimiK3MLP(nn.Module): # DP attention only when driven from the decoder layer (forward_batch # given); the shared-experts instance inside KimiK3MoE passes None and # runs on the already-gathered buffer. - use_dp = self._dp_attention and forward_batch is not None + use_dp = ( + self._dp_attention and forward_batch is not None and not self._dense_attn_tp + ) if use_dp: local_hidden_states = hidden_states hidden_states = get_global_dp_buffer(get_tp_group()) @@ -342,6 +360,8 @@ def _add3( return a + b from sglang.kernels.ops.elementwise import add3 + if not add3.covered(a, b, c): + return (a + b) + c return add3.add3(a, b, c, prefetch_bc=prefetch_bc) @@ -439,6 +459,7 @@ class KimiK3MoE(nn.Module): use_grouped_topk=True, num_expert_group=config.num_expert_group, topk_group=config.topk_group, + scoring_func=config.moe_router_activation_func, correction_bias=self.gate.e_score_correction_bias, quant_config=quant_config, routed_scaling_factor=self.routed_scaling_factor, @@ -488,7 +509,11 @@ class KimiK3MoE(nn.Module): # (DP attention) — with every global token dispatched exactly once. # No DP gather and no TP reduce is needed anywhere in the region. _a2a_backend = get_moe_a2a_backend() - self._ep_a2a = _a2a_backend.is_megamoe() or _a2a_backend.is_deepep() + self._ep_a2a = ( + _a2a_backend.is_megamoe() + or _a2a_backend.is_deepep() + or _a2a_backend.is_ascend_fuseep() + ) # Defer the trtllm-gen finalize (top-k weighted unpermute) out of the # MoE op and fuse it into the push all-reduce's staging pass @@ -506,7 +531,24 @@ class KimiK3MoE(nn.Module): # a2a: the block runs on partial batches (shard / DP-local rows), and # a TP-sharded partial sum could never be reduced across ranks that # hold different tokens. - self._shared_experts_tp1 = self._ep_a2a + self._shared_experts_tp1 = self._ep_a2a and not _k3_shared_experts_attn_tp + # NPU compatibility mode keeps DeepEP's DP-local token dispatch but + # uses the original TP-sharded shared MLP. Gather only that branch's + # inputs, then reduce-scatter its output back to the DP-local rows. + self._shared_experts_attn_tp_comm = ( + _k3_shared_experts_attn_tp + and self._ep_a2a + and self._dp_attention + and get_parallel().attn_tp_size > 1 + ) + shared_experts_tp_kwargs = {} + if self._shared_experts_tp1: + shared_experts_tp_kwargs = dict(tp_rank=0, tp_size=1) + elif self._shared_experts_attn_tp_comm: + shared_experts_tp_kwargs = dict( + tp_rank=get_parallel().attn_tp_rank, + tp_size=get_parallel().attn_tp_size, + ) if self.num_shared_experts is not None and self.num_shared_experts > 0: shared_intermediate_size = moe_intermediate_size * self.num_shared_experts self.shared_experts = KimiK3MLP( @@ -518,7 +560,7 @@ class KimiK3MoE(nn.Module): prefix=f"{prefix}.shared_experts", activation_situ_beta=config.activation_situ_beta, activation_situ_linear_beta=config.activation_situ_linear_beta, - **(dict(tp_rank=0, tp_size=1) if self._shared_experts_tp1 else {}), + **shared_experts_tp_kwargs, ) else: self.shared_experts = None @@ -538,16 +580,26 @@ class KimiK3MoE(nn.Module): # overlap than two streams. self._sbo_shared_overlap = ( self._ep_a2a + and not self._shared_experts_attn_tp_comm and self.shared_experts is not None and self.alt_stream is not None ) if self.use_latent_moe: + latent_quant_config = ( + quant_config + if getattr( + quant_config, + "supports_kimi_k3_quantized_latent_projections", + False, + ) + else None + ) self.routed_expert_down_proj = ReplicatedLinear( hidden_size, self.moe_hidden_size, bias=False, - quant_config=None, + quant_config=latent_quant_config, prefix=f"{prefix}.routed_expert_down_proj", ) self.routed_expert_norm = ( @@ -559,7 +611,7 @@ class KimiK3MoE(nn.Module): self.moe_hidden_size, hidden_size, bias=False, - quant_config=None, + quant_config=latent_quant_config, prefix=f"{prefix}.routed_expert_up_proj", ) else: @@ -608,6 +660,12 @@ class KimiK3MoE(nn.Module): """ if not self.use_latent_moe: return + # These merged layouts feed CUDA-only fused front kernels. Keeping the + # regular parameters on other devices avoids a large transient copy + # during post-load processing and leaves their native kernels in + # control of weight layout. + if _is_npu: + return if self.shared_experts is not None and get_moe_a2a_backend().is_none(): mods = [ self.shared_experts.gate_up_proj, @@ -774,10 +832,11 @@ class KimiK3MoE(nn.Module): if cfg.output_format is not TopKOutputFormat.STANDARD: return False # The kernel implements sigmoid scoring with bias-ranked ungrouped top-k. - # Do NOT test cfg.scoring_func: it defaults to "softmax" and TopK - # documents it as unused. What actually selects sigmoid is the - # grouped-topk-with-correction-bias route (DSv3 noaux_tc), which calls - # biased_grouped_topk and hardwires scoring_func="sigmoid". + # K3 passes moe_router_activation_func explicitly to TopK. The legacy + # GPU biased_grouped_topk path also hardwires sigmoid, but other platform + # implementations consume cfg.scoring_func directly. + if cfg.scoring_func != "sigmoid": + return False if not (cfg.use_grouped_topk and cfg.correction_bias is not None): return False if (cfg.num_expert_group or 1) > 1 or (cfg.topk_group or 1) > 1: @@ -900,8 +959,26 @@ class KimiK3MoE(nn.Module): return self._latent_norm(latent) return self._latent_norm(tensor_model_parallel_all_reduce(latent)) + def _forward_shared_experts(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Run TP-sharded shared experts while DeepEP tokens stay scattered.""" + if not self._shared_experts_attn_tp_comm: + return self.shared_experts(hidden_states) + + group = get_parallel().attn_tp_group + # SP-MoE presents one contiguous token shard per attention-TP rank; + # the DP local buffer is the full reassembled per-replica batch. + gathered_hidden_states = get_local_dp_buffer(group) + attn_tp_all_gather_into_tensor(gathered_hidden_states, hidden_states) + gathered_shared_output = self.shared_experts(gathered_hidden_states) + shared_output = torch.empty_like(hidden_states) + attn_tp_reduce_scatter_tensor(shared_output, gathered_shared_output) + return shared_output + def _forward_unfused( - self, hidden_states: torch.Tensor, *, prefix_sum: Optional[torch.Tensor] + self, + hidden_states: torch.Tensor, + *, + prefix_sum: Optional[torch.Tensor], ) -> torch.Tensor: """Front section with three separate GEMMs, each reading hidden_states: shared-expert MLP, router gate, latent down-proj.""" @@ -924,10 +1001,10 @@ class KimiK3MoE(nn.Module): if self._sbo_shared_overlap: self.alt_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(self.alt_stream): - shared_output = self.shared_experts(hidden_states) + shared_output = self._forward_shared_experts(hidden_states) shared_event = self.alt_stream.record_event() else: - shared_output = self.shared_experts(hidden_states) + shared_output = self._forward_shared_experts(hidden_states) # Front: gate + TopK (+ latent down-proj when the merged front covers it). # The gate and the latent down-proj read the same hidden_states, so the @@ -945,7 +1022,6 @@ class KimiK3MoE(nn.Module): # fp32 logits reach the radix router from moe_fused_gate. router_logits = self.gate(hidden_states) topk_output = self.topk(hidden_states, router_logits) - issue_shared() if not self.use_latent_moe: @@ -968,15 +1044,28 @@ class KimiK3MoE(nn.Module): ) if routed_input is None: - routed_input, _ = self.routed_expert_down_proj(hidden_states) + if hidden_states.shape[0] == 0: + # Idle DP ranks must still enter the EP dispatch below so the + # active replicas can exchange routed tokens. Ascend's + # quantized matmul does not accept an empty activation, so + # materialize its shape-only result without launching GEMM. + routed_input = hidden_states.new_empty((0, self.moe_hidden_size)) + else: + routed_input, _ = self.routed_expert_down_proj(hidden_states) expert_output = ( self._forward_mega_experts(routed_input, topk_output) if self._use_mega_moe else self.experts(routed_input, topk_output) ) - latent = self._reduce_latent(expert_output) - # up_proj is replicated, so the routed output is now fully reduced. - out, _ = self.routed_expert_up_proj(latent) + if expert_output.shape[0] == 0: + # The EP combine returns one row per source token. Keep the + # source-side empty result while avoiding empty RMSNorm/up-proj + # launches; the collective itself has already completed above. + out = hidden_states.new_empty((0, hidden_states.shape[1])) + else: + latent = self._reduce_latent(expert_output) + # up_proj is replicated, so the routed output is now fully reduced. + out, _ = self.routed_expert_up_proj(latent) if shared_event is not None: # SBO join: as late as possible, so the side-stream shared experts # get the whole routed a2a + latent tail to hide under. @@ -984,10 +1073,16 @@ class KimiK3MoE(nn.Module): if shared_output is not None: # tp1 shared experts (SP-MoE) are complete per-rank; TP-sharded # ones need the partial-sum reduction. - if self.tp_size > 1 and not self._shared_experts_tp1: + if ( + self.tp_size > 1 + and not self._shared_experts_tp1 + and not self._shared_experts_attn_tp_comm + ): shared_output = tensor_model_parallel_all_reduce(shared_output) - return _add3(out, shared_output, prefix_sum) - return out if prefix_sum is None else out + prefix_sum + out = _add3(out, shared_output, prefix_sum) + return out + out = out if prefix_sum is None else out + prefix_sum + return out @cached_property def _moe_front_needs_dense_bf16(self) -> bool: @@ -1581,6 +1676,8 @@ class KimiK3DeltaAttention(nn.Module): cuda graph capture).""" if not self.use_full_rank_gate: return + if _is_npu: + return self._bfa_w, sizes = _merge_weights_as_views( [self.f_a_proj, self.b_proj], pad_rows_to=8 ) @@ -1599,6 +1696,8 @@ class KimiK3DeltaAttention(nn.Module): return layer = self.attn w = layer.conv_weights + if _is_npu: + return seg = 12 * 128 # compiled for H = HV = 12 heads of 128 (TP8) if ( w is None @@ -1761,6 +1860,9 @@ class KimiK3MLAAttention(DeepseekV2AttentionMLA): ) -> None: self.all_reduce_fusion = all_reduce_fusion self.use_output_gate = getattr(config, "mla_use_output_gate", False) + # The fused Ascend split+RMSNorm path is not numerically equivalent for + # Kimi-K3. Other MLA models retain the existing fused fast path. + self._disable_npu_fused_split_qk_norm = True super().__init__( layer_id=layer_idx, hidden_size=config.hidden_size, @@ -1966,7 +2068,11 @@ class KimiK3DecoderLayer(nn.Module): # token shard. _a2a_backend = get_moe_a2a_backend() self._sp_moe = ( - (_a2a_backend.is_megamoe() or _a2a_backend.is_deepep()) + ( + _a2a_backend.is_megamoe() + or _a2a_backend.is_deepep() + or _a2a_backend.is_ascend_fuseep() + ) and self._is_moe_layer and get_parallel().attn_tp_group.world_size > 1 ) @@ -2243,7 +2349,6 @@ class KimiK3DecoderLayer(nn.Module): hidden_states, _, _ = self._finish_attn_reduce( hidden_states, allow_scatter=False ) - hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) hidden_states = self.mlp(hidden_states, forward_batch=forward_batch) return hidden_states, residual, False @@ -2963,6 +3068,8 @@ class KimiK3LinearForCausalLM(nn.Module): layer.self_attn, KimiK3DeltaAttention ): continue + if _is_npu: + continue from sglang.kernels.ops.attention.fla.kda import ( precompile_k3_recompute_w_u_kernel, ) @@ -3016,10 +3123,20 @@ class KimiK3ForConditionalGeneration(nn.Module): self.language_model = None if not config.encoder_only: + quant_description = getattr(quant_config, "quant_description", {}) + uses_wrapper_quant_prefix = any( + isinstance(name, str) and name.startswith("language_model.") + for name in quant_description + ) + language_prefix = ( + maybe_prefix(prefix, "language_model") + if uses_wrapper_quant_prefix + else prefix + ) self.language_model = KimiK3LinearForCausalLM( config.text_config, quant_config, - prefix="", + prefix=language_prefix, ) @property diff --git a/python/sglang/srt/speculative/dflash_utils.py b/python/sglang/srt/speculative/dflash_utils.py index 70c747fe6..264b30519 100644 --- a/python/sglang/srt/speculative/dflash_utils.py +++ b/python/sglang/srt/speculative/dflash_utils.py @@ -12,10 +12,13 @@ import triton import triton.language as tl from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod -from sglang.srt.layers.sampler import apply_custom_logit_processor +from sglang.srt.layers.sampler import ( + apply_custom_logit_processor, + top_p_normalize_probs_torch, +) from sglang.srt.managers.schedule_batch import Req from sglang.srt.speculative.spec_utils import sample_simulated_acc_len -from sglang.srt.utils import is_cuda, is_hip, is_musa +from sglang.srt.utils import is_cuda, is_hip, is_musa, is_npu DEFAULT_DFLASH_MASK_TOKEN = "<|MASK|>" @@ -67,6 +70,72 @@ def is_dflash_sampling_verify_available() -> bool: return _DFLASH_SAMPLING_VERIFY_AVAILABLE +def _dflash_npu_top_k_top_p_renorm_prob( + probs: torch.Tensor, + *, + top_ks: Optional[torch.Tensor] = None, + top_ps: Optional[torch.Tensor] = None, +) -> Optional[torch.Tensor]: + if not is_npu() or probs.device.type != "npu": + return None + try: + import torch_npu + except ImportError: + return None + if not hasattr(torch_npu, "npu_top_k_top_p"): + return None + + logits = probs.log() + npu_top_ps = ( + top_ps.reshape(-1).to(device=probs.device, dtype=probs.dtype) + if top_ps is not None + else None + ) + npu_top_ks = ( + top_ks.reshape(-1).to(device=probs.device, dtype=torch.int32) + if top_ks is not None + else None + ) + if npu_top_ks is not None and not bool( + torch.all((npu_top_ks >= 1) & (npu_top_ks <= 1024)).item() + ): + return None + filtered_logits = torch_npu.npu_top_k_top_p(logits, npu_top_ps, npu_top_ks) + return filtered_logits.softmax(dim=-1) + + +def _dflash_top_k_renorm_prob( + probs: torch.Tensor, top_ks: torch.Tensor +) -> torch.Tensor: + if top_k_renorm_prob is not None: + return top_k_renorm_prob(probs, top_ks) + + npu_probs = _dflash_npu_top_k_top_p_renorm_prob(probs, top_ks=top_ks) + if npu_probs is not None: + return npu_probs + + vocab_size = probs.shape[-1] + top_ks = top_ks.reshape(-1).to(device=probs.device, dtype=torch.int64) + top_ks = top_ks.clamp(min=1, max=vocab_size) + max_top_k = int(top_ks.max().item()) + topk_probs, topk_indices = torch.topk(probs, k=max_top_k, dim=-1) + ranks = torch.arange(max_top_k, device=probs.device)[None, :] + topk_probs.masked_fill_(ranks >= top_ks[:, None], 0.0) + topk_probs.div_(topk_probs.sum(dim=-1, keepdim=True)) + return torch.zeros_like(probs).scatter_(1, topk_indices, topk_probs) + + +def _dflash_top_p_renorm_prob( + probs: torch.Tensor, top_ps: torch.Tensor +) -> torch.Tensor: + if top_p_renorm_prob is not None: + return top_p_renorm_prob(probs, top_ps) + npu_probs = _dflash_npu_top_k_top_p_renorm_prob(probs, top_ps=top_ps) + if npu_probs is not None: + return npu_probs + return top_p_normalize_probs_torch(probs, top_ps) + + def dflash_draft_cell_size_per_token( *, draft_model_config: Any, @@ -879,7 +948,7 @@ def build_dflash_verify_target_probs( repeated_top_ps = torch.repeat_interleave( sampling_info.top_ps, draft_token_num, dim=0 ) - topk_probs = top_p_renorm_prob(topk_probs, repeated_top_ps) + topk_probs = _dflash_top_p_renorm_prob(topk_probs, repeated_top_ps) target_probs = torch.zeros_like(scaled_logits, dtype=topk_probs.dtype) target_probs.scatter_(1, topk_indices, topk_probs) @@ -888,12 +957,12 @@ def build_dflash_verify_target_probs( if not sparse_topk_applied: target_probs = F.softmax(scaled_logits, dim=-1) if need_top_k: - target_probs = top_k_renorm_prob( + target_probs = _dflash_top_k_renorm_prob( target_probs, torch.repeat_interleave(sampling_info.top_ks, draft_token_num, dim=0), ) if need_top_p: - target_probs = top_p_renorm_prob( + target_probs = _dflash_top_p_renorm_prob( target_probs, torch.repeat_interleave(sampling_info.top_ps, draft_token_num, dim=0), ) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index c7f27428e..ff5d3396d 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -212,7 +212,8 @@ class DraftBlockProposer: confidence_tap = None folded = False if ( - draft_sampler is not None + envs.SGLANG_DSPARK_FOLDED_PROPOSAL.get() + and draft_sampler is not None and fwd.can_run_graph and (all_greedy or draft_sampler.folded_sampling) ): @@ -322,7 +323,9 @@ class DraftBlockProposer: draft_positions = positions_2d[:, :gamma].reshape(-1) draft_cache_loc = verify_cache_loc_2d[:, :gamma].reshape(-1) - draft_owns_embed = hasattr(self.draft_model, "forward_embed") + draft_owns_embed = envs.SGLANG_DSPARK_EMBED_IN_GRAPH.get() and hasattr( + self.draft_model, "forward_embed" + ) draft_input_embeds: Optional[torch.Tensor] = None if not draft_owns_embed: noise_embedding = embed_module(draft_block_ids) @@ -337,6 +340,7 @@ class DraftBlockProposer: else: raise RuntimeError("DSpark decode expected batch.seq_lens_cpu, got None") + draft_num_tokens = bs * gamma draft_forward_batch = ForwardBatch( forward_mode=ForwardMode.TARGET_VERIFY, batch_size=bs, @@ -351,6 +355,10 @@ class DraftBlockProposer: spec_algorithm=SpeculativeAlgorithm.DSPARK, spec_info=self._draft_block_spec_info, capture_hidden_mode=CaptureHiddenMode.NULL, + num_token_non_padded=torch.tensor( + draft_num_tokens, dtype=torch.int32, device=device + ), + num_token_non_padded_cpu=draft_num_tokens, ) self._fill_dp_moe_sync_metadata(draft_forward_batch, batch) graph_runner = self.draft_model_runner.decode_cuda_graph_runner @@ -377,6 +385,9 @@ class DraftBlockProposer: def _fill_dp_moe_sync_metadata( self, forward_batch: ForwardBatch, batch: ScheduleBatch ) -> None: + # The dense DSpark draft still reuses the target batch's graph tier. + # Set graph eligibility before the DP-MoE-only metadata early return. + forward_batch.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph if not self._dp_moe_sync or batch.global_num_tokens is None: return gnt, gnt_logprob = spec_scale_global_num_tokens( @@ -400,4 +411,3 @@ class DraftBlockProposer: forward_batch.global_num_tokens_for_logprob_gpu = torch.tensor( gnt_logprob, dtype=torch.int64 ).to(device, non_blocking=True) - forward_batch.can_run_dp_cuda_graph = batch.can_run_dp_cuda_graph 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 3a8bb7151..339419741 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_worker_v2.py @@ -350,11 +350,18 @@ class DSparkWorkerV2(BaseSpecWorker): ) with self._draft_context(): if capture_decode_cuda_graph: - self._draft_sampler = self._maybe_build_draft_sampler() - if self._draft_sampler is not None: - self.draft_model_runner.capture_tail_hooks.append( - make_draft_sampler_capture_hook(self._draft_sampler) - ) + # Keep the draft model graph enabled when folded proposal is + # disabled, but do not capture the proposal head as a tail + # hook. The proposer will compute base logits and the Markov + # block eagerly from the graph's hidden states instead. Apart + # from being the intended precision fallback, skipping the + # unused hook avoids paying for two proposal computations. + if envs.SGLANG_DSPARK_FOLDED_PROPOSAL.get(): + self._draft_sampler = self._maybe_build_draft_sampler() + if self._draft_sampler is not None: + self.draft_model_runner.capture_tail_hooks.append( + make_draft_sampler_capture_hook(self._draft_sampler) + ) self._proposer.attach_draft_sampler(self._draft_sampler) self._draft_worker.init_cuda_graphs( capture_decode_cuda_graph=capture_decode_cuda_graph