diff --git a/.github/workflows/nightly-test-npu.yml b/.github/workflows/nightly-test-npu.yml index 6705b0794..bfbbe8041 100644 --- a/.github/workflows/nightly-test-npu.yml +++ b/.github/workflows/nightly-test-npu.yml @@ -55,11 +55,12 @@ jobs: # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy - curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl + curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Print Log Information run: | bash scripts/ci/npu/npu_log_print.sh + - name: Run test timeout-minutes: 240 env: @@ -102,7 +103,7 @@ jobs: # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy - curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl + curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Print Log Information run: | @@ -148,7 +149,7 @@ jobs: # copy required file from our daily cache cp ~/.cache/modelscope/hub/datasets/otavia/ShareGPT_Vicuna_unfiltered/ShareGPT_V3_unfiltered_cleaned_split.json /tmp # copy download through proxy - curl -o /tmp/test.jsonl -L https://gh-proxy.test.osinfra.cn/https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl + curl -o /tmp/test.jsonl -L https://raw.githubusercontent.com/openai/grade-school-math/master/grade_school_math/data/test.jsonl - name: Print Log Information run: | 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 0394a3de3..3a5644177 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -831,6 +831,7 @@ class AscendAttnBackend(AttentionBackend): layer.qk_head_dim <= 128 and causal and forward_batch.encoder_lens is None + and layer.logit_cap == 0 and not getattr(self, "use_native_sdpa", False) ): if not self.use_alibi: @@ -896,6 +897,8 @@ class AscendAttnBackend(AttentionBackend): scaling=layer.scaling, enable_gqa=use_gqa, causal=causal, + logit_cap=layer.logit_cap, + logit_capping_method=layer.logit_capping_method, ) attn_output = attn_output.view( -1, layer.tp_q_head_num * layer.v_head_dim @@ -1022,7 +1025,7 @@ class AscendAttnBackend(AttentionBackend): layer.layer_id ) kv_cache = torch.cat([k_cache, v_cache], dim=-1) - attn_output = self.native_attn._run_sdpa_forward_extend( + attn_output = self.native_attn.run_sdpa_forward_extend( q, attn_output, kv_cache.view(-1, layer.tp_k_head_num, layer.qk_head_dim), @@ -1525,7 +1528,7 @@ class AscendAttnBackend(AttentionBackend): ) # there are some accuracy issues in cross attention scene to use torch_npu._npu_flash_attention_qlens # forward_batch.encoder_lens is not None in cross attention scend, we add native attn to solve accuracy issues - elif forward_batch.encoder_lens is None: + elif forward_batch.encoder_lens is None and layer.logit_cap == 0: query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim) num_tokens = query.shape[0] if not self.use_alibi: @@ -1585,6 +1588,8 @@ class AscendAttnBackend(AttentionBackend): scaling=layer.scaling, enable_gqa=use_gqa, causal=False, + logit_cap=layer.logit_cap, + logit_capping_method=layer.logit_capping_method, ) return attn_output.view(num_tokens, layer.tp_q_head_num * layer.v_head_dim) else: diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py index 4c3b43d69..34bbfc67f 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_torch_native_backend.py @@ -1,5 +1,7 @@ from __future__ import annotations +import math + import torch from torch.nn.functional import scaled_dot_product_attention @@ -8,6 +10,49 @@ class AscendTorchNativeAttnBackend: def __init__(self): pass + def scaled_dot_product_attention_with_softcapping( + self, + query, + key, + value, + attn_mask=None, + is_causal=False, + scale=None, + enable_gqa=False, + logit_cap=0.0, + logit_capping_method="tanh", + ) -> torch.Tensor: + L, S = query.size(-2), key.size(-2) + scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale + attn_bias = torch.zeros(L, S, dtype=query.dtype, device=query.device) + if is_causal: + assert attn_mask is None + temp_mask = torch.ones(L, S, dtype=torch.bool, device=query.device).tril( + diagonal=0 + ) + attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf")) + attn_bias.to(query.dtype) + + if attn_mask is not None: + if attn_mask.dtype == torch.bool: + attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf")) + else: + attn_bias = attn_mask + attn_bias + + if enable_gqa: + key = key.repeat_interleave(query.size(-3) // key.size(-3), -3) + value = value.repeat_interleave(query.size(-3) // value.size(-3), -3) + + attn_weight = query @ key.transpose(-2, -1) * scale_factor + + if logit_cap > 0: + if logit_capping_method == "tanh": + attn_weight = logit_cap * torch.tanh(attn_weight / logit_cap) + + attn_weight += attn_bias + attn_weight = torch.softmax(attn_weight, dim=-1) + return attn_weight @ value + def run_sdpa_forward_extend( self, query: torch.Tensor, @@ -24,6 +69,8 @@ class AscendTorchNativeAttnBackend: scaling=None, enable_gqa=False, causal=False, + logit_cap: float = 0.0, + logit_capping_method: str = "tanh", ): """Run the extend forward by using torch native sdpa op. @@ -94,18 +141,34 @@ class AscendTorchNativeAttnBackend: per_req_key = per_req_key.to(per_req_query.dtype) per_req_value = per_req_value.to(per_req_query.dtype) - per_req_out_redudant = ( - scaled_dot_product_attention( - per_req_query_redudant.unsqueeze(0), - per_req_key.unsqueeze(0), - per_req_value.unsqueeze(0), - enable_gqa=enable_gqa, - scale=scaling, - is_causal=causal, + if logit_cap > 0: + per_req_out_redudant = ( + self.scaled_dot_product_attention_with_softcapping( + per_req_query_redudant.unsqueeze(0), + per_req_key.unsqueeze(0), + per_req_value.unsqueeze(0), + enable_gqa=enable_gqa, + scale=scaling, + is_causal=causal, + logit_cap=logit_cap, + logit_capping_method=logit_capping_method, + ) + .squeeze(0) + .movedim(query.dim() - 2, 0) + ) + else: + per_req_out_redudant = ( + scaled_dot_product_attention( + per_req_query_redudant.unsqueeze(0), + per_req_key.unsqueeze(0), + per_req_value.unsqueeze(0), + enable_gqa=enable_gqa, + scale=scaling, + is_causal=causal, + ) + .squeeze(0) + .movedim(query.dim() - 2, 0) ) - .squeeze(0) - .movedim(query.dim() - 2, 0) - ) output[start_q:end_q, :, :] = per_req_out_redudant[prefill_seq_len_q:, :, :] start_q, start_kv = end_q, end_kv return output @@ -124,6 +187,8 @@ class AscendTorchNativeAttnBackend: scaling=None, enable_gqa=False, causal=False, + logit_cap: float = 0.0, + logit_capping_method: str = "tanh", ): """Run the decode forward by using torch native sdpa op. @@ -180,18 +245,34 @@ class AscendTorchNativeAttnBackend: per_req_key = per_req_key.to(per_req_query.dtype) per_req_value = per_req_value.to(per_req_query.dtype) - per_req_out = ( - scaled_dot_product_attention( - per_req_query.unsqueeze(0), - per_req_key.unsqueeze(0), - per_req_value.unsqueeze(0), - enable_gqa=enable_gqa, - scale=scaling, - is_causal=causal, + if logit_cap > 0: + per_req_out = ( + self.scaled_dot_product_attention_with_softcapping( + per_req_query.unsqueeze(0), + per_req_key.unsqueeze(0), + per_req_value.unsqueeze(0), + enable_gqa=enable_gqa, + scale=scaling, + is_causal=causal, + logit_cap=logit_cap, + logit_capping_method=logit_capping_method, + ) + .squeeze(0) + .movedim(query.dim() - 2, 0) + ) + else: + per_req_out = ( + scaled_dot_product_attention( + per_req_query.unsqueeze(0), + per_req_key.unsqueeze(0), + per_req_value.unsqueeze(0), + enable_gqa=enable_gqa, + scale=scaling, + is_causal=causal, + ) + .squeeze(0) + .movedim(query.dim() - 2, 0) ) - .squeeze(0) - .movedim(query.dim() - 2, 0) - ) output[start_q:end_q, :, :] = per_req_out start_q, start_kv = end_q, end_kv diff --git a/python/sglang/srt/models/grok.py b/python/sglang/srt/models/grok.py index 85155719a..f82642777 100644 --- a/python/sglang/srt/models/grok.py +++ b/python/sglang/srt/models/grok.py @@ -18,6 +18,7 @@ import math from typing import Iterable, Optional, Tuple import torch +import torch.nn.functional as F from torch import nn from transformers import PretrainedConfig @@ -59,7 +60,9 @@ from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode from sglang.srt.model_executor.forward_batch_info import ForwardBatch from sglang.srt.model_loader.loader import DefaultModelLoader from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.utils import add_prefix +from sglang.srt.utils import add_prefix, is_npu + +_is_npu = is_npu() logger = logging.getLogger(__name__) @@ -143,7 +146,7 @@ class Grok1MoE(nn.Module): top_k=top_k, renormalize=False, layer_id=layer_id, - custom_routing_function=custom_routing_function, + custom_routing_function=None if _is_npu else custom_routing_function, ) self.experts = FusedMoE( @@ -162,8 +165,21 @@ class Grok1MoE(nn.Module): ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - topk_output = self.topk(hidden_states, self.gate.weight) - return self.experts(hidden_states, topk_output) + if not _is_npu: + topk_output = self.topk(hidden_states, self.gate.weight) + return self.experts(hidden_states, topk_output) + else: + orig_shape = hidden_states.shape + hidden_states = hidden_states.view(-1, self.hidden_size) + + router_logits, _ = self.gate(hidden_states) + router_logits = self.router_logit_softcapping * F.tanh( + router_logits / self.router_logit_softcapping + ) + topk_output = self.topk(hidden_states, router_logits) + + final_hidden_states = self.experts(hidden_states, topk_output) + return final_hidden_states.view(orig_shape) def _yarn_linear_ramp_mask( @@ -228,6 +244,8 @@ class ScalingRotaryEmbedding(RotaryEmbedding): self.attn_factor = attn_factor self.beta_fast = beta_fast self.beta_slow = beta_slow + if _is_npu: + dtype = torch.float32 # Get n-d magnitude scaling corrected for interpolation self.mscale = float(_yarn_get_mscale(self.scaling_factor) * attn_factor) super().__init__( @@ -396,6 +414,7 @@ class Grok1Attention(nn.Module): max_position=max_position, base=int(self.rope_theta), is_neox_style=True, + dtype=torch.float32 if _is_npu else None, ) pos_encoding_mode = "NONE" @@ -425,7 +444,12 @@ class Grok1Attention(nn.Module): qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.rotary_emb(positions, q, k) + if not _is_npu: + q, k = self.rotary_emb(positions, q, k) + else: + odtype = q.dtype + q, k = self.rotary_emb(positions, q.to(torch.float32), k.to(torch.float32)) + q, k = q.to(odtype), k.to(odtype) attn_output = self.attn(q, k, v, forward_batch) diff --git a/test/registered/ascend/llm_models/test_ascend_grok_2.py b/test/registered/ascend/llm_models/test_ascend_grok_2.py new file mode 100644 index 000000000..9923131f9 --- /dev/null +++ b/test/registered/ascend/llm_models/test_ascend_grok_2.py @@ -0,0 +1,29 @@ +import unittest + +from sglang.test.ascend.gsm8k_ascend_mixin import GSM8KAscendMixin +from sglang.test.ci.ci_register import register_npu_ci +from sglang.test.test_utils import CustomTestCase + +register_npu_ci(est_time=400, suite="nightly-16-npu-a3", nightly=True) + + +class TestGrok2(GSM8KAscendMixin, CustomTestCase): + model = "/root/.cache/modelscope/hub/models/huihui-ai/grok-2" + accuracy = 0.91 + other_args = [ + "--trust-remote-code", + "--mem-fraction-static", + "0.8", + "--attention-backend", + "ascend", + "--disable-radix-cache", + "--disable-cuda-graph", + "--tokenizer-path", + "/root/.cache/modelscope/hub/models/huihui-ai/grok-2/tokenizer.tok.json", + "--tp-size", + "16", + ] + + +if __name__ == "__main__": + unittest.main()