diff --git a/xllm/core/common/global_flags.h b/xllm/core/common/global_flags.h index 46b10f196b..22cf5e6c26 100755 --- a/xllm/core/common/global_flags.h +++ b/xllm/core/common/global_flags.h @@ -374,6 +374,8 @@ DECLARE_string(dit_sparse_attention_version); DECLARE_int64(dit_sparse_attention_mask_refresh_steps); +DECLARE_bool(dit_laser_attention_enabled); + DECLARE_bool(use_audio_in_video); // --- kernel config --- diff --git a/xllm/core/framework/config/dit_config.cpp b/xllm/core/framework/config/dit_config.cpp index 585e219f6c..9e8b6502dd 100644 --- a/xllm/core/framework/config/dit_config.cpp +++ b/xllm/core/framework/config/dit_config.cpp @@ -61,6 +61,11 @@ DEFINE_bool(dit_debug_print, false, "whether print the debug info for dit models"); +DEFINE_bool(dit_laser_attention_enabled, + false, + "whether to use the laser attention kernel (MindIE-SD, tuned for " + "Wan2.2) in place of npu_fusion_attention for DiT attention."); + DEFINE_int64(dit_generation_image_area_max, 0, "Maximum allowed image area (width * height) for image generation " @@ -123,6 +128,7 @@ void DiTConfig::from_flags() { XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_cache_end_blocks); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_debug_print); + XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_laser_attention_enabled); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_generation_image_area_max); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_vae_image_size); XLLM_CONFIG_ASSIGN_FROM_FLAG(dit_enable_vae_tiling); @@ -147,6 +153,7 @@ void DiTConfig::from_json(const JsonReader& json) { XLLM_CONFIG_ASSIGN_FROM_JSON(dit_cache_end_blocks); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_sp_communication_overlap); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_debug_print); + XLLM_CONFIG_ASSIGN_FROM_JSON(dit_laser_attention_enabled); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_generation_image_area_max); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_vae_image_size); XLLM_CONFIG_ASSIGN_FROM_JSON(dit_enable_vae_tiling); @@ -184,6 +191,8 @@ void DiTConfig::append_config_json(nlohmann::ordered_json& config_json) const { config_json, default_config, dit_sp_communication_overlap); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, dit_debug_print); + APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( + config_json, default_config, dit_laser_attention_enabled); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( config_json, default_config, dit_generation_image_area_max); APPEND_CONFIG_JSON_VALUE_IF_NOT_DEFAULT( diff --git a/xllm/core/framework/config/dit_config.h b/xllm/core/framework/config/dit_config.h index c687ba11da..99b461736d 100644 --- a/xllm/core/framework/config/dit_config.h +++ b/xllm/core/framework/config/dit_config.h @@ -53,6 +53,7 @@ class DiTConfig final { "dit_cache_end_blocks", "dit_sp_communication_overlap", "dit_debug_print", + "dit_laser_attention_enabled", "dit_generation_image_area_max", "dit_vae_image_size", "dit_enable_vae_tiling", @@ -89,6 +90,8 @@ class DiTConfig final { PROPERTY(bool, dit_debug_print) = false; + PROPERTY(bool, dit_laser_attention_enabled) = false; + PROPERTY(int64_t, dit_generation_image_area_max) = 0; PROPERTY(int64_t, dit_vae_image_size) = 1048576; diff --git a/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt b/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt index 09bd02d541..459f53b85a 100644 --- a/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt +++ b/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt @@ -8,6 +8,7 @@ cc_library( ../utils.h SRCS replace_token.cpp + laser_attention.cpp top_k_top_p.cpp quant_matmul.cpp quantize.cpp diff --git a/xllm/core/kernels/npu/xllm_ops/laser_attention.cpp b/xllm/core/kernels/npu/xllm_ops/laser_attention.cpp new file mode 100644 index 0000000000..62120db976 --- /dev/null +++ b/xllm/core/kernels/npu/xllm_ops/laser_attention.cpp @@ -0,0 +1,129 @@ +/* Copyright 2026 The xLLM Authors. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ + +#include +#include +#include +#include +#include + +#include +#ifdef TORCH_HIGHER_THAN_PTA6 +#include +#else +#include +#include +#endif + +#include "acl/acl.h" +#include "aclnn_laser_attention.h" +#include "core/common/macros.h" +#include "core/kernels/npu/aclnn/pytorch_npu_helper.hpp" +#include "core/kernels/npu/utils.h" +#include "xllm_ops_api.h" + +namespace xllm::kernel::npu { + +// Laser attention kernel tuned for Wan2.2. q/k/v and output are in BNSD +// layout; output is cast back to the input dtype (kernel emits fp32). +torch::Tensor laser_attention(const torch::Tensor& q_bnsd, + const torch::Tensor& k_bnsd, + const torch::Tensor& v_bnsd, + double scale_value, + int64_t head_num) { + check_tensor(q_bnsd, "query", "laser_attention"); + check_tensor(k_bnsd, "key", "laser_attention"); + check_tensor(v_bnsd, "value", "laser_attention"); + + const auto input_dtype = q_bnsd.scalar_type(); + // Pad seq to a multiple of 256 and head_dim to 128, masking padded keys via + // pre_tokens; slice the padding back off the output afterwards. + namespace F = torch::nn::functional; + const int64_t kSeqBase = 256; + const int64_t kDimBase = 128; + const int64_t kMaxToken = 2147483647; // MAX_TOKEN = 2^31 - 1 + + const int64_t q_seqlen = q_bnsd.size(2); // BNSD: real query seq len + const int64_t kv_seqlen = k_bnsd.size(2); // real key/value seq len + const int64_t head_dim = q_bnsd.size(3); // real head dim + + // 1. cast to fp16 (kernel is fp16-only). + auto q_f16 = q_bnsd.to(torch::kHalf).contiguous(); + auto k_f16 = k_bnsd.to(torch::kHalf).contiguous(); + auto v_f16 = v_bnsd.to(torch::kHalf).contiguous(); + + // 2. zero-pad seq (dim=2) to a multiple of 256 and head_dim (dim=3) to 128. + auto pad_to = [](const torch::Tensor& t, int64_t base, int64_t dim) { + int64_t sz = t.size(dim); + if (sz % base == 0) return t; + int64_t pad = ((sz / base) + 1) * base - sz; + // PadFuncOptions pads from the last dim backwards. + std::vector spec = (dim == 3) ? std::vector{0, pad} + : std::vector{0, 0, 0, pad}; + return F::pad(t, F::PadFuncOptions(spec).mode(torch::kConstant).value(0)) + .contiguous(); + }; + q_f16 = pad_to(pad_to(q_f16, kSeqBase, 2), kDimBase, 3); + k_f16 = pad_to(pad_to(k_f16, kSeqBase, 2), kDimBase, 3); + v_f16 = pad_to(pad_to(v_f16, kSeqBase, 2), kDimBase, 3); + + const auto q_c = q_f16; + const auto k_c = k_f16; + const auto v_c = v_f16; + + // 3. pre_tokens = number of trailing padded keys to exclude from softmax. + int64_t pre_tokens = kMaxToken; + if (kv_seqlen % kSeqBase != 0) { + pre_tokens = (kv_seqlen / kSeqBase + 1) * kSeqBase - kv_seqlen; + } + const int64_t next_tokens = 1; + + // Kernel outputs are fp32. + auto attn_out = + torch::empty({q_c.size(0), q_c.size(1), q_c.size(2), q_c.size(3)}, + q_c.options().dtype(torch::kFloat)); + auto softmax_lms = torch::empty({q_c.size(0), q_c.size(1), q_c.size(2)}, + q_c.options().dtype(torch::kFloat)); + + char input_layout[] = "BNSD"; + const double keep_prob = 1.0; + const bool is_high_precision = true; + const c10::optional no_mask; // atten/alibi/drop masks: unused + + // Two-stage aclnn call (GetWorkspaceSize + execute) handled by the macro. + EXEC_NPU_CMD(aclnnLaserAttention, + q_c, + k_c, + v_c, + no_mask, + no_mask, + no_mask, + scale_value, + head_num, + input_layout, + keep_prob, + pre_tokens, + next_tokens, + is_high_precision, + softmax_lms, + attn_out); + + // Slice off the seq/head_dim padding, then cast back to the caller's dtype. + auto out_real = + attn_out.slice(2, 0, q_seqlen).slice(3, 0, head_dim).contiguous(); + return out_real.to(input_dtype); +} + +} // namespace xllm::kernel::npu diff --git a/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h b/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h index 71f21f1348..194deb7088 100644 --- a/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h +++ b/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h @@ -62,6 +62,14 @@ void replace_token(torch::Tensor& dst, torch::Tensor& src, bool synchronize_stream = true); +// Laser attention (MindIE-SD kernel tuned for Wan2.2). q/k/v in BNSD layout; +// returns attention output in BNSD layout, cast back to the input dtype. +torch::Tensor laser_attention(const torch::Tensor& q_bnsd, + const torch::Tensor& k_bnsd, + const torch::Tensor& v_bnsd, + double scale_value, + int64_t head_num); + void beam_search_rec(const torch::Tensor& logprobs, const torch::Tensor& top_tokens, const torch::Tensor& top_logprobs, diff --git a/xllm/models/dit/transformers/transformer_wan.h b/xllm/models/dit/transformers/transformer_wan.h index e123150be2..9a702ee2af 100644 --- a/xllm/models/dit/transformers/transformer_wan.h +++ b/xllm/models/dit/transformers/transformer_wan.h @@ -52,6 +52,7 @@ using xllm::dit::TpOptions; #endif #include "models/model_registry.h" #if defined(USE_NPU) +#include "core/kernels/npu/xllm_ops/xllm_ops_api.h" #include "torch_npu/csrc/aten/CustomFunctions.h" #endif @@ -672,20 +673,32 @@ class WanAttentionImpl : public torch::nn::Module { const int64_t head_num = q_t.size(1); const int64_t head_dim = q_t.size(-1); - const auto results = at_npu::native::custom_ops::npu_fusion_attention( - q_t, - k_t, - v_t, - head_num, - "BNSD", - torch::nullopt, - torch::nullopt, - torch::nullopt, - std::pow(head_dim, -0.5), - 1.0, - 65535, - 65535); - torch::Tensor out = std::get<0>(results).transpose(1, 2); + torch::Tensor out; + // Laser attention only supports equal-length q/k (self-attention); cross + // attention (q/k different seq len) falls back to npu_fusion_attention. + const bool laser_enable = + DiTConfig::get_instance().dit_laser_attention_enabled() && + q_t.size(2) == k_t.size(2); + if (laser_enable) { + out = xllm::kernel::npu::laser_attention( + q_t, k_t, v_t, std::pow(head_dim, -0.5), head_num) + .transpose(1, 2); + } else { + const auto results = at_npu::native::custom_ops::npu_fusion_attention( + q_t, + k_t, + v_t, + head_num, + "BNSD", + torch::nullopt, + torch::nullopt, + torch::nullopt, + std::pow(head_dim, -0.5), + 1.0, + 65535, + 65535); + out = std::get<0>(results).transpose(1, 2); + } #else constexpr int64_t kAttentionChunkSize = 512; constexpr int64_t kHeadDim = 1;