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 8342e01f7..aee30fac7 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -8,6 +8,7 @@ import torch_npu from sglang.srt.configs.model_config import AttentionArch from sglang.srt.hardware_backend.npu.attention.mla_preprocess import ( + is_fia_nz, is_mla_preprocess_enabled, ) from sglang.srt.layers.attention.base_attn_backend import AttentionBackend @@ -27,6 +28,14 @@ import logging import numpy as np + +def _reshape_kv_for_fia_nz( + tensor: torch.Tensor, num_heads: int, head_dim: int, page_size: int +) -> torch.Tensor: + """Reshapes a tensor for FIA NZ format.""" + return tensor.view(-1, 1, num_heads * head_dim // 16, page_size, 16) + + logger = logging.getLogger(__name__) @@ -881,12 +890,20 @@ class AscendAttnBackend(AttentionBackend): ) c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) - k_rope_cache = k_rope.view( - -1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim - ) - c_kv_cache = c_kv.view( - -1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank - ) + if is_fia_nz(): + k_rope_cache = _reshape_kv_for_fia_nz( + k_rope, layer.tp_k_head_num, self.qk_rope_head_dim, self.page_size + ) + c_kv_cache = _reshape_kv_for_fia_nz( + c_kv, layer.tp_v_head_num, self.kv_lora_rank, self.page_size + ) + else: + k_rope_cache = k_rope.view( + -1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim + ) + c_kv_cache = c_kv.view( + -1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank + ) q_nope = q.view(-1, layer.tp_q_head_num, self.kv_lora_rank).contiguous() q_rope = q_rope.view(-1, layer.tp_q_head_num, self.qk_rope_head_dim) @@ -1082,15 +1099,24 @@ class AscendAttnBackend(AttentionBackend): ) else: c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id) - k_rope_cache = k_rope.view( - -1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim - ) - c_kv_cache = c_kv.view( - -1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank - ) + if is_fia_nz(): + k_rope_cache = _reshape_kv_for_fia_nz( + k_rope, layer.tp_k_head_num, self.qk_rope_head_dim, self.page_size + ) + c_kv_cache = _reshape_kv_for_fia_nz( + c_kv, layer.tp_v_head_num, self.kv_lora_rank, self.page_size + ) + else: + k_rope_cache = k_rope.view( + -1, self.page_size, layer.tp_k_head_num * self.qk_rope_head_dim + ) + c_kv_cache = c_kv.view( + -1, self.page_size, layer.tp_k_head_num * self.kv_lora_rank + ) + + q_nope = q.view(-1, 1, layer.tp_q_head_num, self.kv_lora_rank).contiguous() + q_rope = q_rope.view(-1, 1, layer.tp_q_head_num, self.qk_rope_head_dim) - q_nope = q.view(-1, layer.tp_q_head_num, 1, self.kv_lora_rank).contiguous() - q_rope = q_rope.view(-1, layer.tp_q_head_num, 1, self.qk_rope_head_dim) if self.forward_metadata.seq_lens_cpu_int is None: actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_list else: @@ -1108,7 +1134,7 @@ class AscendAttnBackend(AttentionBackend): num_key_value_heads=layer.tp_k_head_num, block_table=self.forward_metadata.block_tables, block_size=self.page_size, - input_layout="BNSD", + input_layout="BSND", scale=layer.scaling, actual_seq_lengths_kv=actual_seq_len_kv, antiquant_mode=0, @@ -1128,7 +1154,7 @@ class AscendAttnBackend(AttentionBackend): num_key_value_heads=layer.tp_k_head_num, block_table=self.forward_metadata.block_tables, block_size=self.page_size, - input_layout="BNSD", + input_layout="BSND", scale=layer.scaling, actual_seq_lengths_kv=actual_seq_len_kv, antiquant_mode=0, @@ -1243,12 +1269,20 @@ class AscendAttnBackend(AttentionBackend): if self.use_fia and (layer.tp_q_head_num // layer.tp_k_head_num) >= 8: """layer.tp_q_head_num // layer.tp_k_head_num < 8 will support in the later version of CANN""" - kv_c = kv_c.view( - -1, self.page_size, layer.tp_k_head_num * self.kv_lora_rank - ) - k_pe = k_pe.view( - -1, self.page_size, layer.tp_k_head_num * self.qk_rope_head_dim - ) + if is_fia_nz(): + kv_c = _reshape_kv_for_fia_nz( + kv_c, layer.tp_k_head_num, self.kv_lora_rank, self.page_size + ) + k_pe = _reshape_kv_for_fia_nz( + k_pe, layer.tp_k_head_num, self.qk_rope_head_dim, self.page_size + ) + else: + kv_c = kv_c.view( + -1, self.page_size, layer.tp_k_head_num * self.kv_lora_rank + ) + k_pe = k_pe.view( + -1, self.page_size, layer.tp_k_head_num * self.qk_rope_head_dim + ) q = q.view( forward_batch.batch_size, -1, layer.tp_q_head_num, self.kv_lora_rank ) diff --git a/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py b/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py index d1780133e..977ce0744 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py +++ b/python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py @@ -16,6 +16,16 @@ def is_mla_preprocess_enabled() -> bool: return get_bool_env_var("SGLANG_NPU_USE_MLAPO") +@lru_cache(maxsize=1) +def is_fia_nz() -> bool: + is_fia_nz_ = get_bool_env_var("SGLANG_USE_FIA_NZ") + if is_fia_nz_: + assert ( + is_mla_preprocess_enabled() + ), "SGLANG_USE_FIA_NZ must be enable with SGLANG_NPU_USE_MLAPO" + return is_fia_nz_ + + def round_up(val: int, align: int) -> int: if align == 0: return 0 @@ -84,7 +94,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module): self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim def preprocess_weights(self, hidden_states): - self.dummy = torch.empty( + self.dummy = torch.zeros( (hidden_states.shape[-1]), dtype=hidden_states.dtype, device=hidden_states.device, @@ -281,7 +291,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module): -1, 1, 1, self.kv_lora_rank + self.qk_rope_head_dim ) # (B*S,N,1,D) - cache_mode = "PA_BNSD" + cache_mode = "PA_NZ" if is_fia_nz() else "PA_BNSD" self.kvCache = self.kvCache.view( -1, forward_batch.attn_backend.page_size, @@ -328,9 +338,22 @@ class NPUFusedMLAPreprocess(torch.nn.Module): dtype=input_dtype, device=hidden_states.device, ) - + if is_fia_nz(): + kv_shape, kv_rope_shape = k_cache.shape, v_cache.shape + num_blocks, block_size, num_heads, _ = kv_shape + k_cache = k_cache.view( + num_blocks, num_heads * self.kv_lora_rank // 16, block_size, 16 + ) + v_cache = v_cache.view( + num_blocks, num_heads * self.qk_rope_head_dim // 16, block_size, 16 + ) # TODO: dummy inputs to be removed # https://github.com/sgl-project/sgl-kernel-npu/issues/78 + if hasattr(self.q_a_layernorm, "bias"): + q_a_layernorm_bias = self.q_a_layernorm.bias + else: + q_a_layernorm_bias = self.dummy + torch.ops.npu.mla_preprocess( hidden_states, self.dummy, @@ -338,7 +361,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module): self.qkv_a_proj_weight_nz, self.qkv_a_proj_deq_scale_kvq, self.q_a_layernorm.weight, - self.q_a_layernorm.bias, + q_a_layernorm_bias, self.q_b_proj_weight_nz, self.q_b_proj_deq_scale, self.kv_a_layernorm.weight, @@ -354,13 +377,18 @@ class NPUFusedMLAPreprocess(torch.nn.Module): quant_scale1=self.q_b_proj.input_scale, quant_offset1=self.q_b_proj_input_offset, bias1=self.q_b_proj_quant_bias, - cache_mode="krope_ctkv", + cache_mode="nzcache" if is_fia_nz() else "krope_ctkv", quant_mode="per_tensor_quant_asymm", q_out0=q_nope_out, kv_cache_out0=k_cache, q_out1=q_rope_out, kv_cache_out1=v_cache, ) + + if is_fia_nz(): + k_cache = k_cache.view(kv_shape) + v_cache = v_cache.view(kv_rope_shape) + return ( q_rope_out, v_cache, 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 58cd9194e..6626c468f 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 @@ -1,9 +1,11 @@ from typing import TYPE_CHECKING import torch +import torch_npu from sglang.srt.hardware_backend.npu.attention.mla_preprocess import ( NPUFusedMLAPreprocess, + is_fia_nz, is_mla_preprocess_enabled, ) from sglang.srt.layers.attention.nsa.utils import ( @@ -62,26 +64,39 @@ def forward_mha_prepare_npu( kv_a, _ = latent_cache.split([m.kv_lora_rank, m.qk_rope_head_dim], dim=-1) latent_cache = latent_cache.unsqueeze(1) - kv_a = m.kv_a_layernorm(kv_a) - kv = m.kv_b_proj(kv_a)[0] + B, S = q.shape[0], 1 + cos, sin = m.rotary_emb.get_cos_sin_cache( + positions, hidden_states.dtype, offsets=None + ) + q_pe = torch_npu.npu_interleave_rope( + q_pe.reshape(B, -1, S, m.qk_rope_head_dim), + cos, + sin, + ) + q_pe = q_pe.reshape(B, -1, m.qk_rope_head_dim) + + ckv_cache, k_rope_cache = forward_batch.token_to_kv_pool.get_kv_buffer(m.layer_id) + _, _, k_pe, kv_a = torch_npu.npu_kv_rmsnorm_rope_cache( + latent_cache.view(-1, 1, 1, m.kv_lora_rank + m.qk_rope_head_dim), # bnsd + m.kv_a_layernorm.weight, + cos, + sin, + forward_batch.out_cache_loc.to(torch.int64), + k_rope_cache, + ckv_cache, + k_rope_scale=None, + c_kv_scale=None, + k_rope_offset=None, + c_kv_offset=None, + epsilon=m.kv_a_layernorm.variance_epsilon, + cache_mode="PA_NZ" if is_fia_nz() else "PA_BNSD", + is_output_kv=True, + ) # adapter NZ + + k_pe = k_pe.reshape(B, -1, m.qk_rope_head_dim) - k_pe = latent_cache[:, :, m.kv_lora_rank :] - if m.rotary_emb is not None: - q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe) q[..., m.qk_nope_head_dim :] = q_pe - m._set_mla_kv_buffer(latent_cache, kv_a, k_pe, forward_batch) - if forward_batch.mha_one_shot and sum(forward_batch.extend_prefix_lens_cpu) != 0: - if m.use_nsa and m.kv_cache_dtype == "fp8_e4m3": - # FP8 path: dequantize NSA-specific FP8 format to BF16 - kv_a, k_pe = m._get_mla_kv_buffer_from_fp8(forward_batch) - else: - # BF16/FP16 path: directly fetch from cache - kv_a, k_pe = m._get_mla_kv_buffer( - forward_batch.fetch_mha_one_shot_kv_indices(), - q.dtype, - forward_batch, - ) kv = m.kv_b_proj(kv_a)[0] kv = kv.view(-1, m.num_local_heads, m.qk_nope_head_dim + m.v_head_dim) k_nope = kv[..., : m.qk_nope_head_dim] @@ -233,15 +248,15 @@ 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), + (attn_output.shape[0], m.num_local_heads, m.v_head_dim), dtype=attn_output.dtype, device=attn_output.device, ) - torch.bmm( - attn_output.transpose(0, 1), - m.w_vc, - out=attn_bmm_output.view(-1, m.num_local_heads, m.v_head_dim).transpose(0, 1), - ) + + attn_output = attn_output.contiguous() + torch.ops.npu.batch_matmul_transpose(attn_output, m.w_vc, attn_bmm_output) + + attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim) output, _ = m.o_proj(attn_bmm_output) return output diff --git a/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py b/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py index 279700b3e..4c5cbea48 100644 --- a/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py +++ b/python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py @@ -324,7 +324,7 @@ class NPUW8A8Int8DynamicMoEMethod(FusedMoEMethodBase): class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase): def __init__(self) -> None: - self.group_size = 256 + self.group_size = 0 self.tp_size = 1 def create_weights( @@ -338,6 +338,7 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase): ) -> None: from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported + self.is_per_channel_weight = self.group_size == 0 self.num_experts = num_experts extra_weight_attrs.update( {"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value} @@ -399,52 +400,55 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase): set_weight_attrs(w2_weight_offset, extra_weight_attrs) # >>> special param for w4a8 - w13_weight_scale_second = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size // self.group_size, - dtype=torch.float32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_scale_second", w13_weight_scale_second) - set_weight_attrs(w13_weight_scale_second, extra_weight_attrs) - w13_weight_offset_second = torch.nn.Parameter( - torch.empty( - num_experts, - 2 * intermediate_size_per_partition, - hidden_size // self.group_size, - dtype=torch.float32, - ), - requires_grad=False, - ) - layer.register_parameter("w13_weight_offset_second", w13_weight_offset_second) - set_weight_attrs(w13_weight_offset_second, extra_weight_attrs) + if not self.is_per_channel_weight: + w13_weight_scale_second = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.group_size, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w13_weight_scale_second", w13_weight_scale_second) + set_weight_attrs(w13_weight_scale_second, extra_weight_attrs) + w13_weight_offset_second = torch.nn.Parameter( + torch.empty( + num_experts, + 2 * intermediate_size_per_partition, + hidden_size // self.group_size, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter( + "w13_weight_offset_second", w13_weight_offset_second + ) + set_weight_attrs(w13_weight_offset_second, extra_weight_attrs) - w2_weight_scale_second = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition // self.group_size, - dtype=torch.float32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_scale_second", w2_weight_scale_second) - set_weight_attrs(w2_weight_scale_second, extra_weight_attrs) + w2_weight_scale_second = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // self.group_size, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_scale_second", w2_weight_scale_second) + set_weight_attrs(w2_weight_scale_second, extra_weight_attrs) - w2_weight_offset_second = torch.nn.Parameter( - torch.empty( - num_experts, - hidden_size, - intermediate_size_per_partition // self.group_size, - dtype=torch.float32, - ), - requires_grad=False, - ) - layer.register_parameter("w2_weight_offset_second", w2_weight_offset_second) - set_weight_attrs(w2_weight_offset_second, extra_weight_attrs) + w2_weight_offset_second = torch.nn.Parameter( + torch.empty( + num_experts, + hidden_size, + intermediate_size_per_partition // self.group_size, + dtype=torch.float32, + ), + requires_grad=False, + ) + layer.register_parameter("w2_weight_offset_second", w2_weight_offset_second) + set_weight_attrs(w2_weight_offset_second, extra_weight_attrs) w13_scale_bias = torch.nn.Parameter( torch.empty( @@ -466,6 +470,11 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase): def process_scale(self, weight: torch.Tensor, scale, per_group_scale): scale = scale.transpose(1, 2).contiguous() + if self.is_per_channel_weight: + scale_np = scale.cpu().numpy() + scale_np.dtype = np.uint32 + scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu() + return scale_uint64_tensor, None per_group_scale = per_group_scale.transpose(1, 2).contiguous() group_num, k, n = weight.shape # the weight of the new version is reduced by half by pack n, so it needs to be restored diff --git a/python/sglang/srt/layers/rotary_embedding.py b/python/sglang/srt/layers/rotary_embedding.py index e5b129f24..0ee89d102 100644 --- a/python/sglang/srt/layers/rotary_embedding.py +++ b/python/sglang/srt/layers/rotary_embedding.py @@ -806,6 +806,10 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding): / yarn_get_mscale(self.scaling_factor, float(mscale_all_dim)) * attn_factor ) + self.cos_cached_total = None + self.sin_cached_total = None + self.cos_cached = None + self.sin_cached = None self.device = device super().__init__( head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype @@ -854,8 +858,41 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding): cos = freqs.cos() * self.mscale sin = freqs.sin() * self.mscale cache = torch.cat((cos, sin), dim=-1) + + emb = torch.cat((freqs, freqs), dim=-1) + self.cos_cached_total = torch.cos(emb) * self.mscale + self.sin_cached_total = torch.sin(emb) * self.mscale return cache + def get_cos_cached_total(self): + return self.cos_cached_total + + def get_sin_cached_total(self): + return self.sin_cached_total + + def get_cos_sin_cache( + self, positions, dtype, offsets: Optional[torch.Tensor] = None + ): + self.cos_cached = ( + self.cos_cached_total[ + torch.add(positions, offsets) if offsets is not None else positions + ] + .unsqueeze(-2) + .unsqueeze(-2) + .to(dtype) + ) + self.sin_cached = ( + self.sin_cached_total[ + torch.add(positions, offsets) if offsets is not None else positions + ] + .unsqueeze(-2) + .unsqueeze(-2) + .to(dtype) + ) + cos = self.cos_cached.to(positions.device) + sin = self.sin_cached.to(positions.device) + return cos, sin + def forward_native( self, positions: torch.Tensor, @@ -906,14 +943,7 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding): num_tokens, num_q_heads, _ = query.shape num_k_heads = key.shape[1] - self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(positions.device) - cos_sin = self.cos_sin_cache[ - torch.add(positions, offsets) if offsets is not None else positions - ] - cos, sin = cos_sin.chunk(2, dim=-1) - # Reshape to [batchsize, head_dim, seq, rotary_dim] - cos = cos.repeat(1, 2).unsqueeze(-2).unsqueeze(-2) - sin = sin.repeat(1, 2).unsqueeze(-2).unsqueeze(-2) + cos, sin = self.get_cos_sin_cache(positions, query.dtype, offsets) query_rot = query[..., : self.rotary_dim] key_rot = key[..., : self.rotary_dim] diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 931445cc5..5fb402f66 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -541,11 +541,6 @@ class ForwardBatch: num_tokens_per_dp=num_tokens_per_dp, ) - self.num_token_non_padded_cpu = compute_local_num_token_non_padded( - global_num_token_non_padded=self.num_token_non_padded_cpu, - num_tokens_per_dp=num_tokens_per_dp, - ) - def merge_mm_inputs(self) -> Optional[MultimodalInputs]: """ Merge all multimodal inputs in the batch into a single MultiModalInputs object.