diff --git a/.gitmodules b/.gitmodules index 08b00197e0..a8a9a3dc43 100755 --- a/.gitmodules +++ b/.gitmodules @@ -20,7 +20,7 @@ fetchRecurseSubmodules = false [submodule "third_party/xllm_ops"] path = third_party/xllm_ops - url = https://gitcode.com/xLLM-AI/xllm_ops.git + url = https://github.com/xLLM-AI/xllm-ops.git fetchRecurseSubmodules = true [submodule "third_party/etcd_cpp_apiv3"] path = third_party/etcd_cpp_apiv3 diff --git a/third_party/xllm_ops b/third_party/xllm_ops index 818f04b2fc..fa7859e666 160000 --- a/third_party/xllm_ops +++ b/third_party/xllm_ops @@ -1 +1 @@ -Subproject commit 818f04b2fc5aed617fcd2e750936ef7d19956ad5 +Subproject commit fa7859e66694a2ece7e41492c3a8e95a415e15b5 diff --git a/xllm/core/kernels/npu/npu_causal_conv1d.cpp b/xllm/core/kernels/npu/npu_causal_conv1d.cpp index cca76c7f3a..94fc367bdb 100644 --- a/xllm/core/kernels/npu/npu_causal_conv1d.cpp +++ b/xllm/core/kernels/npu/npu_causal_conv1d.cpp @@ -83,4 +83,70 @@ torch::Tensor causal_conv1d(const torch::Tensor& x, return output; } +std::tuple causal_conv1d_qkv( + const torch::Tensor& x, + const torch::Tensor& weight, + const torch::Tensor& conv_state, + const torch::IntArrayRef query_start_loc_opt, + const torch::IntArrayRef cache_indices_opt, + const torch::IntArrayRef initial_state_mode_opt, + int64_t num_qk_heads, + int64_t num_v_heads, + int64_t head_k_dim, + int64_t head_v_dim) { + constexpr int64_t kPackedQkvActivationMode = 2; + constexpr int64_t kPadSlotId = -1; + constexpr int64_t kForwardRunMode = 0; + + check_tensor(x, "x", "causal_conv1d_qkv"); + check_tensor(weight, "weight", "causal_conv1d_qkv"); + check_tensor(conv_state, "conv_state", "causal_conv1d_qkv"); + CHECK_EQ(x.dim(), 2) << "causal_conv1d_qkv expects x with shape [T, D]"; + CHECK_GT(num_qk_heads, 0) << "num_qk_heads must be positive"; + CHECK_GT(num_v_heads, 0) << "num_v_heads must be positive"; + CHECK_EQ(head_k_dim, 128) << "packed QKV requires head_k_dim=128"; + CHECK_EQ(head_v_dim, head_k_dim) + << "packed QKV requires equal Q/K and V head dimensions"; + const int64_t q_elements_per_token = num_qk_heads * head_k_dim; + const int64_t k_elements_per_token = num_qk_heads * head_k_dim; + const int64_t v_elements_per_token = num_v_heads * head_v_dim; + CHECK_EQ(x.size(1), + q_elements_per_token + k_elements_per_token + v_elements_per_token) + << "causal_conv1d_qkv input width does not match the local Q/K/V layout"; + + auto packed = torch::empty(x.sizes(), x.options().dtype(torch::kBFloat16)); + c10::optional bias_opt = c10::nullopt; + torch::IntArrayRef num_accepted_tokens_opt; + EXEC_NPU_CMD(aclnnCausalConv1dQkv, + x, + weight, + bias_opt, + conv_state, + query_start_loc_opt, + cache_indices_opt, + initial_state_mode_opt, + num_accepted_tokens_opt, + kPackedQkvActivationMode, + kPadSlotId, + kForwardRunMode, + q_elements_per_token, + k_elements_per_token, + v_elements_per_token, + head_k_dim, + packed); + + const int64_t num_tokens = x.size(0); + auto packed_flat = packed.view({-1}); + const int64_t q_elements = num_tokens * q_elements_per_token; + const int64_t k_elements = num_tokens * k_elements_per_token; + const int64_t v_elements = num_tokens * v_elements_per_token; + auto q = packed_flat.narrow(0, 0, q_elements) + .view({1, num_tokens, num_qk_heads, head_k_dim}); + auto k = packed_flat.narrow(0, q_elements, k_elements) + .view({1, num_tokens, num_qk_heads, head_k_dim}); + auto v = packed_flat.narrow(0, q_elements + k_elements, v_elements) + .view({1, num_tokens, num_v_heads, head_v_dim}); + return {q, k, v}; +} + } // namespace xllm::kernel::npu diff --git a/xllm/core/kernels/npu/npu_ops_api.h b/xllm/core/kernels/npu/npu_ops_api.h index c6da986340..c862e2ac75 100644 --- a/xllm/core/kernels/npu/npu_ops_api.h +++ b/xllm/core/kernels/npu/npu_ops_api.h @@ -370,6 +370,18 @@ torch::Tensor causal_conv1d(const torch::Tensor& x, int64_t pad_slot_id, int64_t run_mode); +std::tuple causal_conv1d_qkv( + const torch::Tensor& x, + const torch::Tensor& weight, + const torch::Tensor& conv_state, + const torch::IntArrayRef query_start_loc_opt, + const torch::IntArrayRef cache_indices_opt, + const torch::IntArrayRef initial_state_mode_opt, + int64_t num_qk_heads, + int64_t num_v_heads, + int64_t head_k_dim, + int64_t head_v_dim); + void causal_conv1d_out(const torch::Tensor& output, const torch::Tensor& x, const torch::Tensor& weight, diff --git a/xllm/core/kernels/npu/xllm_ops/npu_mega_chunk_gdn.cpp b/xllm/core/kernels/npu/xllm_ops/npu_mega_chunk_gdn.cpp index 9df87489fb..b721a2c718 100644 --- a/xllm/core/kernels/npu/xllm_ops/npu_mega_chunk_gdn.cpp +++ b/xllm/core/kernels/npu/xllm_ops/npu_mega_chunk_gdn.cpp @@ -31,7 +31,8 @@ constexpr int64_t kMegaChunkSize = 128; struct MaskCache { torch::Tensor mask_lower; torch::Tensor mask_full; - torch::Tensor minus_identity; + torch::Tensor minus_identity_fp16; + torch::Tensor minus_identity_bf16; }; std::unordered_map g_mask_cache; @@ -53,10 +54,14 @@ MaskCache get_or_create_masks(const torch::Device& device) { torch::ones({kMegaChunkSize, kMegaChunkSize}, torch::TensorOptions(device).dtype(torch::kFloat32)), /*diagonal=*/0); - cache.minus_identity = + cache.minus_identity_fp16 = torch::zeros({kMegaChunkSize, kMegaChunkSize}, torch::TensorOptions(device).dtype(torch::kFloat16)); - cache.minus_identity.diagonal().fill_(-1); + cache.minus_identity_fp16.diagonal().fill_(-1); + cache.minus_identity_bf16 = + torch::zeros({kMegaChunkSize, kMegaChunkSize}, + torch::TensorOptions(device).dtype(torch::kBFloat16)); + cache.minus_identity_bf16.diagonal().fill_(-1); g_mask_cache[device_index] = cache; return cache; } @@ -83,11 +88,17 @@ std::pair npu_mega_chunk_gdn( k_normalized = npu_l2norm_last_dim(k); } - auto q_fp16 = q_normalized.to(torch::kFloat16); - auto k_fp16 = k_normalized.to(torch::kFloat16); - auto v_fp16 = v.to(torch::kFloat16); + const bool use_bf16_compute = + q_normalized.scalar_type() == torch::kBFloat16 && + k_normalized.scalar_type() == torch::kBFloat16 && + v.scalar_type() == torch::kBFloat16; + const auto compute_dtype = + use_bf16_compute ? torch::kBFloat16 : torch::kFloat16; + auto q_compute = q_normalized.to(compute_dtype); + auto k_compute = k_normalized.to(compute_dtype); + auto v_compute = v.to(compute_dtype); auto g_fp32 = g.to(torch::kFloat32); - auto beta_fp16 = beta.to(torch::kFloat16); + auto beta_compute = beta.to(compute_dtype); torch::Tensor cu_seqlens_int32; int64_t num_sequences = 0; @@ -125,6 +136,8 @@ std::pair npu_mega_chunk_gdn( const int64_t num_matrices = num_chunks * num_value_heads; auto masks = get_or_create_masks(q.device()); + const auto& minus_identity = + use_bf16_compute ? masks.minus_identity_bf16 : masks.minus_identity_fp16; const int64_t B = q.size(0); const int64_t T = q.size(1); @@ -137,41 +150,41 @@ std::pair npu_mega_chunk_gdn( ? scale.value() : std::pow(static_cast(K), -0.5f); - auto opts_fp16 = torch::TensorOptions(q.device()).dtype(torch::kFloat16); + auto opts_compute = torch::TensorOptions(q.device()).dtype(compute_dtype); auto opts_fp32 = torch::TensorOptions(q.device()).dtype(torch::kFloat32); - auto out = torch::empty({B, T, H, V}, opts_fp16); + auto out = torch::empty({B, T, H, V}, opts_compute); auto g_sum = torch::empty({B, T, H}, opts_fp32); auto g_t = torch::empty({H, T}, opts_fp32); - auto beta_t = torch::empty({H, T}, opts_fp16); - auto a = torch::zeros({B, T, H, kMegaChunkSize}, opts_fp16); - auto a_inv_f32 = torch::zeros({B, T, H, kMegaChunkSize}, opts_fp32); - auto a_inv = torch::zeros({B, T, H, kMegaChunkSize}, opts_fp16); - auto w = torch::empty({B, T, H, V}, opts_fp16); - auto u = torch::empty({B, T, H, V}, opts_fp16); - auto h = torch::zeros({num_matrices, K, V}, opts_fp16); - auto v_new = torch::empty({B, T, H, V}, opts_fp16); + auto beta_t = torch::empty({H, T}, opts_compute); + auto a = torch::empty({B, T, H, kMegaChunkSize}, opts_compute); + auto a_inv_f32 = torch::empty({B, T, H, kMegaChunkSize}, opts_fp32); + auto a_inv = torch::empty({B, T, H, kMegaChunkSize}, opts_compute); + auto w = torch::empty({B, T, H, V}, opts_compute); + auto u = torch::empty({B, T, H, V}, opts_compute); + auto h = torch::empty({num_matrices, K, V}, opts_compute); + auto v_new = torch::empty({B, T, H, V}, opts_compute); torch::Tensor initial_state_arg; bool has_initial_state = false; if (initial_state.has_value() && initial_state->defined()) { - initial_state_arg = initial_state->to(torch::kFloat16); + initial_state_arg = initial_state->to(compute_dtype); has_initial_state = true; } else { - initial_state_arg = torch::zeros({num_sequences, H, K, V}, opts_fp16); + initial_state_arg = torch::zeros({num_sequences, H, K, V}, opts_compute); } - auto final_state = torch::zeros({num_sequences * H, K, V}, opts_fp16); + auto final_state = torch::empty({num_sequences * H, K, V}, opts_compute); EXEC_NPU_CMD(aclnnMegaChunkGdn, - q_fp16, - k_fp16, - v_fp16, + q_compute, + k_compute, + v_compute, g_fp32, - beta_fp16, + beta_compute, masks.mask_lower, masks.mask_full, - masks.minus_identity, + minus_identity, cu_seqlens_int32, initial_state_arg, num_matrices, diff --git a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp index 48b58d4867..988e1df225 100644 --- a/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp +++ b/xllm/core/layers/npu_torch/qwen3_gated_delta_net_base.cpp @@ -18,6 +18,7 @@ limitations under the License. #include #include +#include "xllm/core/kernels/npu/npu_ops_api.h" #include "xllm/core/kernels/ops_api.h" #include "xllm/core/platform/npu/acl_graph_task_update_context.h" @@ -540,6 +541,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( const bool is_any_prefill = attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; torch::Tensor mixed_qkv, z, b, a; + torch::Tensor processed_q, processed_k, processed_v; int64_t batch_size = 0; int64_t seq_len = 0; @@ -578,6 +580,13 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( a = a.view({batch_size, seq_len, num_v_heads_ / tp_size_}); } + const bool fla_ssm_state_layout = use_fla_ssm_state_layout(); + const int64_t local_q_heads = num_k_heads_ / tp_size_; + const int64_t local_v_heads = num_v_heads_ / tp_size_; + const int64_t local_conv_dim = + 2 * local_q_heads * head_k_dim_ + local_v_heads * head_v_dim_; + bool used_direct_prefill_qkv = false; + torch::Tensor conv_cache = kv_cache.get_conv_cache(); torch::Tensor ssm_cache = kv_cache.get_ssm_cache(); torch::Device device = mixed_qkv.device(); @@ -598,21 +607,76 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( input_params.embedding.linear_state_ids.begin(), input_params.embedding.linear_state_ids.end()); torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv); - mixed_qkv = xllm::kernel::causal_conv1d( - conv_input, - conv_weight, - conv_cache, - std::optional(), // bias (no bias for qwen3) - torch::IntArrayRef(input_params.parallel.query_start_loc), - torch::IntArrayRef(linear_state_indices_vec), - torch::IntArrayRef(input_params.parallel.has_initial_state), - num_accepted_tokens_opt, - xllm::npu::kCausalConv1dActivationSilu, - xllm::npu::kCausalConv1dGraphPadSlotId, - xllm::npu::kCausalConv1dRunModeForward); - mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv); - mixed_qkv = mixed_qkv.transpose(1, 2); + const bool direct_qkv_model_supported = + fla_ssm_state_layout && num_k_heads_ % tp_size_ == 0 && + num_v_heads_ % tp_size_ == 0 && local_q_heads > 0 && + local_v_heads > 0 && head_k_dim_ == 128 && head_v_dim_ == 128; + const bool direct_qkv_metadata_available = + attn_metadata.q_seq_lens_vec.size() == + static_cast(batch_size) && + input_params.parallel.query_start_loc.size() == + static_cast(batch_size + 1) && + input_params.embedding.linear_state_ids.size() == + static_cast(batch_size) && + input_params.parallel.has_initial_state.size() == + static_cast(batch_size); + int64_t total_valid_tokens = 0; + bool direct_qkv_lengths_valid = direct_qkv_metadata_available; + if (direct_qkv_metadata_available) { + for (const int32_t valid_len : attn_metadata.q_seq_lens_vec) { + direct_qkv_lengths_valid = + direct_qkv_lengths_valid && valid_len >= 0 && valid_len <= seq_len; + total_valid_tokens += valid_len; + } + } + const bool direct_qkv_sequence_supported = + direct_qkv_model_supported && direct_qkv_lengths_valid && + conv_input.dim() == 2 && total_valid_tokens == conv_input.size(0); + const bool direct_qkv_shape_supported = + direct_qkv_sequence_supported && conv_input.size(1) == local_conv_dim && + conv_weight.dim() == 2 && conv_weight.size(0) == 4 && + conv_weight.size(1) == local_conv_dim && conv_cache.dim() == 3 && + conv_cache.size(1) >= 3 && conv_cache.size(2) == local_conv_dim; + const bool direct_qkv_dtype_supported = + direct_qkv_shape_supported && + conv_input.scalar_type() == torch::kBFloat16 && + conv_weight.scalar_type() == torch::kBFloat16 && + conv_cache.scalar_type() == torch::kBFloat16; + const bool use_direct_prefill_qkv = + direct_qkv_dtype_supported && conv_input.is_contiguous() && + conv_weight.is_contiguous() && conv_cache.is_contiguous(); + if (use_direct_prefill_qkv) { + std::tie(processed_q, processed_k, processed_v) = + xllm::kernel::npu::causal_conv1d_qkv( + conv_input, + conv_weight, + conv_cache, + torch::IntArrayRef(input_params.parallel.query_start_loc), + torch::IntArrayRef(linear_state_indices_vec), + torch::IntArrayRef(input_params.parallel.has_initial_state), + local_q_heads, + local_v_heads, + head_k_dim_, + head_v_dim_); + used_direct_prefill_qkv = true; + } else { + mixed_qkv = xllm::kernel::causal_conv1d( + conv_input, + conv_weight, + conv_cache, + std::optional(), // bias (no bias for qwen3) + torch::IntArrayRef(input_params.parallel.query_start_loc), + torch::IntArrayRef(linear_state_indices_vec), + torch::IntArrayRef(input_params.parallel.has_initial_state), + num_accepted_tokens_opt, + xllm::npu::kCausalConv1dActivationSilu, + xllm::npu::kCausalConv1dGraphPadSlotId, + xllm::npu::kCausalConv1dRunModeForward); + + mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv); + mixed_qkv = mixed_qkv.transpose(1, 2); + } } else { if (use_spec_verify) { CHECK(input_params.num_accepted_tokens.defined()) @@ -694,7 +758,6 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv); mixed_qkv = mixed_qkv.transpose(1, 2); } - const bool fla_ssm_state_layout = use_fla_ssm_state_layout(); const bool use_fused_sigmoid_gdn_decode = fla_ssm_state_layout && !use_spec_verify && !is_any_prefill && checkpoint_stride == 1; @@ -733,7 +796,10 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( gdn_params.threshold = 20.0f; std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params); } - auto [processed_q, processed_k, processed_v] = process_mixed_qkv(mixed_qkv); + if (!used_direct_prefill_qkv) { + std::tie(processed_q, processed_k, processed_v) = + process_mixed_qkv(mixed_qkv); + } torch::Tensor core_attn_out; torch::Tensor last_recurrent_state; // Apply chunked or recurrent gated-delta attention and update caches. @@ -796,20 +862,28 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( packed_beta.reserve(batch_size); for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx]; - packed_q.emplace_back( - processed_q[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len)); - packed_k.emplace_back( - processed_k[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len)); - packed_v.emplace_back( - processed_v[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len)); + if (!used_direct_prefill_qkv) { + packed_q.emplace_back(processed_q[batch_idx].narrow( + /*dim=*/0, /*start=*/0, valid_len)); + packed_k.emplace_back(processed_k[batch_idx].narrow( + /*dim=*/0, /*start=*/0, valid_len)); + packed_v.emplace_back(processed_v[batch_idx].narrow( + /*dim=*/0, /*start=*/0, valid_len)); + } packed_g.emplace_back( g[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len)); packed_beta.emplace_back( beta[batch_idx].narrow(/*dim=*/0, /*start=*/0, valid_len)); } - packed_processed_q = torch::cat(packed_q, 0).unsqueeze(0); - packed_processed_k = torch::cat(packed_k, 0).unsqueeze(0); - packed_processed_v = torch::cat(packed_v, 0).unsqueeze(0); + if (used_direct_prefill_qkv) { + packed_processed_q = processed_q; + packed_processed_k = processed_k; + packed_processed_v = processed_v; + } else { + packed_processed_q = torch::cat(packed_q, 0).unsqueeze(0); + packed_processed_k = torch::cat(packed_k, 0).unsqueeze(0); + packed_processed_v = torch::cat(packed_v, 0).unsqueeze(0); + } packed_g_tensor = torch::cat(packed_g, 0).unsqueeze(0); packed_beta_tensor = torch::cat(packed_beta, 0).unsqueeze(0); } @@ -842,7 +916,7 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( mega_chunk_gdn_params.cu_seqlens = attn_metadata.q_cu_seq_lens; mega_chunk_gdn_params.q_seq_lens = c10::ArrayRef( attn_metadata.q_seq_lens_vec.data(), static_cast(batch_size)); - mega_chunk_gdn_params.use_qk_l2norm_in_kernel = true; + mega_chunk_gdn_params.use_qk_l2norm_in_kernel = !used_direct_prefill_qkv; torch::Tensor packed_core_attn_out; std::tie(packed_core_attn_out, last_recurrent_state) = xllm::kernel::mega_chunk_gdn(mega_chunk_gdn_params); @@ -852,7 +926,11 @@ torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( core_attn_out = core_attn_out.to(processed_v.scalar_type()); } } else { - core_attn_out = torch::zeros_like(processed_v); + core_attn_out = + used_direct_prefill_qkv + ? torch::zeros({batch_size, seq_len, local_v_heads, head_v_dim_}, + z.options()) + : torch::zeros_like(processed_v); int64_t packed_offset = 0; for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { const int64_t valid_len = attn_metadata.q_seq_lens_vec[batch_idx];