From 4e1f8381c8087d3906213b2c8d83ffb0b5d1eb4b Mon Sep 17 00:00:00 2001 From: wangyifan194 Date: Tue, 18 Aug 2026 18:06:23 +0800 Subject: [PATCH] feat: support Qwen3.5 eagle-mode on NPU with Python executor. --- tests/python/test_registry.py | 6 +- xllm/core/kernels/npu/npu_ops_library.cpp | 229 ++++++++++++++++++ xllm/core/kernels/npu/tilelang/CMakeLists.txt | 2 + .../tilelang/causal_conv1d_update_wrapper.cpp | 176 ++++++++++++++ .../kernels/npu/tilelang/tilelang_ops_api.h | 17 ++ xllm/core/kernels/ops_api.cpp | 158 +----------- xllm/python/kernels_cuda/__init__.py | 2 + xllm/python/kernels_cuda/gated_delta_net.py | 39 +++ xllm/python/kernels_npu/__init__.py | 8 +- xllm/python/kernels_npu/_custom_op.py | 103 ++++++++ xllm/python/kernels_npu/causal_conv1d.py | 47 ++-- xllm/python/kernels_npu/gated_delta_net.py | 144 +++++++---- xllm/python/kernels_npu/normalization.py | 14 +- xllm/python/layers/gated_delta_net.py | 60 +++-- xllm/python/model_platform_support.py | 2 +- 15 files changed, 739 insertions(+), 268 deletions(-) create mode 100644 xllm/core/kernels/npu/tilelang/causal_conv1d_update_wrapper.cpp diff --git a/tests/python/test_registry.py b/tests/python/test_registry.py index c61b534a03..999e9569b7 100644 --- a/tests/python/test_registry.py +++ b/tests/python/test_registry.py @@ -21,10 +21,10 @@ def test_unsupported_model_fails_before_import(monkeypatch: pytest.MonkeyPatch) -> None: import_model = Mock() - monkeypatch.setattr(registry.current_platform, "device_type", lambda: "npu") + monkeypatch.setattr(registry.current_platform, "device_type", lambda: "cuda") monkeypatch.setattr(registry, "import_module", import_model) - with pytest.raises(NotImplementedError, match="qwen3_5.*npu"): - registry.get_model_class("qwen3_5") + with pytest.raises(NotImplementedError, match="qwen3_vl.*cuda"): + registry.get_model_class("qwen3_vl") import_model.assert_not_called() diff --git a/xllm/core/kernels/npu/npu_ops_library.cpp b/xllm/core/kernels/npu/npu_ops_library.cpp index 1c9ae56305..ec8fc1f8f9 100644 --- a/xllm/core/kernels/npu/npu_ops_library.cpp +++ b/xllm/core/kernels/npu/npu_ops_library.cpp @@ -27,6 +27,8 @@ limitations under the License. #include "kernels/npu/xllm_ops/xllm_ops_api.h" #include "npu_ops_api.h" +#include "tilelang/tilelang_ops_api.h" +#include "triton_npu/torch_api/triton_ops_api.h" namespace xllm { @@ -38,6 +40,194 @@ torch::Tensor rms_norm_npu(const torch::Tensor& input, return xllm::kernel::npu::rms_norm(input, weight, eps, "rmsnorm"); } +torch::Tensor rms_norm_gated_npu(const torch::Tensor& input, + const torch::Tensor& gate, + const torch::Tensor& weight, + double eps) { + return xllm::kernel::npu::layer_norm_fwd_aclnn(input, + weight, + /*bias=*/torch::Tensor(), + eps, + /*z=*/gate, + /*group_size=*/input.size(-1), + /*norm_before_gate=*/true, + /*is_rms_norm=*/true); +} + +torch::Tensor l2_norm_npu(torch::Tensor input, double eps) { + return xllm::kernel::npu::npu_l2norm_last_dim(input, eps); +} + +torch::Tensor causal_conv1d_update_npu(torch::Tensor x, + torch::Tensor conv_state, + torch::Tensor weight, + torch::Tensor state_indices) { + // Python layer stores weight as [dim, kernel_width]; tilelang expects + // [kernel_width, dim]. Transpose if the first dim is larger. + if (weight.size(0) > weight.size(1)) { + weight = weight.t().contiguous(); + } + return xllm::kernel::npu::tilelang::causal_conv1d_update( + x, + conv_state, + weight, + /*bias=*/std::nullopt, + /*conv_state_indices=*/state_indices, + /*query_start_loc=*/std::nullopt, + /*max_query_len=*/1, + /*activation=*/true); +} + +torch::Tensor causal_conv1d_prefill_npu(torch::Tensor x, + torch::Tensor weight, + torch::Tensor conv_state, + torch::Tensor state_indices, + torch::Tensor has_initial_state, + torch::Tensor query_start_loc) { + // Python layer stores weight as [dim, kernel_width]; CANN expects + // [kernel_width, dim]. + if (weight.size(0) > weight.size(1)) { + weight = weight.t().contiguous(); + } + + // Convert device tensors to host vectors for IntArrayRef parameters. + auto qsl_cpu = query_start_loc.to(torch::kCPU, torch::kInt64).contiguous(); + auto si_cpu = state_indices.to(torch::kCPU, torch::kInt64).contiguous(); + auto ism_cpu = has_initial_state.to(torch::kCPU, torch::kInt64).contiguous(); + + std::vector qsl_vec(qsl_cpu.data_ptr(), + qsl_cpu.data_ptr() + qsl_cpu.numel()); + std::vector si_vec(si_cpu.data_ptr(), + si_cpu.data_ptr() + si_cpu.numel()); + std::vector ism_vec(ism_cpu.data_ptr(), + ism_cpu.data_ptr() + ism_cpu.numel()); + + constexpr int64_t kActivationSilu = 1; + constexpr int64_t kPadSlotId = -1; + constexpr int64_t kRunModeForward = 0; + + return xllm::kernel::npu::causal_conv1d( + x, + weight, + conv_state, + /*bias_opt=*/std::nullopt, + torch::IntArrayRef(qsl_vec), + torch::IntArrayRef(si_vec), + torch::IntArrayRef(ism_vec), + /*num_accepted_tokens_opt=*/torch::IntArrayRef{}, + kActivationSilu, + kPadSlotId, + kRunModeForward); +} + +std::tuple +causal_conv1d_qkv_prefill_npu(torch::Tensor x, + torch::Tensor weight, + torch::Tensor conv_state, + torch::Tensor state_indices, + torch::Tensor has_initial_state, + torch::Tensor query_start_loc, + int64_t num_qk_heads, + int64_t num_v_heads, + int64_t head_k_dim, + int64_t head_v_dim) { + // Python layer stores weight as [dim, kernel_width]; CANN expects + // [kernel_width, dim]. + if (weight.size(0) > weight.size(1)) { + weight = weight.t().contiguous(); + } + + auto qsl_cpu = query_start_loc.to(torch::kCPU, torch::kInt64).contiguous(); + auto si_cpu = state_indices.to(torch::kCPU, torch::kInt64).contiguous(); + auto ism_cpu = has_initial_state.to(torch::kCPU, torch::kInt64).contiguous(); + + std::vector qsl_vec(qsl_cpu.data_ptr(), + qsl_cpu.data_ptr() + qsl_cpu.numel()); + std::vector si_vec(si_cpu.data_ptr(), + si_cpu.data_ptr() + si_cpu.numel()); + std::vector ism_vec(ism_cpu.data_ptr(), + ism_cpu.data_ptr() + ism_cpu.numel()); + + return xllm::kernel::npu::causal_conv1d_qkv(x, + weight, + conv_state, + torch::IntArrayRef(qsl_vec), + torch::IntArrayRef(si_vec), + torch::IntArrayRef(ism_vec), + num_qk_heads, + num_v_heads, + head_k_dim, + head_v_dim); +} + +std::tuple chunk_gated_delta_rule_npu( + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor g, + torch::Tensor beta, + torch::Tensor initial_state, + torch::Tensor cu_seqlens) { + return xllm::kernel::npu::npu_mega_chunk_gdn( + q, + k, + v, + g, + beta, + /*scale=*/std::nullopt, + /*initial_state=*/initial_state, + /*output_final_state=*/true, + /*cu_seqlens=*/cu_seqlens, + /*q_seq_lens=*/{}, + /*use_qk_l2norm_in_kernel=*/true); +} + +std::tuple fused_gdn_gating_npu( + torch::Tensor a_log, + torch::Tensor a, + torch::Tensor b, + torch::Tensor dt_bias) { + auto [g, beta] = xllm::kernel::npu::tilelang::fused_gdn_gating( + a_log.to(torch::kFloat32), + a, + b, + dt_bias.to(torch::kFloat32), + /*softplus_beta=*/1.0f, + /*softplus_threshold=*/20.0f); + return std::make_tuple(g, beta); +} + +torch::Tensor fused_sigmoid_gating_delta_rule_decode_npu( + torch::Tensor a_log, + torch::Tensor a, + torch::Tensor dt_bias, + torch::Tensor q, + torch::Tensor k, + torch::Tensor v, + torch::Tensor b, + torch::Tensor ssm_state, + torch::Tensor state_indices, + torch::Tensor cu_seqlens, + double scale) { + auto a_log_f32 = a_log.to(torch::kFloat32); + auto dt_bias_f32 = dt_bias.to(torch::kFloat32); + return xllm::kernel::npu::npu_fused_sigmoid_gating_delta_rule_update( + a_log_f32, + a, + dt_bias_f32, + q, + k, + v, + b, + ssm_state, + state_indices, + cu_seqlens, + /*scale=*/static_cast(scale), + /*use_qk_l2norm_in_kernel=*/true, + /*softplus_beta=*/1.0f, + /*softplus_threshold=*/20.0f); +} + std::tuple fused_add_rms_norm_npu( torch::Tensor& input, torch::Tensor& residual, @@ -279,6 +469,35 @@ void ensure_xllm_ops_registered() { // compiled only under USE_NPU (mutually exclusive with USE_CUDA). TORCH_LIBRARY(xllm_ops, m) { m.def("rms_norm(Tensor input, Tensor weight, float eps) -> Tensor"); + m.def( + "rms_norm_gated(Tensor input, Tensor gate, Tensor weight, float eps) -> " + "Tensor"); + m.def("l2_norm(Tensor input, float eps) -> Tensor"); + m.def( + "causal_conv1d_update(Tensor x, Tensor(a!) conv_state, Tensor weight, " + "Tensor state_indices) -> Tensor"); + m.def( + "chunk_gated_delta_rule(Tensor q, Tensor k, Tensor v, Tensor g, " + "Tensor beta, Tensor initial_state, Tensor cu_seqlens) -> " + "(Tensor, Tensor)"); + m.def( + "causal_conv1d_prefill(Tensor x, Tensor weight, Tensor(a!) conv_state, " + "Tensor state_indices, Tensor has_initial_state, " + "Tensor query_start_loc) -> Tensor"); + m.def( + "causal_conv1d_qkv_prefill(Tensor x, Tensor weight, " + "Tensor(a!) conv_state, Tensor state_indices, " + "Tensor has_initial_state, Tensor query_start_loc, " + "int num_qk_heads, int num_v_heads, " + "int head_k_dim, int head_v_dim) -> (Tensor, Tensor, Tensor)"); + m.def( + "fused_gdn_gating(Tensor a_log, Tensor a, Tensor b, Tensor dt_bias) -> " + "(Tensor, Tensor)"); + m.def( + "fused_sigmoid_gating_delta_rule_decode(Tensor a_log, Tensor a, " + "Tensor dt_bias, Tensor q, Tensor k, Tensor v, Tensor b, " + "Tensor(a!) ssm_state, Tensor state_indices, Tensor cu_seqlens, " + "float scale) -> Tensor"); m.def( "fused_add_rms_norm(Tensor(a!) input, Tensor(b!) residual, Tensor " "weight, " @@ -353,6 +572,16 @@ TORCH_LIBRARY(xllm_ops, m) { TORCH_LIBRARY_IMPL(xllm_ops, PrivateUse1, m) { m.impl("rms_norm", TORCH_FN(xllm::rms_norm_npu)); + m.impl("rms_norm_gated", TORCH_FN(xllm::rms_norm_gated_npu)); + m.impl("l2_norm", TORCH_FN(xllm::l2_norm_npu)); + m.impl("causal_conv1d_update", TORCH_FN(xllm::causal_conv1d_update_npu)); + m.impl("chunk_gated_delta_rule", TORCH_FN(xllm::chunk_gated_delta_rule_npu)); + m.impl("causal_conv1d_prefill", TORCH_FN(xllm::causal_conv1d_prefill_npu)); + m.impl("causal_conv1d_qkv_prefill", + TORCH_FN(xllm::causal_conv1d_qkv_prefill_npu)); + m.impl("fused_gdn_gating", TORCH_FN(xllm::fused_gdn_gating_npu)); + m.impl("fused_sigmoid_gating_delta_rule_decode", + TORCH_FN(xllm::fused_sigmoid_gating_delta_rule_decode_npu)); m.impl("fused_add_rms_norm", TORCH_FN(xllm::fused_add_rms_norm_npu)); m.impl("silu_and_mul", TORCH_FN(xllm::silu_and_mul_npu)); m.impl("reshape_paged_cache", TORCH_FN(xllm::reshape_paged_cache_npu)); diff --git a/xllm/core/kernels/npu/tilelang/CMakeLists.txt b/xllm/core/kernels/npu/tilelang/CMakeLists.txt index e8e37c7713..9df65926ac 100644 --- a/xllm/core/kernels/npu/tilelang/CMakeLists.txt +++ b/xllm/core/kernels/npu/tilelang/CMakeLists.txt @@ -180,6 +180,7 @@ cc_library( tilelang_ops_api.h SRCS ${TILELANG_KERNEL_SRCS} + causal_conv1d_update_wrapper.cpp DEPS torch torch_npu @@ -200,6 +201,7 @@ target_link_libraries(tilelang_kernels "$ENV{NPU_HOME_PATH}/lib64/libascend_dump.so" "$ENV{NPU_HOME_PATH}/lib64/libprofapi.so" "$ENV{NPU_HOME_PATH}/lib64/libmmpa.so" + "$ENV{NPU_HOME_PATH}/lib64/libunified_dlog.so" m dl ) diff --git a/xllm/core/kernels/npu/tilelang/causal_conv1d_update_wrapper.cpp b/xllm/core/kernels/npu/tilelang/causal_conv1d_update_wrapper.cpp new file mode 100644 index 0000000000..aec2110183 --- /dev/null +++ b/xllm/core/kernels/npu/tilelang/causal_conv1d_update_wrapper.cpp @@ -0,0 +1,176 @@ +/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved. + +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 "core/kernels/npu/tilelang/tilelang_ops_api.h" + +namespace xllm::kernel::npu::tilelang { + +torch::Tensor causal_conv1d_update( + torch::Tensor x, + torch::Tensor conv_state, + torch::Tensor weight, + const std::optional& bias, + const std::optional& conv_state_indices, + const std::optional& query_start_loc, + int32_t max_query_len, + bool activation, + const std::optional& initial_state_idx, + const std::optional& block_idx_last_scheduled_token, + const std::optional& initial_state_mode_opt) { + const bool has_silu = activation; + const int32_t dim = static_cast(x.size(-1)); + + auto bias_work = bias.has_value() && bias.value().defined() + ? bias.value() + : torch::zeros({dim}, x.options()); + + auto cu_seqlens = + query_start_loc.has_value() + ? query_start_loc.value().to(torch::kInt32) + : torch::arange( + 0, + x.size(0) + 1, + std::max(max_query_len, int32_t{1}), + torch::TensorOptions().dtype(torch::kInt32).device(x.device())); + + int64_t batch = cu_seqlens.size(0) - 1; + if (batch <= 0) { + return x; + } + + auto i32_opts = + torch::TensorOptions().dtype(torch::kInt32).device(x.device()); + + torch::Tensor init_indices; + torch::Tensor current_indices; + if (conv_state_indices.has_value()) { + auto ci = conv_state_indices.value().to(torch::kInt32); + if (ci.dim() == 1) { + init_indices = ci; + current_indices = ci; + } else { + auto ci_0 = ci.select(1, 0); + auto ci_1 = ci.select(1, 1); + if (initial_state_idx.has_value()) { + auto isi = initial_state_idx.value().to(torch::kInt32); + init_indices = torch::where(isi == 0, ci_0, ci_1); + } else { + init_indices = ci_0; + } + if (block_idx_last_scheduled_token.has_value()) { + auto bilt = block_idx_last_scheduled_token.value().to(torch::kInt32); + current_indices = torch::where(bilt == 0, ci_0, ci_1); + } else { + current_indices = ci_0; + } + } + } else { + init_indices = torch::arange(batch, i32_opts); + current_indices = init_indices; + } + + torch::Tensor initial_state_mode; + if (initial_state_mode_opt.has_value()) { + initial_state_mode = initial_state_mode_opt.value().to(torch::kInt32); + } else { + initial_state_mode = torch::ones({batch}, i32_opts); + } + + const bool is_3d = (x.dim() == 3); + auto x_flat = is_3d ? x.reshape({-1, dim}) : x; + + if (has_causal_conv1d_decode_specialization(batch, dim, has_silu)) { + auto conv_state_nonconst = conv_state; + auto y = causal_conv1d_decode( + /*conv_state=*/conv_state_nonconst, + /*x=*/x_flat, + /*weight=*/weight, + /*bias=*/bias_work, + /*init_indices=*/init_indices, + /*current_indices=*/current_indices, + /*initial_state_mode=*/initial_state_mode, + /*has_silu=*/has_silu); + + if (is_3d) { + y = y.view(x.sizes()); + } + return y; + } + + // Fallback: per-batch loop using causal_conv1d (batch=1 kernel, fp16). + auto original_dtype = x.scalar_type(); + bool need_cast = (original_dtype != torch::kFloat16); + + auto x_fp16 = need_cast ? x_flat.to(torch::kFloat16) : x_flat; + auto weight_fp16 = need_cast ? weight.to(torch::kFloat16) : weight; + auto conv_state_fp16 = + need_cast ? conv_state.to(torch::kFloat16).clone() : conv_state.clone(); + auto bias_fp16 = need_cast ? bias_work.to(torch::kFloat16) : bias_work; + + auto y_fp16 = torch::empty({x_flat.size(0), dim}, x_fp16.options()); + auto cu_seqlens_cpu = cu_seqlens.to(torch::kCPU); + const int32_t* cu_ptr = cu_seqlens_cpu.data_ptr(); + + for (int64_t b = 0; b < batch; ++b) { + int32_t seq_start_b = cu_ptr[b]; + int32_t seq_end_b = cu_ptr[b + 1]; + int32_t sb_len = seq_end_b - seq_start_b; + if (sb_len <= 0) { + continue; + } + + auto x_b = x_fp16.slice(0, seq_start_b, seq_end_b); + auto init_b = init_indices.slice(0, b, b + 1); + auto curr_b = current_indices.slice(0, b, b + 1); + auto ism_b = initial_state_mode.slice(0, b, b + 1); + + auto cu_b = torch::tensor( + {0, sb_len}, + torch::TensorOptions().dtype(torch::kInt32).device(x.device())); + + auto y_b = causal_conv1d(conv_state_fp16, + x_b, + weight_fp16, + bias_fp16, + cu_b, + init_b, + curr_b, + ism_b, + has_silu); + + y_fp16.slice(0, seq_start_b, seq_end_b).copy_(y_b); + } + + if (need_cast) { + conv_state.copy_(conv_state_fp16.to(original_dtype)); + } else { + conv_state.copy_(conv_state_fp16); + } + auto y = need_cast ? y_fp16.to(original_dtype) : y_fp16; + + if (is_3d) { + y = y.view(x.sizes()); + } + return y; +} + +} // namespace xllm::kernel::npu::tilelang diff --git a/xllm/core/kernels/npu/tilelang/tilelang_ops_api.h b/xllm/core/kernels/npu/tilelang/tilelang_ops_api.h index 9d42dac5e8..630e66ab2f 100644 --- a/xllm/core/kernels/npu/tilelang/tilelang_ops_api.h +++ b/xllm/core/kernels/npu/tilelang/tilelang_ops_api.h @@ -190,6 +190,23 @@ torch::Tensor causal_conv1d(torch::Tensor& conv_state, const torch::Tensor& initial_state_mode, bool has_silu); +// High-level causal conv1d update for decode/spec-verify on NPU. +// Handles parameter assembly (cu_seqlens, indices) and dispatches to +// causal_conv1d_decode specialization or fallback per-batch loop. +torch::Tensor causal_conv1d_update( + torch::Tensor x, + torch::Tensor conv_state, + torch::Tensor weight, + const std::optional& bias, + const std::optional& conv_state_indices, + const std::optional& query_start_loc, + int32_t max_query_len, + bool activation, + const std::optional& initial_state_idx = std::nullopt, + const std::optional& block_idx_last_scheduled_token = + std::nullopt, + const std::optional& initial_state_mode_opt = std::nullopt); + // Run fused sigmoid-gating delta-rule SSM scan on NPU. // Returns (out, final_state). // out: [T_padded, nv, dv] (padded token dim; caller strips padding) diff --git a/xllm/core/kernels/ops_api.cpp b/xllm/core/kernels/ops_api.cpp index 1eef606ad8..637657d679 100644 --- a/xllm/core/kernels/ops_api.cpp +++ b/xllm/core/kernels/ops_api.cpp @@ -1544,152 +1544,18 @@ torch::Tensor causal_conv1d_update(CausalConv1dUpdateParams& params) { params.block_idx_last_scheduled_token, params.initial_state_idx); #elif defined(USE_NPU) - const bool has_silu = params.activation; - - auto x_work = params.x; - auto weight_work = params.weight; - auto conv_state_work = params.conv_state; - - const int32_t dim = static_cast(x_work.size(1)); - - auto bias_work = params.bias.has_value() && params.bias.value().defined() - ? params.bias.value() - : torch::zeros({dim}, x_work.options()); - - auto conv_state_t = conv_state_work; - auto weight_t = weight_work; - - auto cu_seqlens = - params.query_start_loc.has_value() - ? params.query_start_loc.value().to(torch::kInt32) - : torch::arange(0, - x_work.size(0) + 1, - std::max(params.max_query_len, int32_t{1}), - torch::TensorOptions() - .dtype(torch::kInt32) - .device(x_work.device())); - - int64_t batch = cu_seqlens.size(0) - 1; - if (batch <= 0) { - return x_work; - } - - auto i32_opts = - torch::TensorOptions().dtype(torch::kInt32).device(x_work.device()); - - torch::Tensor init_indices; - torch::Tensor current_indices; - if (params.conv_state_indices.has_value()) { - auto ci = params.conv_state_indices.value().to(torch::kInt32); - if (ci.dim() == 1) { - init_indices = ci; - current_indices = ci; - } else { - auto ci_0 = ci.select(1, 0); - auto ci_1 = ci.select(1, 1); - if (params.initial_state_idx.has_value()) { - auto isi = params.initial_state_idx.value().to(torch::kInt32); - init_indices = torch::where(isi == 0, ci_0, ci_1); - } else { - init_indices = ci_0; - } - if (params.block_idx_last_scheduled_token.has_value()) { - auto bilt = - params.block_idx_last_scheduled_token.value().to(torch::kInt32); - current_indices = torch::where(bilt == 0, ci_0, ci_1); - } else { - current_indices = ci_0; - } - } - } else { - init_indices = torch::arange(batch, i32_opts); - current_indices = init_indices; - } - - torch::Tensor initial_state_mode; - if (params.initial_state_mode.has_value()) { - initial_state_mode = params.initial_state_mode.value().to(torch::kInt32); - } else { - initial_state_mode = torch::ones({batch}, i32_opts); - } - - const bool is_3d = (x_work.dim() == 3); - auto x_flat = is_3d ? x_work.reshape({-1, dim}) : x_work; - - if (npu::tilelang::has_causal_conv1d_decode_specialization( - batch, dim, has_silu)) { - auto conv_state_t_nonconst = conv_state_t; - auto y = npu::tilelang::causal_conv1d_decode( - /*conv_state=*/conv_state_t_nonconst, - /*x=*/x_flat, - /*weight=*/weight_t, - /*bias=*/bias_work, - /*init_indices=*/init_indices, - /*current_indices=*/current_indices, - /*initial_state_mode=*/initial_state_mode, - /*has_silu=*/has_silu); - - if (is_3d) { - y = y.view(x_work.sizes()); - } - return y; - } - - // Fallback: per-batch loop using causal_conv1d (batch=1 kernel, fp16). - auto original_dtype = x_work.scalar_type(); - bool need_cast = (original_dtype != torch::kFloat16); - - auto x_fp16 = need_cast ? x_flat.to(torch::kFloat16) : x_flat; - auto weight_fp16 = need_cast ? weight_work.to(torch::kFloat16) : weight_work; - auto conv_state_fp16 = need_cast ? conv_state_work.to(torch::kFloat16).clone() - : conv_state_work.clone(); - auto bias_fp16 = need_cast ? bias_work.to(torch::kFloat16) : bias_work; - - auto y_fp16 = torch::empty({x_flat.size(0), dim}, x_fp16.options()); - auto cu_seqlens_cpu = cu_seqlens.to(torch::kCPU); - const int32_t* cu_ptr = cu_seqlens_cpu.data_ptr(); - - for (int64_t b = 0; b < batch; ++b) { - int32_t seq_start_b = cu_ptr[b]; - int32_t seq_end_b = cu_ptr[b + 1]; - int32_t sb_len = seq_end_b - seq_start_b; - if (sb_len <= 0) { - continue; - } - - auto x_b = x_fp16.slice(0, seq_start_b, seq_end_b); - auto init_b = init_indices.slice(0, b, b + 1); - auto curr_b = current_indices.slice(0, b, b + 1); - auto ism_b = initial_state_mode.slice(0, b, b + 1); - - auto cu_b = torch::tensor( - {0, sb_len}, - torch::TensorOptions().dtype(torch::kInt32).device(x_work.device())); - - auto y_b = npu::tilelang::causal_conv1d(conv_state_fp16, - x_b, - weight_fp16, - bias_fp16, - cu_b, - init_b, - curr_b, - ism_b, - has_silu); - - y_fp16.slice(0, seq_start_b, seq_end_b).copy_(y_b); - } - - if (need_cast) { - params.conv_state.copy_(conv_state_fp16.to(original_dtype)); - } else { - params.conv_state.copy_(conv_state_fp16); - } - auto y = need_cast ? y_fp16.to(original_dtype) : y_fp16; - - if (is_3d) { - y = y.view(x_work.sizes()); - } - return y; + return npu::tilelang::causal_conv1d_update( + params.x, + params.conv_state, + params.weight, + params.bias, + params.conv_state_indices, + params.query_start_loc, + params.max_query_len, + params.activation, + params.initial_state_idx, + params.block_idx_last_scheduled_token, + params.initial_state_mode); #elif defined(USE_MUSA) // Default path has no capture-safe output buffer. MUSA GDN layers that need // a persistent buffer must call musa::causal_conv1d_update(params, diff --git a/xllm/python/kernels_cuda/__init__.py b/xllm/python/kernels_cuda/__init__.py index 36cbeca543..592c57fccb 100644 --- a/xllm/python/kernels_cuda/__init__.py +++ b/xllm/python/kernels_cuda/__init__.py @@ -48,6 +48,7 @@ chunk_gated_delta_rule, fused_gdn_prefill_post_conv, fused_recurrent_gated_delta_rule_packed_decode, + gdn_prefill_prepare, resolve_gdn_prefill_backend, ) from .linear import prepare_row_parallel_weight @@ -115,6 +116,7 @@ "causal_conv1d_prefill", "causal_conv1d_decode", "resolve_gdn_prefill_backend", + "gdn_prefill_prepare", "fused_gdn_prefill_post_conv", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", diff --git a/xllm/python/kernels_cuda/gated_delta_net.py b/xllm/python/kernels_cuda/gated_delta_net.py index b7409c3bb7..932570befc 100644 --- a/xllm/python/kernels_cuda/gated_delta_net.py +++ b/xllm/python/kernels_cuda/gated_delta_net.py @@ -199,9 +199,48 @@ def _chunk_gated_delta_rule_fake( return torch.empty_like(v), torch.empty_like(initial_state) +def gdn_prefill_prepare( + mixed_qkv: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + a_log: torch.Tensor, + dt_bias: torch.Tensor, + num_key_heads: int, + num_value_heads: int, + key_head_dim: int, + value_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused conv + split + l2norm + gating for prefill. + + Encapsulates the CUDA-optimal fusion strategy: causal_conv1d_prefill + produces packed convolved output, then fused_gdn_prefill_post_conv + splits into Q/K/V with l2norm and computes gating in one kernel. + + Returns: + (q, k, v, g, beta) with shapes [T, H, D] / [T, H]. + """ + from .causal_conv1d import causal_conv1d_prefill + + convolved = causal_conv1d_prefill( + mixed_qkv, weight, conv_state, state_indices, + has_initial_state, cu_seqlens, + ) + q, k, v, g, beta = fused_gdn_prefill_post_conv( + convolved, a, b, a_log, dt_bias, + num_key_heads, key_head_dim, value_head_dim, + ) + return q, k, v, g, beta + + __all__ = [ "GdnPrefillBackend", "resolve_gdn_prefill_backend", + "gdn_prefill_prepare", "fused_gdn_prefill_post_conv", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", diff --git a/xllm/python/kernels_npu/__init__.py b/xllm/python/kernels_npu/__init__.py index 29b3c4e1c7..5a73ad2bac 100644 --- a/xllm/python/kernels_npu/__init__.py +++ b/xllm/python/kernels_npu/__init__.py @@ -33,11 +33,12 @@ "update_decode_graph_metadata", "vision_fusion_attention", ), - "causal_conv1d": ("causal_conv1d_decode", "causal_conv1d_prefill"), + "causal_conv1d": ("causal_conv1d_decode", "causal_conv1d_qkv_prefill"), "gated_delta_net": ( "chunk_gated_delta_rule", - "fused_gdn_prefill_post_conv", + "fused_gdn_gating", "fused_recurrent_gated_delta_rule_packed_decode", + "gdn_prefill_prepare", "resolve_gdn_prefill_backend", ), "linear": ("prepare_row_parallel_weight",), @@ -99,10 +100,9 @@ "scatter_nd_update", "sparse_flash_attention", "sparse_flash_attention_out", - "causal_conv1d_prefill", "causal_conv1d_decode", "resolve_gdn_prefill_backend", - "fused_gdn_prefill_post_conv", + "gdn_prefill_prepare", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", ] diff --git a/xllm/python/kernels_npu/_custom_op.py b/xllm/python/kernels_npu/_custom_op.py index 6b9d91183b..283d7903a6 100644 --- a/xllm/python/kernels_npu/_custom_op.py +++ b/xllm/python/kernels_npu/_custom_op.py @@ -66,6 +66,97 @@ def _rms_norm_fake( return torch.empty_like(input) +def _rms_norm_gated_fake( + input: torch.Tensor, + gate: torch.Tensor, + weight: torch.Tensor, + eps: float, +) -> torch.Tensor: + del gate, weight, eps + return torch.empty_like(input) + + +def _l2_norm_fake( + input: torch.Tensor, + eps: float, +) -> torch.Tensor: + del eps + return torch.empty_like(input) + + +def _causal_conv1d_update_fake( + x: torch.Tensor, + conv_state: torch.Tensor, + weight: torch.Tensor, + state_indices: torch.Tensor, +) -> torch.Tensor: + del conv_state, weight, state_indices + return torch.empty_like(x) + + +def _chunk_gated_delta_rule_fake( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + del q, k, g, beta, cu_seqlens + num_seqs = initial_state.shape[0] + return torch.empty_like(v), torch.empty_like(initial_state) + + +def _causal_conv1d_qkv_prefill_fake( + x: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + query_start_loc: torch.Tensor, + num_qk_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + del weight, conv_state, state_indices, has_initial_state, query_start_loc + num_tokens = x.shape[0] + opts = x.options().dtype(torch.bfloat16) + q = torch.empty(1, num_tokens, num_qk_heads, head_k_dim, **opts) + k = torch.empty(1, num_tokens, num_qk_heads, head_k_dim, **opts) + v = torch.empty(1, num_tokens, num_v_heads, head_v_dim, **opts) + return q, k, v + + +def _fused_gdn_gating_fake( + a_log: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + dt_bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + del a_log, dt_bias + return torch.empty_like(a), torch.empty_like(b) + + +def _fused_sigmoid_gating_delta_rule_decode_fake( + a_log: torch.Tensor, + a: torch.Tensor, + dt_bias: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + b: torch.Tensor, + ssm_state: torch.Tensor, + state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + scale: float, +) -> torch.Tensor: + del a_log, a, dt_bias, k, b, ssm_state, state_indices, cu_seqlens, scale + # Output shape matches v: [num_tokens, num_value_heads, value_dim] + return torch.empty_like(v) + + def _fused_add_rms_norm_fake( input: torch.Tensor, residual: torch.Tensor, @@ -327,6 +418,18 @@ def _sparse_flash_attention_out_fake( register_fake("xllm_ops::rms_norm", _rms_norm_fake) +register_fake("xllm_ops::rms_norm_gated", _rms_norm_gated_fake) +register_fake("xllm_ops::l2_norm", _l2_norm_fake) +register_fake("xllm_ops::causal_conv1d_update", _causal_conv1d_update_fake) +register_fake("xllm_ops::chunk_gated_delta_rule", _chunk_gated_delta_rule_fake) +register_fake( + "xllm_ops::causal_conv1d_qkv_prefill", _causal_conv1d_qkv_prefill_fake +) +register_fake("xllm_ops::fused_gdn_gating", _fused_gdn_gating_fake) +register_fake( + "xllm_ops::fused_sigmoid_gating_delta_rule_decode", + _fused_sigmoid_gating_delta_rule_decode_fake, +) register_fake("xllm_ops::fused_add_rms_norm", _fused_add_rms_norm_fake) register_fake("xllm_ops::silu_and_mul", _silu_and_mul_fake) register_fake("xllm_ops::reshape_paged_cache", _reshape_paged_cache_fake) diff --git a/xllm/python/kernels_npu/causal_conv1d.py b/xllm/python/kernels_npu/causal_conv1d.py index 5f2fbe3a19..c4d99eb2ef 100644 --- a/xllm/python/kernels_npu/causal_conv1d.py +++ b/xllm/python/kernels_npu/causal_conv1d.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""NPU causal-convolution kernels. +"""NPU causal-convolution kernels (PyTorch small-op implementation). -Neither has an NPU kernel yet. The signatures are the contract an NPU -implementation has to meet; see ``kernels_cuda/causal_conv1d.py`` and the -Triton launcher it calls for the reference behaviour. +Implements the same semantics as the CUDA Triton reference in +``kernels_cuda/triton/causal_conv1d.py`` using only standard PyTorch +operations. Performance is not optimized; correctness and precision +alignment are the goals. """ from __future__ import annotations @@ -24,32 +25,28 @@ import torch -def causal_conv1d_prefill( +def causal_conv1d_qkv_prefill( value: torch.Tensor, weight: torch.Tensor, conv_state: torch.Tensor, state_indices: torch.Tensor, has_initial_state: torch.Tensor, query_start_loc: torch.Tensor, -) -> torch.Tensor: - """Convolve a variable-length batch and update the convolution states. - - Args: - value: Packed activations of shape ``[num_tokens, channels]``. - weight: Depthwise kernel of shape ``[channels, kernel_size]``. - conv_state: Per-sequence convolution state, updated in place. - state_indices: State slot of every sequence. - has_initial_state: Whether a sequence continues an earlier state. - query_start_loc: Start offset of every sequence in ``value``. + num_qk_heads: int, + num_v_heads: int, + head_k_dim: int, + head_v_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused conv + split into Q/K/V for prefill. Returns: - Convolved activations with the shape and dtype of ``value``. + (q, k, v) with shapes [1, T, num_qk_heads, head_k_dim], + [1, T, num_qk_heads, head_k_dim], [1, T, num_v_heads, head_v_dim]. """ - del value, weight, conv_state, state_indices, has_initial_state - del query_start_loc - raise NotImplementedError( - "causal_conv1d_prefill has no NPU kernel; see " - "kernels_cuda/triton/causal_conv1d.py for the reference implementation" + return torch.ops.xllm_ops.causal_conv1d_qkv_prefill( + value, weight, conv_state, state_indices, + has_initial_state.to(torch.int64), query_start_loc, + num_qk_heads, num_v_heads, head_k_dim, head_v_dim, ) @@ -70,11 +67,9 @@ def causal_conv1d_decode( Returns: Convolved activations with the shape and dtype of ``value``. """ - del value, weight, conv_state, state_indices - raise NotImplementedError( - "causal_conv1d_decode has no NPU kernel; see " - "kernels_cuda/triton/causal_conv1d.py for the reference implementation" + return torch.ops.xllm_ops.causal_conv1d_update( + value, conv_state, weight, state_indices ) -__all__ = ["causal_conv1d_prefill", "causal_conv1d_decode"] +__all__ = ["causal_conv1d_qkv_prefill", "causal_conv1d_decode"] diff --git a/xllm/python/kernels_npu/gated_delta_net.py b/xllm/python/kernels_npu/gated_delta_net.py index 3fa3577a38..9eddc53a2a 100644 --- a/xllm/python/kernels_npu/gated_delta_net.py +++ b/xllm/python/kernels_npu/gated_delta_net.py @@ -12,11 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""NPU gated-delta-network kernels. +"""NPU gated-delta-network kernels (PyTorch small-op implementation). -None has an NPU kernel yet. The signatures are the contract an NPU -implementation has to meet; see ``kernels_cuda/gated_delta_net.py`` and the -Triton launchers it calls for the reference behaviour. +Implements the same semantics as the CUDA Triton references in +``kernels_cuda/triton/gdn_prefill.py`` and ``kernels_cuda/triton/gated_delta_net.py`` +using only standard PyTorch operations. Performance is not optimized; +correctness and precision alignment are the goals. """ from __future__ import annotations @@ -25,64 +26,68 @@ import torch -GdnPrefillBackend = Literal["flashinfer", "triton"] +GdnPrefillBackend = Literal["pytorch_naive"] def resolve_gdn_prefill_backend( capability: tuple[int, int] | None = None, ) -> GdnPrefillBackend: - """Select the prefill backend of the active device. + """Select the prefill backend for NPU. Args: - capability: Device capability to resolve for; ``None`` reads it from - the current device. + capability: Ignored on NPU. Returns: - The name to pass as ``backend`` to :func:`chunk_gated_delta_rule`. + The backend name to pass to :func:`chunk_gated_delta_rule`. """ del capability - raise NotImplementedError( - "resolve_gdn_prefill_backend has no NPU implementation; gated delta networks are not supported on NPU yet" - ) + return "pytorch_naive" + + +def fused_gdn_gating( + a_log: torch.Tensor, + a: torch.Tensor, + b: torch.Tensor, + dt_bias: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute decay gate g and beta from raw projections via TileLang kernel.""" + return torch.ops.xllm_ops.fused_gdn_gating(a_log, a, b, dt_bias) -def fused_gdn_prefill_post_conv( +def gdn_prefill_prepare( mixed_qkv: torch.Tensor, + weight: torch.Tensor, + conv_state: torch.Tensor, + state_indices: torch.Tensor, + has_initial_state: torch.Tensor, + cu_seqlens: torch.Tensor, a: torch.Tensor, b: torch.Tensor, a_log: torch.Tensor, dt_bias: torch.Tensor, num_key_heads: int, + num_value_heads: int, key_head_dim: int, value_head_dim: int, -) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, -]: - """Split the post-convolution projection and build the recurrence gates. +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fused conv + split + l2norm + gating for prefill. - Args: - mixed_qkv: Packed projection of shape ``[num_tokens, qkv_size]``. - a: Gate projection of shape ``[num_tokens, num_value_heads]``. - b: Beta projection of shape ``[num_tokens, num_value_heads]``. - a_log: Per-head log decay of shape ``[num_value_heads]``. - dt_bias: Per-head timestep bias of shape ``[num_value_heads]``. - num_key_heads: Key heads on this rank. - key_head_dim: Size of one key head. - value_head_dim: Size of one value head. + Encapsulates the NPU-optimal fusion strategy: causal_conv1d_qkv does + conv + split + l2norm in one kernel, then fused_gdn_gating computes + decay and beta independently. Returns: - Query, key, value, the decay gate and beta. + (q, k, v, g, beta) with shapes [T, H, D] / [T, H]. """ - del mixed_qkv, a, b, a_log, dt_bias - del num_key_heads, key_head_dim, value_head_dim - raise NotImplementedError( - "fused_gdn_prefill_post_conv has no NPU kernel; see " - "kernels_cuda/triton/gdn_prefill.py for the reference implementation" + from .causal_conv1d import causal_conv1d_qkv_prefill + + q, k, v = causal_conv1d_qkv_prefill( + mixed_qkv, weight, conv_state, state_indices, + has_initial_state, cu_seqlens, + num_key_heads, num_value_heads, key_head_dim, value_head_dim, ) + g, beta = fused_gdn_gating(a_log, a, b, dt_bias) + return q.squeeze(0), k.squeeze(0), v.squeeze(0), g, beta def fused_recurrent_gated_delta_rule_packed_decode( @@ -104,17 +109,54 @@ def fused_recurrent_gated_delta_rule_packed_decode( a_log: Per-head log decay of shape ``[num_value_heads]``. dt_bias: Per-head timestep bias of shape ``[num_value_heads]``. initial_state: Recurrent state pool, updated in place. + Shape ``[num_slots, num_value_heads, value_dim, key_dim]``. state_indices: State slot of every sequence. scale: Query scale. Returns: Output of shape ``[batch_size, 1, num_value_heads, value_head_dim]``. """ - del mixed_qkv, a, b, a_log, dt_bias, initial_state, state_indices, scale - raise NotImplementedError( - "fused_recurrent_gated_delta_rule_packed_decode has no NPU kernel; see " - "kernels_cuda/triton/gated_delta_net.py for the reference implementation" + batch = mixed_qkv.shape[0] + num_value_heads, value_dim, key_dim = initial_state.shape[-3:] + qkv_dim = mixed_qkv.shape[1] + query_key_dim = qkv_dim - num_value_heads * value_dim + query_dim = query_key_dim // 2 + num_key_heads = query_dim // key_dim + + # Split mixed_qkv + q_flat = mixed_qkv[:, :query_dim] + k_flat = mixed_qkv[:, query_dim : 2 * query_dim] + v_flat = mixed_qkv[:, 2 * query_dim:] + + q = q_flat.view(batch, num_key_heads, key_dim) + k = k_flat.view(batch, num_key_heads, key_dim) + v = v_flat.view(batch, num_value_heads, value_dim) + + # Kernel expects [batch, seq_len, heads, dim] — add seq dim for decode + q = q.unsqueeze(1) # [batch, 1, num_key_heads, key_dim] + k = k.unsqueeze(1) # [batch, 1, num_key_heads, key_dim] + v = v.unsqueeze(1) # [batch, 1, num_value_heads, value_dim] + + # Kernel does l2norm, gating, GQA expansion, and recurrence internally. + # cu_seqlens for decode: each seq has 1 token. + cu_seqlens = torch.arange( + batch + 1, dtype=torch.int32, device=mixed_qkv.device + ) + + output = torch.ops.xllm_ops.fused_sigmoid_gating_delta_rule_decode( + a_log, + a.unsqueeze(1), + dt_bias, + q.contiguous(), + k.contiguous(), + v.contiguous(), + b.unsqueeze(1), + initial_state, + state_indices, + cu_seqlens, + scale, ) + return output.unsqueeze(1) def chunk_gated_delta_rule( @@ -136,22 +178,34 @@ def chunk_gated_delta_rule( g: Decay gate of shape ``[num_tokens, num_value_heads]``. beta: Beta with the shape of ``g``. initial_state: Recurrent state each sequence starts from. + Shape ``[batch, num_value_heads, value_dim, key_dim]``. cu_seqlens: Cumulative sequence lengths. - backend: Name returned by :func:`resolve_gdn_prefill_backend`. + backend: Ignored on NPU. Returns: The output with the shape of ``v`` and the final recurrent state. """ - del q, k, v, g, beta, initial_state, cu_seqlens, backend - raise NotImplementedError( - "chunk_gated_delta_rule has no NPU kernel; see kernels_cuda/triton/fla/ for the reference implementation" + del backend + # npu_mega_chunk_gdn expects [B, T, H, D] layout with B=1 for packed input + # Cast g and beta to match C++ layer behavior (bf16 round-trip) + g_input = g.to(v.dtype) + beta_input = beta.to(v.dtype) + output, final_state = torch.ops.xllm_ops.chunk_gated_delta_rule( + q.unsqueeze(0), + k.unsqueeze(0), + v.unsqueeze(0), + g_input.unsqueeze(0), + beta_input.unsqueeze(0), + initial_state, + cu_seqlens, ) + return output.squeeze(0), final_state __all__ = [ "GdnPrefillBackend", "resolve_gdn_prefill_backend", - "fused_gdn_prefill_post_conv", + "gdn_prefill_prepare", "fused_recurrent_gated_delta_rule_packed_decode", "chunk_gated_delta_rule", ] diff --git a/xllm/python/kernels_npu/normalization.py b/xllm/python/kernels_npu/normalization.py index 75a0aa98f2..4c94d94e7d 100644 --- a/xllm/python/kernels_npu/normalization.py +++ b/xllm/python/kernels_npu/normalization.py @@ -32,10 +32,7 @@ def l2_norm(value: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: Returns: A tensor with the shape and dtype of ``value``. """ - del value, eps - raise NotImplementedError( - "l2_norm has no NPU kernel; see kernels_cuda/triton/l2_norm.py for the reference implementation" - ) + return torch.ops.xllm_ops.l2_norm(value, eps) def rms_norm_gated( @@ -44,21 +41,18 @@ def rms_norm_gated( weight: torch.Tensor, eps: float = 1e-6, ) -> torch.Tensor: - """Apply RMSNorm to ``value`` and gate the result with ``gate``. + """Apply RMSNorm to ``value`` and gate the result with ``silu(gate)``. Args: value: Tensor to normalize. - gate: Gate applied after normalization, same shape as ``value``. + gate: Gate applied after normalization (SiLU is applied internally). weight: RMSNorm weight over the last dimension. eps: RMSNorm epsilon. Returns: A tensor with the shape and dtype of ``value``. """ - del value, gate, weight, eps - raise NotImplementedError( - "rms_norm_gated has no NPU kernel; see kernels_cuda/triton/rms_norm.py for the reference implementation" - ) + return torch.ops.xllm_ops.rms_norm_gated(value, gate, weight, eps) __all__ = ["rms_norm", "fused_add_rms_norm", "l2_norm", "rms_norm_gated"] diff --git a/xllm/python/layers/gated_delta_net.py b/xllm/python/layers/gated_delta_net.py index cd081b396e..9e2233bed5 100644 --- a/xllm/python/layers/gated_delta_net.py +++ b/xllm/python/layers/gated_delta_net.py @@ -125,15 +125,25 @@ def _conv_prefill( state_indices: torch.Tensor, has_initial_state: torch.Tensor, cu_seqlens: torch.Tensor, - ) -> torch.Tensor: + a: torch.Tensor, + b: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: self._conv_state_dim_first(conv_state) - return kernels.causal_conv1d_prefill( + return kernels.gdn_prefill_prepare( mixed_qkv, self.conv1d_weight, conv_state, state_indices, has_initial_state, cu_seqlens, + a, + b, + self.A_log, + self.dt_bias, + self.num_k_heads, + self.num_v_heads, + self.key_head_dim, + self.value_head_dim, ) def _conv_decode( @@ -152,17 +162,16 @@ def _conv_decode( def _gdn_prefill( self, - mixed_qkv: torch.Tensor, - a: torch.Tensor, - b: torch.Tensor, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, ssm_state: torch.Tensor, state_indices: torch.Tensor, has_initial_state: torch.Tensor, cu_seqlens: torch.Tensor, ) -> torch.Tensor: - # TODO: Fuse cache gather/zero/scatter and null-row output masking into - # the GDN kernels. The current staging is intentionally kept for this - # correctness PR and should be removed in the next performance PR. non_null_state = state_indices > 0 use_initial_state = non_null_state & has_initial_state cache_indices = state_indices.to(torch.long) @@ -173,16 +182,6 @@ def _gdn_prefill( initial_state, torch.zeros_like(initial_state), ) - q, k, v, g, beta = kernels.fused_gdn_prefill_post_conv( - mixed_qkv=mixed_qkv, - a=a, - b=b, - a_log=self.A_log, - dt_bias=self.dt_bias, - num_key_heads=self.num_k_heads, - key_head_dim=self.key_head_dim, - value_head_dim=self.value_head_dim, - ) output, final_state = kernels.chunk_gated_delta_rule( q, k, @@ -193,11 +192,12 @@ def _gdn_prefill( cu_seqlens, self.gdn_prefill_backend, ) + num_tokens = q.shape[0] sequence_lengths = cu_seqlens.diff().to(dtype=torch.long) token_mask = torch.repeat_interleave( non_null_state, sequence_lengths, - output_size=mixed_qkv.shape[0], + output_size=num_tokens, ) output = torch.where(token_mask[:, None, None], output, 0.0) ssm_state.index_copy_( @@ -255,26 +255,20 @@ def forward(self, hidden: torch.Tensor) -> torch.Tensor: raise RuntimeError("has_initial_state is required by Qwen3.5 prefill") has_initial_state = has_initial_state.to(device=hidden.device, dtype=torch.bool) if has_initial_state.shape != state_indices.shape: - raise ValueError("has_initial_state must match linear_state_indices") - mixed_qkv = self._conv_prefill( - mixed_qkv, - conv_state, - state_indices, - has_initial_state, - cu_seqlens, + raise ValueError( + "has_initial_state must match linear_state_indices" + ) + q, k, v, g, beta = self._conv_prefill( + mixed_qkv, conv_state, state_indices, + has_initial_state, cu_seqlens, a, b, ) else: mixed_qkv = self._conv_decode(mixed_qkv, conv_state, state_indices) if is_prefill: output = self._gdn_prefill( - mixed_qkv, - a, - b, - ssm_state, - state_indices, - has_initial_state, - cu_seqlens, + q, k, v, g, beta, + ssm_state, state_indices, has_initial_state, cu_seqlens, ) else: output = self._gdn_decode(mixed_qkv, a, b, ssm_state, state_indices) diff --git a/xllm/python/model_platform_support.py b/xllm/python/model_platform_support.py index caae712526..711d648c90 100644 --- a/xllm/python/model_platform_support.py +++ b/xllm/python/model_platform_support.py @@ -16,7 +16,7 @@ MODEL_PLATFORM_SUPPORT: dict[str, dict[str, bool]] = { "qwen3": {"cuda": True, "npu": True}, - "qwen3_5": {"cuda": True, "npu": False}, + "qwen3_5": {"cuda": True, "npu": True}, "qwen3_vl": {"cuda": False, "npu": True}, "deepseek_v32": {"cuda": False, "npu": True}, "glm5_2": {"cuda": False, "npu": True},