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 3a5644177..ccbeeb453 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -11,6 +11,7 @@ from sgl_kernel_npu.attention.sinks_attention import ( ) from sglang.srt.configs.model_config import AttentionArch +from sglang.srt.dllm.config import DllmConfig from sglang.srt.hardware_backend.npu.attention.ascend_torch_native_backend import ( AscendTorchNativeAttnBackend, ) @@ -252,6 +253,13 @@ class AscendAttnBackend(AttentionBackend): if self.use_mla: self.ringmla_mask = self.ascend_attn_mask_builder.ringmla_mask + # dllm model config + self.dllm_config = DllmConfig.from_server_args(model_runner.server_args) + self.is_dllm_model = False + if self.dllm_config is not None: + self.is_dllm_model = True + self.dllm_block_size = self.dllm_config.block_size + def get_verify_buffers_to_fill_after_draft(self): """ Return buffers for verify attention kernels that needs to be filled after draft. @@ -353,6 +361,20 @@ class AscendAttnBackend(AttentionBackend): metadata = ForwardMetadata() metadata.block_tables = self.graph_metadata["block_tables"][:bs, :] + if self.is_dllm_model: + max_len = int(seq_lens[:bs].max().item()) + max_seq_pages = (max_len + self.page_size - 1) // self.page_size + metadata.block_tables[:bs, :max_seq_pages].copy_( + ( + self.req_to_token[req_pool_indices[:bs], :max_len][ + :, :: self.page_size + ] + // self.page_size + ).to(torch.int32) + ) + metadata.block_tables[:bs, max_seq_pages:].fill_(0) + metadata.block_tables[bs:, :].fill_(0) + metadata.seq_lens_cpu_list = seq_lens.cpu().int().tolist() metadata.seq_lens = seq_lens if ( @@ -374,6 +396,15 @@ class AscendAttnBackend(AttentionBackend): dtype=torch.int32, device=seq_lens.device, ) + if forward_mode.is_dllm_extend(): + extend_seq_lens_cpu_int = torch.tensor( + [self.dllm_block_size for i in range(bs)], + dtype=torch.int32, + device=seq_lens.device, + ) + metadata.seq_lens_list_cumsum = ( + torch.cumsum(extend_seq_lens_cpu_int, dim=0).int().tolist() + ) self.graph_metadata[bs] = metadata self.forward_metadata = metadata @@ -401,8 +432,10 @@ class AscendAttnBackend(AttentionBackend): self.req_to_token[req_pool_indices[:bs], :max_len][:, :: self.page_size] // self.page_size ) + metadata.block_tables[:bs, max_seq_pages:].fill_(0) metadata.block_tables[bs:, :].fill_(0) + if forward_mode.is_target_verify(): seq_lens = seq_lens + self.speculative_num_draft_tokens metadata.seq_lens[:bs].copy_(seq_lens[:bs]) @@ -729,6 +762,17 @@ class AscendAttnBackend(AttentionBackend): if is_mla_preprocess_enabled(): # MLAPO and MLAPROLOG do save kv_cache save_kv_cache = False + if self.is_dllm_model: + return self.forward_dllm( + q, + k, + v, + layer, + forward_batch, + save_kv_cache, + q_rope=q_rope, + k_rope=k_rope, + ) if topk_indices is not None: return self.forward_sparse( q, @@ -823,7 +867,6 @@ class AscendAttnBackend(AttentionBackend): or layer.attn_type == AttentionType.ENCODER_ONLY ): causal = False - # 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 # Model skywork-reward-gemma2-2-27B also suffers from precision anomalies, thus the torch native backend becomes beneficial approach. @@ -841,7 +884,6 @@ class AscendAttnBackend(AttentionBackend): dtype=query.dtype, device=query.device, ) - torch_npu._npu_flash_attention_qlens( query=query, key_cache=k_cache, @@ -1085,6 +1127,65 @@ class AscendAttnBackend(AttentionBackend): return attn_output + def forward_dllm( + self, + q, + k, + v, + layer: RadixAttention, + forward_batch: ForwardBatch, + save_kv_cache: bool = True, + # For multi_head latent attention + q_rope: Optional[torch.Tensor] = None, + k_rope: Optional[torch.Tensor] = None, + topk_indices: Optional[torch.Tensor] = None, + ): + if save_kv_cache: + forward_batch.token_to_kv_pool.set_kv_buffer( + layer, forward_batch.out_cache_loc, k, v + ) + + k_cache = forward_batch.token_to_kv_pool.get_key_buffer(layer.layer_id) + v_cache = forward_batch.token_to_kv_pool.get_value_buffer(layer.layer_id) + query = q.reshape(-1, layer.tp_q_head_num, layer.qk_head_dim) + + if self.forward_metadata.seq_lens_cpu_int is None: + # capture + actual_seq_lengths_kv = self.forward_metadata.seq_lens_cpu_list + else: + # eagle + actual_seq_lengths_kv = ( + self.forward_metadata.seq_lens_cpu_int.cpu().int().tolist() + ) + + if self.forward_metadata.extend_seq_lens_cpu_int is None: + # capture & replay + actual_seq_lengths = self.forward_metadata.seq_lens_list_cumsum + else: + actual_seq_lengths = ( + torch.cumsum(self.forward_metadata.extend_seq_lens_cpu_int, dim=0) + .int() + .tolist() + ) + + attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score( + query, + k_cache.view(-1, self.page_size, layer.tp_k_head_num * layer.qk_head_dim), + v_cache.view(-1, self.page_size, layer.tp_v_head_num * layer.v_head_dim), + block_table=self.forward_metadata.block_tables, + block_size=self.page_size, + num_heads=layer.tp_q_head_num, + num_key_value_heads=layer.tp_k_head_num, + input_layout="TND", + atten_mask=None, + scale=layer.scaling, + actual_seq_lengths=actual_seq_lengths, + actual_seq_lengths_kv=actual_seq_lengths_kv, + ) + attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim) + + return attn_output + def forward_mtp( self, q, diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py index 896667d33..dc40a0fb8 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py @@ -83,10 +83,16 @@ class NPUGraphRunner(CudaGraphRunner): self.use_fia = get_bool_env_var("ASCEND_USE_FIA", "False") def _init_arch_map(self): - self.attr_name: Dict[str, str] = { - AttentionArch.MLA: "actual_seq_lengths_kv", - AttentionArch.MHA: "context_lens", - } + if self.is_dllm: + self.attr_name: Dict[str, str] = { + AttentionArch.MLA: "actual_seq_lengths_kv", + AttentionArch.MHA: "actual_seq_lengths_kv", + } + else: + self.attr_name: Dict[str, str] = { + AttentionArch.MLA: "actual_seq_lengths_kv", + AttentionArch.MHA: "context_lens", + } self.attr_type: Dict[str, Union[list, torch.Tensor]] = { AttentionArch.MLA: [], AttentionArch.MHA: torch.Tensor(), @@ -188,8 +194,15 @@ class NPUGraphRunner(CudaGraphRunner): output = self.output_buffers[self.bs] if isinstance(output, LogitsProcessorOutput): + if self.is_dllm: + next_token_logits = None + full_logits = output.full_logits[: self.raw_num_token] + else: + full_logits = None + next_token_logits = output.next_token_logits[: self.raw_num_token] return LogitsProcessorOutput( - next_token_logits=output.next_token_logits[: self.raw_num_token], + next_token_logits=next_token_logits, + full_logits=full_logits, hidden_states=( output.hidden_states[: self.raw_num_token] if output.hidden_states is not None diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 77b65379e..c4cfcb2c2 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -209,6 +209,7 @@ from sglang.srt.utils import ( set_random_seed, suppress_other_loggers, ) +from sglang.srt.utils.common import is_npu from sglang.srt.utils.hf_transformers_utils import ( get_processor, get_tokenizer, @@ -224,6 +225,8 @@ TEST_RETRACT = envs.SGLANG_TEST_RETRACT.get() TEST_RETRACT_INTERVAL = envs.SGLANG_TEST_RETRACT_INTERVAL.get() TEST_RETRACT_NO_PREFILL_BS = envs.SGLANG_TEST_RETRACT_NO_PREFILL_BS.get() +_is_npu = is_npu() + @dataclass class EmbeddingBatchResult: @@ -410,6 +413,24 @@ class Scheduler( def init_model_config(self): self.model_config = ModelConfig.from_server_args(self.server_args) + if _is_npu: + # make sure the page size is not larger than block_size and chunked_prefill_size on NPU backend + # the npu backend request the defined page size to be no larger than block_size and chunked_prefill_size + from sglang.srt.dllm.config import DllmConfig + + self.dllm_config = ( # For diffusion LLM + DllmConfig.from_server_args(self.server_args) + if self.server_args.dllm_algorithm is not None + else None + ) + if self.dllm_config: + if self.dllm_config.block_size < self.page_size: + logger.warning( + "WARNING: " + f"The page size {self.page_size} should not be larger than dllm block size {self.dllm_config.block_size}." + f"Page size now falls back to {self.dllm_config.block_size}" + ) + self.page_size = self.dllm_config.block_size def init_ipc_channels(self, port_args: PortArgs): context = zmq.Context(2) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index aaeb4fa7f..cc283a5d2 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -470,7 +470,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): if batch.dllm_config is not None: block_size = batch.dllm_config.block_size # Use int64 for AMD rotary embedding kernel compatibility - positions_dtype = torch.int64 if is_hip() else torch.int32 + positions_dtype = torch.int64 if is_hip() or _is_npu else torch.int32 ret.positions = torch.tensor( [ i diff --git a/python/sglang/srt/models/llada2.py b/python/sglang/srt/models/llada2.py index 5f5cc438c..7294524e3 100644 --- a/python/sglang/srt/models/llada2.py +++ b/python/sglang/srt/models/llada2.py @@ -77,11 +77,18 @@ from sglang.srt.models.utils import ( enable_fused_set_kv_buffer, ) from sglang.srt.server_args import get_global_server_args -from sglang.srt.utils import add_prefix, is_cuda, is_non_idle_and_non_empty, make_layers +from sglang.srt.utils import ( + add_prefix, + is_cuda, + is_non_idle_and_non_empty, + is_npu, + make_layers, +) LoraConfig = None logger = logging.getLogger(__name__) _is_cuda = is_cuda() +_is_npu = is_npu() class LLaDA2MoeMLP(nn.Module): @@ -190,6 +197,11 @@ class LLaDA2MoeSparseMoeBlock(nn.Module): self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) self.score_function = getattr(config, "score_function", None) + # fused_topk_npu() conducting norm before scale with routed_scaling_factor by default + # norm_topk_prob=True will renorm the routed_scaling_factor thus need to keep norm_topk_prob=False + if _is_npu: + self.norm_topk_prob = False + if config.hidden_act != "silu": raise ValueError( f"Unsupported activation: {config.hidden_act}. " diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index f2dfe68d1..4fe003cc6 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3124,6 +3124,12 @@ class ServerArgs: "Attention backend is set to triton for diffusion LLM inference on AMD GPUs" ) self.attention_backend = "triton" + elif is_npu(): + if self.attention_backend != "ascend": + logger.warning( + "Attention backend is overridden to 'ascend' when running on NPU for diffusion LLM inference." + ) + self.attention_backend = "ascend" elif not self.disable_cuda_graph: if self.attention_backend != "flashinfer": logger.warning( diff --git a/test/srt/ascend/test_llada2_mini_ascend.py b/test/srt/ascend/test_llada2_mini_ascend.py new file mode 100644 index 000000000..a1fdafcc5 --- /dev/null +++ b/test/srt/ascend/test_llada2_mini_ascend.py @@ -0,0 +1,87 @@ +import os +import unittest +from types import SimpleNamespace + +from sglang.srt.utils import kill_process_tree +from sglang.test.few_shot_gsm8k import run_eval as run_eval_few_shot_gsm8k +from sglang.test.send_one import BenchArgs, send_one_prompt +from sglang.test.test_utils import ( + DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + DEFAULT_URL_FOR_TEST, + CustomTestCase, + is_in_ci, + popen_launch_server, + write_github_step_summary, +) + + +class TestLLaDA2Mini(CustomTestCase): + @classmethod + def setUpClass(cls): + cls._old_disable_acl = os.environ.get("SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT") + os.environ["SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT"] = "1" + + cls.model = "/root/.cache/modelscope/hub/models/inclusionAI/LLaDA2.0-mini" + cls.base_url = DEFAULT_URL_FOR_TEST + + other_args = [ + "--trust-remote-code", + "--disable-radix-cache", + "--mem-fraction-static", + "0.9", + "--max-running-requests", + "1", + "--attention-backend", + "ascend", + "--dllm-algorithm", + "LowConfidence", # TODO: Add dLLM configurations + ] + + cls.process = popen_launch_server( + cls.model, + cls.base_url, + timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH, + other_args=other_args, + ) + + @classmethod + def tearDownClass(cls): + kill_process_tree(cls.process.pid) + + if cls._old_disable_acl is None: + os.environ.pop("SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT", None) + else: + os.environ["SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT"] = cls._old_disable_acl + + def test_gsm8k(self): + args = SimpleNamespace( + num_shots=5, + data_path=None, + num_questions=200, + max_new_tokens=512, + parallel=128, + host="http://127.0.0.1", + port=int(self.base_url.split(":")[-1]), + ) + metrics = run_eval_few_shot_gsm8k(args) + print(f"{metrics=}") + + self.assertGreater(metrics["accuracy"], 0.88) + self.assertGreater(metrics["output_throughput"], 70) + + def test_bs_1_speed(self): + args = BenchArgs(port=int(self.base_url.split(":")[-1]), max_new_tokens=2048) + acc_length, speed = send_one_prompt(args) + + print(f"{speed=:.2f}") + + if is_in_ci(): + write_github_step_summary( + f"### test_bs_1_speed (llada2-mini) with tp1\n" + f"{speed=:.2f} token/s\n" + ) + self.assertGreater(speed, 130) + + +if __name__ == "__main__": + unittest.main() diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index c27decf36..8adbc6064 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -107,6 +107,7 @@ suite_ascend = { TestFile("ascend/test_ascend_hicache_mla.py", 400), TestFile("ascend/test_ascend_tp4_bf16.py", 400), TestFile("ascend/test_ascend_w4a4_quantization.py", 600), + TestFile("ascend/test_llada2_mini_ascend.py", 800), ], "per-commit-16-npu-a3": [ TestFile("ascend/test_ascend_deepep.py", 3600),