diff --git a/ex_engine/csrc/ilu/ixformer.h b/ex_engine/csrc/ilu/ixformer.h index 57ce66dc..83bad88e 100644 --- a/ex_engine/csrc/ilu/ixformer.h +++ b/ex_engine/csrc/ilu/ixformer.h @@ -1,4 +1,4 @@ -/* Copyright 2025 The xLLM Authors. All Rights Reserved. +/* Copyright 2025-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. diff --git a/ex_engine/csrc/ilu/utils.h b/ex_engine/csrc/ilu/utils.h index e8af0c3c..9fd15298 100644 --- a/ex_engine/csrc/ilu/utils.h +++ b/ex_engine/csrc/ilu/utils.h @@ -1,4 +1,4 @@ -/* Copyright 2025 The xLLM Authors. All Rights Reserved. +/* Copyright 2025-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. diff --git a/ex_engine/csrc/ilu_layer_attention.h b/ex_engine/csrc/ilu_layer_attention.h index a971835f..bf4b59ba 100644 --- a/ex_engine/csrc/ilu_layer_attention.h +++ b/ex_engine/csrc/ilu_layer_attention.h @@ -1,4 +1,4 @@ -/* Copyright 2025 The xLLM Authors. All Rights Reserved. +/* Copyright 2025-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. diff --git a/ex_engine/csrc/ilu_layer_fused_moe.h b/ex_engine/csrc/ilu_layer_fused_moe.h index 3e477064..8d4e9da9 100644 --- a/ex_engine/csrc/ilu_layer_fused_moe.h +++ b/ex_engine/csrc/ilu_layer_fused_moe.h @@ -1,4 +1,4 @@ -/* Copyright 2025 The xLLM Authors. All Rights Reserved. +/* Copyright 2025-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. diff --git a/ex_engine/csrc/moe/fused_moe_xllm.cpp b/ex_engine/csrc/moe/fused_moe_xllm.cpp new file mode 100644 index 00000000..3462842a --- /dev/null +++ b/ex_engine/csrc/moe/fused_moe_xllm.cpp @@ -0,0 +1,124 @@ +/* Copyright 2025-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 "kernels/cuda/cuda_ops_api.h" +#include "kernels/cuda/utils.h" +#include "platform/device.h" +#include "platform/platform.h" + +namespace xllm::kernel::cuda { + +torch::Tensor cutlass_fused_moe( + const torch::Tensor& input, // [num_tokens, hidden] + const torch::Tensor& token_selected_experts, // [num_tokens, top_k] + const torch::Tensor& token_final_scales, // [num_tokens, top_k] + const torch::Tensor& + fc1_expert_weights, // [num_experts, inter_dim, hidden] + const torch::Tensor& + fc2_expert_weights, // [num_experts, hidden, inter_dim] + torch::ScalarType output_dtype, + const std::vector& quant_scales, + int32_t tp_size, + int32_t tp_rank, + int32_t ep_size, + int32_t ep_rank, + int32_t cluster_size, + int32_t cluster_rank, + const std::optional& fc1_expert_biases, + const std::optional& fc2_expert_biases, + const std::optional& input_sf, + const std::optional& swiglu_alpha, + const std::optional& swiglu_beta, + const std::optional& swiglu_limit, + const std::optional& output, + bool enable_alltoall, + bool use_deepseek_fp8_block_scale, + bool use_w4_group_scaling, + bool use_mxfp8_act_scaling, + bool min_latency_mode, + bool use_packed_weights, + int32_t tune_max_num_tokens, + ActivationType activation_type) { + int64_t num_rows = input.size(0); + int64_t hidden_size = fc2_expert_weights.size(1); + + if (min_latency_mode) { + num_rows *= fc2_expert_weights.size(0); + } + + std::vector output_shape = {num_rows, hidden_size}; + torch::Tensor result_output; + if (output.has_value() && output.value().defined()) { + result_output = output.value(); + } else { + torch::TensorOptions options = input.options().dtype(output_dtype); + result_output = torch::empty(output_shape, options); + } + + std::string fused_moe_uri = "fused_moe"; + if (Platform::is_support_sm90a()) { + fused_moe_uri += "_90"; + } else if (Platform::is_support_sm100a() || Platform::is_support_sm100f()) { + fused_moe_uri += "_100"; + } else if (Platform::is_support_sm120a()) { + fused_moe_uri += "_120"; + } else { + LOG(FATAL) << "FusedMoE is only supported on sm90, sm100, sm120."; + } + + bind_tvmffi_stream_to_current_torch_stream(input.device()); + + ffi::Module fused_moe_runner = + get_function(fused_moe_uri, "init")( + to_dl_data_type(input.scalar_type()), + to_dl_data_type(fc1_expert_weights.scalar_type()), + to_dl_data_type(output_dtype), + use_deepseek_fp8_block_scale, + use_w4_group_scaling, + use_mxfp8_act_scaling, + use_packed_weights) + .cast(); + + fused_moe_runner->GetFunction("run_moe").value()( + to_ffi_tensor(result_output), + to_ffi_tensor(input), + to_ffi_tensor(token_selected_experts), + to_ffi_optional_tensor(token_final_scales), + to_ffi_tensor(fc1_expert_weights), + to_ffi_optional_tensor(fc1_expert_biases), + to_ffi_tensor(fc2_expert_weights), + to_ffi_optional_tensor(fc2_expert_biases), + to_ffi_optional_array_tensors(quant_scales), + to_ffi_optional_tensor(input_sf), + to_ffi_optional_tensor(swiglu_alpha), + to_ffi_optional_tensor(swiglu_beta), + to_ffi_optional_tensor(swiglu_limit), + tp_size, + tp_rank, + ep_size, + ep_rank, + cluster_size, + cluster_rank, + enable_alltoall, + min_latency_mode, + /*profile_ids=*/ffi::Optional>(), // TODO: support + // auto tuning + // profile ids + support_pdl(), + activation_type); + + return result_output; +} +} // namespace xllm::kernel::cuda diff --git a/ex_engine/csrc/moe/moe_combine.cu b/ex_engine/csrc/moe/moe_combine.cu new file mode 100755 index 00000000..f4f21c69 --- /dev/null +++ b/ex_engine/csrc/moe/moe_combine.cu @@ -0,0 +1,105 @@ +/* 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. +==============================================================================*/ + +// Fused MoE combine kernel — reorder + weighted sum in one pass. +// Replaces: torch::zeros + index_copy_ + view + multiply + sum +// +// Algorithm per token (each block handles one token): +// 1. For each of its topk experts, read gemm2 at flat_idx directly +// (gemm2 is flat-index-ordered after scatter via index_copy_ with dst_src) +// 2. Multiply by router weight +// 3. Accumulate into output[token] +// +// Grid: num_tokens (N) blocks +// Block: HIDDEN_DIM / HIDDEN_TILE threads + +#include + +#include "device_utils.cuh" +#include "kernels/cuda/cuda_ops_api.h" + +namespace xllm::kernel::cuda { + +constexpr int32_t kCombineBlockSize = 256; + +template +__global__ void XLLM_KERNEL_ATTR(kCombineBlockSize) moe_combine_kernel( + const scalar_t* __restrict__ gemm2, // [N*topk, H] flat-index-ordered + const float* __restrict__ reduce_weight, // [N, topk] + scalar_t* __restrict__ output, // [N, H] + int64_t N, + int32_t topk, + int64_t H) { + int64_t token_id = blockIdx.x; // 0 .. N-1 + if (token_id >= N) return; + + int32_t tid = threadIdx.x; + int32_t stride = kCombineBlockSize; + + // Accumulate over topk experts for this token + for (int64_t h = tid; h < H; h += stride) { + float acc = 0.0f; + for (int32_t k = 0; k < topk; ++k) { + int64_t flat_idx = token_id * topk + k; + float w = reduce_weight[flat_idx]; + acc += w * static_cast(gemm2[flat_idx * H + h]); + } + output[token_id * H + h] = static_cast(acc); + } +} + +// ---- Host-side orchestrator ---- +torch::Tensor moe_combine_result( + const torch::Tensor& gemm2, // [N*topk, H] flat-index-ordered + const torch::Tensor& reduce_weight, // [N, topk] float or same as gemm2 + int64_t N, + int32_t topk) { + auto stream = at::cuda::getCurrentCUDAStream(); + int64_t H = gemm2.size(1); + auto dtype = gemm2.scalar_type(); + + auto output = torch::empty({N, H}, gemm2.options()); + auto rw = reduce_weight.to(gemm2.device(), torch::kFloat32).contiguous(); + + if (dtype == torch::kFloat16) { + moe_combine_kernel + <<>>(gemm2.data_ptr(), + rw.data_ptr(), + output.data_ptr(), + N, + topk, + H); + } else if (dtype == torch::kBFloat16) { + moe_combine_kernel + <<>>(gemm2.data_ptr(), + rw.data_ptr(), + output.data_ptr(), + N, + topk, + H); + } else { + moe_combine_kernel + <<>>(gemm2.data_ptr(), + rw.data_ptr(), + output.data_ptr(), + N, + topk, + H); + } + + return output; +} + +} // namespace xllm::kernel::cuda diff --git a/ex_engine/csrc/moe/moe_compute_index.cu b/ex_engine/csrc/moe/moe_compute_index.cu new file mode 100644 index 00000000..5e15a442 --- /dev/null +++ b/ex_engine/csrc/moe/moe_compute_index.cu @@ -0,0 +1,155 @@ +/* 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. +==============================================================================*/ + +// Fused MoE token index computation — 3 kernels replacing: +// torch::bincount + 2 × torch::argsort + torch::cumsum + CPU sync +// +// Phase 1 histogram: atomicAdd per-expert token counts +// Phase 2 prefix_sum: 1 block, exclusive scan → expert_offsets +// Phase 3 place_indices: atomicAdd on offsets, write dst_src + src_dst +// +// expert_sizes = per-expert token count [num_experts] (preserved) +// expert_offsets = exclusive prefix sum of counts (scratch, reused) + +#include + +#include + +#include "kernels/cuda/cuda_ops_api.h" + +namespace xllm::kernel::cuda { + +constexpr int32_t kMoeIndexBlock = 256; + +// ---- Phase 1: histogram ---- +__global__ void +#ifdef USE_DCU +__launch_bounds__(kMoeIndexBlock, 1) +#endif + moe_histogram_kernel(const int32_t* __restrict__ expert_id, + int32_t* __restrict__ expert_sizes, + int64_t num_elements, + int32_t num_experts) { + int64_t tid = int64_t(blockIdx.x) * kMoeIndexBlock + threadIdx.x; + if (tid < num_elements) { + int32_t eid = expert_id[tid]; + if (eid >= 0 && eid < num_experts) { + atomicAdd(&expert_sizes[eid], 1); + } + } +} + +// ---- Phase 2: exclusive prefix sum (1 block) ---- +// input: expert_sizes (per-expert counts) +// output: expert_offsets (exclusive scan of counts) +// total_out (total number of tokens, scalar) +__global__ void +#ifdef USE_DCU +__launch_bounds__(kMoeIndexBlock, 1) +#endif + moe_prefix_sum_kernel(const int32_t* __restrict__ expert_sizes, + int32_t* __restrict__ expert_offsets, + int32_t num_experts, + int64_t* __restrict__ total_out) { + using BlockScan = cub::BlockScan; + __shared__ typename BlockScan::TempStorage s_scan; + + int32_t val = (threadIdx.x < num_experts) ? expert_sizes[threadIdx.x] : 0; + int32_t offset; + BlockScan(s_scan).ExclusiveSum(val, offset); + __syncthreads(); + + // total = all elements sum = last thread's exclusive output + its input + int32_t total = offset + val; + + if (threadIdx.x < num_experts) { + expert_offsets[threadIdx.x] = offset; + } + if (threadIdx.x == 0 && total_out != nullptr) { + *total_out = total; + } +} + +// ---- Phase 3: place indices ---- +// atomicAdd on expert_offsets to assign a unique position within +// [start(e), start(e)+count(e)), then write both direction mappings. +__global__ void +#ifdef USE_DCU +__launch_bounds__(kMoeIndexBlock, 1) +#endif + moe_place_indices_kernel(const int32_t* __restrict__ expert_id, + int32_t* __restrict__ expert_offsets, + int32_t* __restrict__ dst_src, + int32_t* __restrict__ src_dst, + int64_t num_elements, + int32_t num_experts) { + int64_t flat_idx = int64_t(blockIdx.x) * kMoeIndexBlock + threadIdx.x; + if (flat_idx >= num_elements) return; + + int32_t eid = expert_id[flat_idx]; + if (eid < 0 || eid >= num_experts) return; + + int32_t pos = atomicAdd(&expert_offsets[eid], 1); + dst_src[pos] = static_cast(flat_idx); + src_dst[flat_idx] = pos; +} + +// ---- Host-side orchestrator ---- +// Returns {src_dst, dst_src, expert_sizes} +std::tuple moe_compute_index( + const torch::Tensor& expert_id, + int64_t num_experts) { + auto device = expert_id.device(); + auto stream = at::cuda::getCurrentCUDAStream(); + int64_t N = expert_id.numel(); + int32_t E = static_cast(num_experts); + CHECK_LE(E, kMoeIndexBlock) << "num_experts cannot exceed " << kMoeIndexBlock; + auto expert_id_i32 = expert_id.to(torch::kInt32).contiguous(); + auto opt_i32 = expert_id_i32.options(); + + auto expert_sizes = torch::zeros({num_experts}, opt_i32); + auto expert_offsets = torch::empty({num_experts}, opt_i32); + auto dst_src = torch::empty({N}, opt_i32); + auto src_dst = torch::empty({N}, opt_i32); + + int64_t grid = (N + kMoeIndexBlock - 1) / kMoeIndexBlock; + + // Phase 1: histogram + moe_histogram_kernel<<>>( + expert_id_i32.data_ptr(), + expert_sizes.data_ptr(), + N, + E); + + // Phase 2: prefix sum (1 block) + moe_prefix_sum_kernel<<<1, kMoeIndexBlock, 0, stream>>>( + expert_sizes.data_ptr(), + expert_offsets.data_ptr(), + E, + nullptr); + + // Phase 3: place indices + moe_place_indices_kernel<<>>( + expert_id_i32.data_ptr(), + expert_offsets.data_ptr(), + dst_src.data_ptr(), + src_dst.data_ptr(), + N, + E); + + return std::make_tuple(src_dst, dst_src, expert_sizes); +} + +} // namespace xllm::kernel::cuda diff --git a/ex_engine/csrc/qwen3_gated_delta_net_base.cpp b/ex_engine/csrc/qwen3_gated_delta_net_base.cpp new file mode 100644 index 00000000..7f8b4b5c --- /dev/null +++ b/ex_engine/csrc/qwen3_gated_delta_net_base.cpp @@ -0,0 +1,1164 @@ +/* Copyright 2025-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 "qwen3_gated_delta_net_base.h" + +#include +#include + +#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" + +namespace xllm { +namespace layer { + +namespace { +torch::Tensor l2norm(const torch::Tensor& x, int64_t dim, double eps = 1e-6) { + auto norm = torch::sqrt(torch::sum(torch::square(x), dim, true) + eps); + return x / norm; +} + +torch::Tensor repeat_tensor_heads(const torch::Tensor& tensor, + int64_t target_heads, + int64_t head_dim) { + const int64_t current_heads = tensor.size(head_dim); + if (current_heads == target_heads) { + return tensor; + } + CHECK_GT(current_heads, 0) << "current heads must be positive"; + CHECK_EQ(target_heads % current_heads, 0) + << "target heads must be divisible by current heads, target_heads=" + << target_heads << ", current_heads=" << current_heads; + + const int64_t repeats = target_heads / current_heads; + std::vector view_shape = tensor.sizes().vec(); + view_shape.insert(view_shape.begin() + head_dim + 1, 1); + std::vector expand_shape = view_shape; + expand_shape[head_dim + 1] = repeats; + std::vector output_shape = tensor.sizes().vec(); + output_shape[head_dim] = target_heads; + return tensor.unsqueeze(head_dim + 1) + .expand(expand_shape) + .reshape(output_shape) + .contiguous(); +} + +std::tuple torch_recurrent_gated_delta_rule( + torch::Tensor query, + torch::Tensor key, + torch::Tensor value, + torch::Tensor g, + torch::Tensor beta, + std::optional initial_state, + bool output_final_state = true, + bool use_qk_l2norm_in_kernel = true) { + auto initial_dtype = query.dtype(); + + if (use_qk_l2norm_in_kernel) { + query = l2norm(query, -1, 1e-6); + key = l2norm(key, -1, 1e-6); + } + + auto to_float32_and_transpose = [](torch::Tensor x) { + return x.transpose(1, 2).contiguous().to(torch::kFloat32); + }; + query = to_float32_and_transpose(query); + key = to_float32_and_transpose(key); + value = to_float32_and_transpose(value); + beta = to_float32_and_transpose(beta); + g = to_float32_and_transpose(g); + const int64_t value_num_heads = value.size(1); + query = repeat_tensor_heads(query, value_num_heads, 1); + key = repeat_tensor_heads(key, value_num_heads, 1); + + int64_t batch_size = key.size(0); + int64_t num_heads = key.size(1); + int64_t sequence_length = key.size(2); + int64_t k_head_dim = key.size(3); + int64_t v_head_dim = value.size(3); + + float scale_val = 1.0 / std::sqrt(static_cast(query.size(-1))); + torch::Tensor scale = torch::tensor(scale_val, query.options()); + query = query * scale; + torch::Tensor core_attn_out = torch::zeros( + {batch_size, num_heads, sequence_length, v_head_dim}, + torch::TensorOptions().dtype(torch::kFloat32).device(value.device())); + torch::Tensor last_recurrent_state; + if (!initial_state.has_value()) { + last_recurrent_state = torch::zeros( + {batch_size, num_heads, k_head_dim, v_head_dim}, + torch::TensorOptions().dtype(torch::kFloat32).device(value.device())); + } else { + last_recurrent_state = + initial_state.value().to(value.device(), torch::kFloat32); + } + + for (int64_t i = 0; i < sequence_length; ++i) { + torch::Tensor q_t = query.select(2, i); + torch::Tensor k_t = key.select(2, i); + torch::Tensor v_t = value.select(2, i); + torch::Tensor g_t = g.select(2, i).exp().unsqueeze(-1).unsqueeze(-1); + torch::Tensor beta_t = beta.select(2, i).unsqueeze(-1); + last_recurrent_state = last_recurrent_state * g_t; + torch::Tensor kv_mem = + torch::sum(last_recurrent_state * k_t.unsqueeze(-1), -2); + torch::Tensor delta = (v_t - kv_mem) * beta_t; + last_recurrent_state = + last_recurrent_state + k_t.unsqueeze(-1) * delta.unsqueeze(-2); + core_attn_out.select(2, i) = + torch::sum(last_recurrent_state * q_t.unsqueeze(-1), -2); + } + + core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype); + return std::make_tuple(core_attn_out, last_recurrent_state); +} + +std::tuple torch_chunk_gated_delta_rule( + torch::Tensor query, + torch::Tensor key, + torch::Tensor value, + torch::Tensor g, + torch::Tensor beta, + int64_t chunk_size = 64, + c10::optional initial_state = c10::nullopt, + bool output_final_state = true, + bool use_qk_l2norm_in_kernel = true) { + auto initial_dtype = query.dtype(); + if (use_qk_l2norm_in_kernel) { + query = l2norm(query, -1, 1e-6); + key = l2norm(key, -1, 1e-6); + } + auto to_float32 = [](torch::Tensor x) { + return x.transpose(1, 2).contiguous().to(torch::kFloat32); + }; + + query = to_float32(query); + key = to_float32(key); + value = to_float32(value); + beta = to_float32(beta); + g = to_float32(g); + const int64_t value_num_heads = value.size(1); + query = repeat_tensor_heads(query, value_num_heads, 1); + key = repeat_tensor_heads(key, value_num_heads, 1); + + int64_t batch_size = query.size(0); + int64_t num_heads = query.size(1); + int64_t sequence_length = query.size(2); + int64_t k_head_dim = key.size(-1); + int64_t v_head_dim = value.size(-1); + + int64_t pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size; + query = torch::nn::functional::pad( + query, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size})); + key = torch::nn::functional::pad( + key, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size})); + value = torch::nn::functional::pad( + value, torch::nn::functional::PadFuncOptions({0, 0, 0, pad_size})); + beta = torch::nn::functional::pad( + beta, torch::nn::functional::PadFuncOptions({0, pad_size})); + g = torch::nn::functional::pad( + g, torch::nn::functional::PadFuncOptions({0, pad_size})); + + int64_t total_sequence_length = sequence_length + pad_size; + float scale = 1.0 / std::sqrt(static_cast(query.size(-1))); + query = query * scale; + auto v_beta = value * beta.unsqueeze(-1); + auto k_beta = key * beta.unsqueeze(-1); + auto reshape_to_chunks = [chunk_size](torch::Tensor x) { + auto shape = x.sizes(); + std::vector new_shape = { + shape[0], shape[1], shape[2] / chunk_size, chunk_size, shape[3]}; + return x.reshape(new_shape); + }; + + query = reshape_to_chunks(query); + key = reshape_to_chunks(key); + value = reshape_to_chunks(value); + k_beta = reshape_to_chunks(k_beta); + v_beta = reshape_to_chunks(v_beta); + + auto g_shape = g.sizes(); + std::vector g_new_shape = { + g_shape[0], g_shape[1], g_shape[2] / chunk_size, chunk_size}; + g = g.reshape(g_new_shape); + auto mask = torch::triu( + torch::ones( + {chunk_size, chunk_size}, + torch::TensorOptions().dtype(torch::kBool).device(query.device())), + 0); + + g = g.cumsum(-1); + auto g_diff = g.unsqueeze(-1) - g.unsqueeze(-2); + auto decay_mask = g_diff.tril().exp().to(torch::kFloat32); + decay_mask = decay_mask.tril(); + auto attn = -(torch::matmul(k_beta, key.transpose(-1, -2)) * decay_mask) + .masked_fill(mask, 0.0); + for (int64_t i = 1; i < chunk_size; ++i) { + if (!attn.is_contiguous()) { + attn = attn.contiguous(); + } + auto row = attn.slice(-2, i, i + 1) + .slice(-1, 0, i) + .squeeze(-2) + .clone() + .contiguous(); + auto sub = attn.slice(-2, 0, i).slice(-1, 0, i).clone().contiguous(); + auto row_unsq = row.unsqueeze(-1).contiguous(); + auto row_sub_mul = (row_unsq * sub).contiguous(); + auto row_sub_sum = row_sub_mul.sum(-2).contiguous(); + auto row_final = (row + row_sub_sum).contiguous(); + attn.index_put_({torch::indexing::Ellipsis, + torch::indexing::Slice(i, i + 1), + torch::indexing::Slice(0, i)}, + row_final.unsqueeze(-2)); + } + + attn = attn + + torch::eye( + chunk_size, + torch::TensorOptions().dtype(attn.dtype()).device(attn.device())); + value = torch::matmul(attn, v_beta); + auto k_cumdecay = torch::matmul(attn, (k_beta * g.exp().unsqueeze(-1))); + torch::Tensor last_recurrent_state; + if (!initial_state.has_value()) { + last_recurrent_state = torch::zeros( + {batch_size, num_heads, k_head_dim, v_head_dim}, + torch::TensorOptions().dtype(value.dtype()).device(value.device())); + } else { + last_recurrent_state = initial_state.value().to(value); + } + auto core_attn_out = torch::zeros_like(value); + mask = torch::triu( + torch::ones( + {chunk_size, chunk_size}, + torch::TensorOptions().dtype(torch::kBool).device(query.device())), + 1); + int64_t num_chunks = total_sequence_length / chunk_size; + for (int64_t i = 0; i < num_chunks; ++i) { + auto q_i = query.select(2, i); + auto k_i = key.select(2, i); + auto v_i = value.select(2, i); + auto attn_i = + (torch::matmul(q_i, k_i.transpose(-1, -2)) * decay_mask.select(2, i)) + .masked_fill_(mask, 0.0); + auto v_prime = torch::matmul(k_cumdecay.select(2, i), last_recurrent_state); + auto v_new = v_i - v_prime; + auto attn_inter = torch::matmul(q_i * g.select(2, i).unsqueeze(-1).exp(), + last_recurrent_state); + core_attn_out.select(2, i) = attn_inter + torch::matmul(attn_i, v_new); + auto g_i_last = g.select(2, i).select(-1, -1).unsqueeze(-1); + auto g_exp_term = (g_i_last - g.select(2, i)).exp().unsqueeze(-1); + auto k_g_exp = (k_i * g_exp_term).transpose(-1, -2).contiguous(); + last_recurrent_state = last_recurrent_state * g_i_last.unsqueeze(-1).exp() + + torch::matmul(k_g_exp, v_new); + } + auto core_attn_out_shape = core_attn_out.sizes(); + std::vector reshape_shape = { + core_attn_out_shape[0], + core_attn_out_shape[1], + core_attn_out_shape[2] * core_attn_out_shape[3], + core_attn_out_shape[4]}; + core_attn_out = core_attn_out.reshape(reshape_shape); + core_attn_out = core_attn_out.slice(2, 0, sequence_length); + core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype); + return std::make_tuple(core_attn_out, last_recurrent_state); +} + +int64_t get_checkpoint_stride(const torch::Tensor& conv_cache, + const torch::Tensor& ssm_cache) { + if (!conv_cache.defined() || !ssm_cache.defined() || + conv_cache.numel() == 0 || ssm_cache.numel() == 0) { + return 1; + } + CHECK_GT(conv_cache.size(0), 0) << "conv cache must have positive batch dim"; + CHECK_EQ(ssm_cache.size(0) % conv_cache.size(0), 0) + << "ssm cache checkpoint layout mismatch, ssm_rows=" << ssm_cache.size(0) + << ", conv_rows=" << conv_cache.size(0); + return ssm_cache.size(0) / conv_cache.size(0); +} + +torch::Tensor build_linear_state_base_indices( + const torch::Tensor& logical_state_indices, + int64_t checkpoint_stride) { + if (checkpoint_stride == 1) { + return logical_state_indices; + } + return logical_state_indices * checkpoint_stride; +} + +torch::Tensor expand_sequence_tensor_to_batch(const torch::Tensor& tensor, + int64_t target_batch, + const char* tensor_name) { + CHECK(tensor.defined()) << tensor_name << " must be defined"; + CHECK_EQ(tensor.dim(), 1) << tensor_name << " must be a 1D tensor."; + const int64_t source_batch = tensor.size(0); + if (source_batch == target_batch) { + return tensor.contiguous(); + } + CHECK_GT(source_batch, 0) << tensor_name << " must not be empty."; + CHECK_EQ(target_batch % source_batch, 0) + << tensor_name << " cannot be expanded from " << source_batch << " to " + << target_batch; + const int64_t repeat_count = target_batch / source_batch; + return tensor.unsqueeze(1) + .expand({source_batch, repeat_count}) + .reshape({target_batch}) + .contiguous(); +} + +torch::Tensor run_causal_conv1d_graph_update( + const std::shared_ptr& graph_context, + const torch::Tensor& x, + const torch::Tensor& weight, + const torch::Tensor& conv_state, + const std::optional& bias, + const std::vector& query_start_loc, + const std::vector& cache_indices, + const std::vector& num_accepted_tokens, + xllm::npu::CausalConv1dGraphBranch branch) { + CHECK(graph_context != nullptr && graph_context->capturing) + << "causal_conv1d graph update can only be registered during capture"; + + c10_npu::NPUStream stream = c10_npu::getCurrentNPUStream(); + auto event = std::make_shared(ACL_EVENT_EXTERNAL); + event->block(stream); + event->reset(stream); + + torch::Tensor output; + c10_npu::graph_task_group_begin(stream); + const std::vector empty_host_args; + CHECK(!query_start_loc.empty()) + << "query_start_loc must be populated for causal_conv1d graph update"; + CHECK_EQ(query_start_loc.back(), x.size(0)) + << "query_start_loc must be padded to x.shape[0] during graph capture"; + CHECK_EQ(cache_indices.size() + 1, query_start_loc.size()) + << "cache_indices must be sequence-scoped"; + if (branch == xllm::npu::CausalConv1dGraphBranch::kSpecVerify) { + CHECK_EQ(num_accepted_tokens.size(), cache_indices.size()) + << "num_accepted_tokens must be sequence-scoped for spec verify"; + } + + output = torch::empty_like(x); + xllm::kernel::causal_conv1d_out(output, + x, + weight, + conv_state, + bias, + torch::IntArrayRef(query_start_loc), + torch::IntArrayRef(cache_indices), + torch::IntArrayRef(empty_host_args), + torch::IntArrayRef(num_accepted_tokens), + xllm::npu::kCausalConv1dActivationSilu, + xllm::npu::kCausalConv1dGraphPadSlotId, + xllm::npu::kCausalConv1dRunModeUpdate); + c10_npu::NPUTaskGroupHandle handle = c10_npu::graph_task_group_end(stream); + + xllm::npu::CausalConv1dGraphTask task; + task.output = output; + task.x = x; + task.weight = weight; + task.conv_state = conv_state; + task.bias = bias; + task.activation_mode = xllm::npu::kCausalConv1dActivationSilu; + task.pad_slot_id = xllm::npu::kCausalConv1dGraphPadSlotId; + task.run_mode = xllm::npu::kCausalConv1dRunModeUpdate; + task.branch = branch; + task.handle = handle; + task.event = std::move(event); + graph_context->causal_conv1d_tasks.emplace_back(std::move(task)); + return output; +} + +torch::Tensor run_spec_verify_gated_delta_rule( + torch::Tensor query, + torch::Tensor key, + torch::Tensor value, + torch::Tensor g, + torch::Tensor beta, + torch::Tensor& ssm_cache, + const torch::Tensor& checkpoint_indices, + const torch::Tensor& num_accepted_tokens, + const torch::Tensor& cu_seq_lens, + const std::vector& q_seq_lens_vec, + double scale) { + const auto device = value.device(); + const int64_t batch_size = value.size(0); + const int64_t seq_len = value.size(1); + const int64_t total_seq_len = batch_size * seq_len; + CHECK_EQ(cu_seq_lens.numel(), batch_size + 1) + << "GDN spec verify cu_seq_lens must be cumulative."; + CHECK_EQ(q_seq_lens_vec.size(), static_cast(batch_size)) + << "GDN spec verify q_seq_lens_vec must be per sequence."; + for (int64_t batch_idx = 0; batch_idx < batch_size; ++batch_idx) { + CHECK_EQ(q_seq_lens_vec[batch_idx], seq_len) + << "Qwen3.5 spec verify fused recurrent path expects dense " + "same-length validate tokens."; + } + + xllm::kernel::FusedRecurrentGatedDeltaRuleParams params; + params.q = query.reshape({1, total_seq_len, query.size(-2), query.size(-1)}) + .contiguous(); + params.k = + key.reshape({1, total_seq_len, key.size(-2), key.size(-1)}).contiguous(); + params.v = value.reshape({1, total_seq_len, value.size(-2), value.size(-1)}) + .contiguous(); + params.g = g.to(torch::kFloat32) + .reshape({1, total_seq_len, g.size(-1)}) + .contiguous(); + params.beta = beta.reshape({1, total_seq_len, beta.size(-1)}).contiguous(); + params.scale = static_cast(scale); + params.initial_state = ssm_cache; + params.inplace_final_state = true; + params.cu_seqlens = cu_seq_lens.to(torch::kLong).contiguous(); + params.ssm_state_indices = checkpoint_indices.contiguous(); + params.num_accepted_tokens = + num_accepted_tokens.to(device, torch::kInt32).contiguous(); + params.use_qk_l2norm_in_kernel = true; + + auto output_and_state = + xllm::kernel::fused_recurrent_gated_delta_rule(params); + return output_and_state.first.view( + {batch_size, seq_len, value.size(-2), value.size(-1)}); +} + +} // namespace + +Qwen3GatedDeltaNetBaseImpl::Qwen3GatedDeltaNetBaseImpl( + const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options) { + tp_size_ = parallel_args.tp_group_->world_size(); + rank_ = parallel_args.tp_group_->rank(); + num_k_heads_ = args.linear_num_key_heads(); + num_v_heads_ = args.linear_num_value_heads(); + head_k_dim_ = args.linear_key_head_dim(); + head_v_dim_ = args.linear_value_head_dim(); + k_size_ = num_k_heads_ * head_k_dim_; + v_size_ = num_v_heads_ * head_v_dim_; + conv_kernel_size_ = args.linear_conv_kernel_dim(); + + // Shared causal conv projection over mixed QKV states. + conv1d_ = register_module("conv1d", + ColumnParallelLinear(args.linear_conv_kernel_dim(), + k_size_ * 2 + v_size_, + /*bias=*/false, + /*gather_output=*/false, + quant_args, + parallel_args.tp_group_, + options)); + + auto opts = options.dtype(torch::kFloat32); + dt_bias_ = register_parameter("dt_bias", + torch::ones({num_v_heads_ / tp_size_}, opts), + /*requires_grad=*/false); + + A_log_ = register_parameter("A_log", + torch::empty({num_v_heads_ / tp_size_}, opts), + /*requires_grad=*/false); + + // Output projection and gated RMSNorm shared by hybrid variants. + o_proj_ = register_module("out_proj", + RowParallelLinear(v_size_, + args.hidden_size(), + /*bias=*/false, + /*input_is_parallelized=*/true, + /*if_reduce_results=*/true, + quant_args, + parallel_args.tp_group_, + options)); + + norm_ = register_module( + "norm", RmsNormGated(head_v_dim_, args.rms_norm_eps(), options)); +} + +void Qwen3GatedDeltaNetBaseImpl::load_common_state_dict( + const StateDict& state_dict) { + const int64_t rank = rank_; + const int64_t world_size = tp_size_; + const int32_t shard_tensor_count = 3; + const std::vector shard_sizes = { + k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_}; + + if (auto w = state_dict.get_tensor("conv1d.weight"); w.defined()) { + conv1d_->load_state_dict( + StateDict({{"weight", w.squeeze(1)}}, + static_cast(state_dict.prefix()) + "conv1d."), + shard_tensor_count, + shard_sizes); + conv1d_->weight().set_(conv1d_->weight().transpose(0, 1).contiguous()); + } + o_proj_->load_state_dict(state_dict.get_dict_with_prefix("out_proj.")); + if (auto w = state_dict.get_tensor("norm.weight"); w.defined()) { + norm_->load_state_dict(StateDict({{"weight", w}})); + } + LOAD_SHARDED_WEIGHT(dt_bias, 0); + LOAD_SHARDED_WEIGHT(A_log, 0); +} + +void Qwen3GatedDeltaNetBaseImpl::verify_common_loaded_weights( + const std::string& prefix) const { + CHECK(dt_bias_is_loaded_) + << "Missing required weight after all shards loaded: " << prefix + << "dt_bias"; + CHECK(A_log_is_loaded_) << "Missing required weight after all shards loaded: " + << prefix << "A_log"; +} + +std::pair +Qwen3GatedDeltaNetBaseImpl::project_padded_inputs( + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata) { + if (attn_metadata.is_prefill || attn_metadata.is_chunked_prefill) { + auto [qkvz_flat, ba_flat] = project_flat_inputs(hidden_states); + return {reshape_projected_tokens_with_pad(attn_metadata, qkvz_flat), + reshape_projected_tokens_with_pad(attn_metadata, ba_flat)}; + } + return project_decode_inputs(hidden_states); +} + +torch::Tensor Qwen3GatedDeltaNetBaseImpl::forward( + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params) { + // Early-return on dummy shards. Under dp>1, an empty shard is padded with a + // fake token by worker_impl but its GDN state tensors (kv_cache_tokens_nums, + // linear_state_ids etc.) are left undefined. This mirrors the is_dummy + // early-return in Attention::forward (npu_torch/attention.cpp). Uses + // zeros_like rather than empty_like so downstream post-norm / mlp do not + // read uninitialized data. Placed before FlashComm1 sequence gather so + // dummy shards do not enter the collective and waste bandwidth. + if (attn_metadata.is_dummy) { + return torch::zeros_like(hidden_states); + } + const FlashComm1Context* fc1_ctx = get_current_flash_comm1_context(); + torch::Tensor h = hidden_states; + if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { + h = gather_sequence(hidden_states, *fc1_ctx); + } + + // Save the gathered hidden-state size for potential padding later. + const int64_t original_num_tokens = h.size(0); + const bool use_spec_verify = input_params.is_spec_verify; + 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; + + // Qwen3.5 stores qkv, z, b, and a as separate projection weights, so it can + // use their outputs directly in every forward mode. Qwen3Next stores qkvz + // and ba as packed weights and uses the fused-split fallback below. + auto split_inputs = project_split_inputs(h, attn_metadata); + if (split_inputs.has_value()) { + std::tie(mixed_qkv, z, b, a) = split_inputs.value(); + batch_size = mixed_qkv.size(0); + seq_len = mixed_qkv.size(1); + } else { + auto [qkvz_padded, ba_padded] = project_padded_inputs(h, attn_metadata); + batch_size = qkvz_padded.size(0); + seq_len = qkvz_padded.size(1); + + torch::Tensor qkvz_flat = + qkvz_padded.view({batch_size * seq_len, qkvz_padded.size(-1)}); + torch::Tensor ba_flat = + ba_padded.view({batch_size * seq_len, ba_padded.size(-1)}); + xllm::kernel::FusedQkvzbaSplitReshapeParams fused_params; + fused_params.mixed_qkvz = qkvz_flat; + fused_params.mixed_ba = ba_flat; + fused_params.num_heads_qk = static_cast(num_k_heads_ / tp_size_); + fused_params.num_heads_v = static_cast(num_v_heads_ / tp_size_); + fused_params.head_qk = static_cast(head_k_dim_); + fused_params.head_v = static_cast(head_v_dim_); + + std::tie(mixed_qkv, z, b, a) = + xllm::kernel::fused_qkvzba_split_reshape_cat(fused_params); + + mixed_qkv = mixed_qkv.view({batch_size, seq_len, mixed_qkv.size(-1)}); + z = z.view({batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); + b = b.view({batch_size, seq_len, num_v_heads_ / tp_size_}); + 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(); + torch::Tensor conv_weight = conv1d_->weight(); + torch::Tensor logical_state_indices = + get_linear_state_indices(input_params, device); + const int64_t checkpoint_stride = + get_checkpoint_stride(conv_cache, ssm_cache); + torch::Tensor linear_state_base_indices = + build_linear_state_base_indices(logical_state_indices, checkpoint_stride); + auto graph_context = input_params.graph.acl_graph_task_update_context; + const bool register_conv1d_graph_update = + graph_context != nullptr && graph_context->capturing; + + if (!use_spec_verify && is_any_prefill) { + torch::IntArrayRef num_accepted_tokens_opt; + std::vector linear_state_indices_vec( + 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); + + 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.linear_state_validity_mask.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.linear_state_validity_mask), + 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.linear_state_validity_mask), + 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()) + << "num_accepted_tokens must be populated for Qwen3.5 spec verify"; + } + torch::Tensor conv_input = reshape_qkvz_unpad(attn_metadata, mixed_qkv); + const auto& num_accepted = use_spec_verify + ? input_params.num_accepted_tokens_host + : std::vector(); + const std::vector linear_state_indices_host( + input_params.embedding.linear_state_ids.begin(), + input_params.embedding.linear_state_ids.end()); + if (register_conv1d_graph_update) { + if (use_spec_verify) { + const auto conv1d_branch = + xllm::npu::CausalConv1dGraphBranch::kSpecVerify; + mixed_qkv = run_causal_conv1d_graph_update( + graph_context, + conv_input, + conv_weight, + conv_cache, + std::optional(), + input_params.parallel.query_start_loc, + linear_state_indices_host, + num_accepted, + conv1d_branch); + } else { + auto conv_input_2d = conv_input.dim() == 3 + ? conv_input.reshape({-1, conv_input.size(-1)}) + : conv_input; + xllm::kernel::CausalConv1dUpdateParams conv1d_params; + conv1d_params.x = conv_input_2d; + conv1d_params.conv_state = conv_cache; + conv1d_params.weight = conv_weight; + conv1d_params.conv_state_indices = logical_state_indices; + conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens; + conv1d_params.max_query_len = attn_metadata.max_query_len; + mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params); + if (conv_input.dim() == 3) { + mixed_qkv = + mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)}); + } + } + } else { + if (use_spec_verify) { + torch::Tensor output = torch::empty_like(conv_input); + xllm::kernel::causal_conv1d_out( + output, + conv_input, + conv_weight, + conv_cache, + std::optional(), + torch::IntArrayRef(input_params.parallel.query_start_loc), + torch::IntArrayRef(linear_state_indices_host), + torch::IntArrayRef(std::vector()), + torch::IntArrayRef(num_accepted), + xllm::npu::kCausalConv1dActivationSilu, + xllm::npu::kCausalConv1dGraphPadSlotId, + xllm::npu::kCausalConv1dRunModeUpdate); + mixed_qkv = output; + } else { + auto conv_input_2d = conv_input.dim() == 3 + ? conv_input.reshape({-1, conv_input.size(-1)}) + : conv_input; + xllm::kernel::CausalConv1dUpdateParams conv1d_params; + conv1d_params.x = conv_input_2d; + conv1d_params.conv_state = conv_cache; + conv1d_params.weight = conv_weight; + conv1d_params.conv_state_indices = logical_state_indices; + conv1d_params.query_start_loc = attn_metadata.q_cu_seq_lens; + conv1d_params.max_query_len = attn_metadata.max_query_len; + mixed_qkv = xllm::kernel::causal_conv1d_update(conv1d_params); + if (conv_input.dim() == 3) { + mixed_qkv = + mixed_qkv.view({conv_input.size(0), -1, mixed_qkv.size(-1)}); + } + } + } + mixed_qkv = reshape_projected_tokens_with_pad(attn_metadata, mixed_qkv); + mixed_qkv = mixed_qkv.transpose(1, 2); + } + const bool use_fused_sigmoid_gdn_decode = + fla_ssm_state_layout && !use_spec_verify && !is_any_prefill && + checkpoint_stride == 1; + torch::Tensor g; + torch::Tensor beta; + // Compute gated delta net decay and beta terms. + if (use_spec_verify || attn_metadata.is_chunked_prefill || + checkpoint_stride > 1) { + beta = torch::sigmoid(b); + torch::Tensor A_log_exp = A_log_.exp(); + torch::Tensor a_float = a.to(torch::kFloat32); + torch::Tensor a_plus_dt = a_float + dt_bias_; + torch::Tensor softplus_out = torch::nn::functional::softplus( + a_plus_dt, + torch::nn::functional::SoftplusFuncOptions().beta(1.0).threshold(20.0)); + g = -A_log_exp * softplus_out; + g = g.to(a.dtype()).contiguous(); + } else if (attn_metadata.is_prefill) { + xllm::kernel::FusedGdnGatingParams gdn_params; + gdn_params.A_log = A_log_; + gdn_params.a = a.contiguous().view({-1, a.size(-1)}); + gdn_params.b = b.contiguous().view({-1, b.size(-1)}); + gdn_params.dt_bias = dt_bias_; + gdn_params.beta = 1.0f; + gdn_params.threshold = 20.0f; + std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params); + g = g.squeeze(0).contiguous().view({batch_size, seq_len, a.size(-1)}); + beta = beta.squeeze(0).contiguous().view({batch_size, seq_len, b.size(-1)}); + } else if (!use_fused_sigmoid_gdn_decode) { + xllm::kernel::FusedGdnGatingParams gdn_params; + gdn_params.A_log = A_log_; + gdn_params.a = a.view({-1, a.size(-1)}); + gdn_params.b = b.view({-1, b.size(-1)}); + gdn_params.dt_bias = dt_bias_; + gdn_params.beta = 1.0f; + gdn_params.threshold = 20.0f; + std::tie(g, beta) = xllm::kernel::fused_gdn_gating(gdn_params); + } + 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. + if (use_spec_verify) { + torch::Tensor spec_num_accepted_tokens = expand_sequence_tensor_to_batch( + input_params.num_accepted_tokens.to(device, torch::kInt32), + batch_size, + "num_accepted_tokens"); + torch::Tensor spec_linear_state_base_indices = + expand_sequence_tensor_to_batch( + linear_state_base_indices, batch_size, "linear_state_base_indices"); + torch::Tensor step_offsets = + torch::arange(seq_len, + torch::TensorOptions() + .dtype(spec_linear_state_base_indices.dtype()) + .device(device)); + torch::Tensor checkpoint_indices = + spec_linear_state_base_indices.unsqueeze(1) + step_offsets; + double scale = 1.0 / std::sqrt(static_cast(processed_q.size(-1))); + core_attn_out = + run_spec_verify_gated_delta_rule(processed_q, + processed_k, + processed_v, + g, + beta, + ssm_cache, + checkpoint_indices, + spec_num_accepted_tokens, + attn_metadata.q_cu_seq_lens, + attn_metadata.q_seq_lens_vec, + scale); + } else if (is_any_prefill) { + CHECK_GE(attn_metadata.q_seq_lens_vec.size(), + static_cast(batch_size)) + << "q_seq_lens_vec must be populated for Qwen3.5 prefill."; + const bool use_single_prefill_pack = + batch_size == 1 && attn_metadata.q_seq_lens_vec.size() == 1 && + attn_metadata.q_seq_lens_vec[0] == seq_len; + torch::Tensor packed_processed_q; + torch::Tensor packed_processed_k; + torch::Tensor packed_processed_v; + torch::Tensor packed_g_tensor; + torch::Tensor packed_beta_tensor; + if (use_single_prefill_pack) { + packed_processed_q = processed_q; + packed_processed_k = processed_k; + packed_processed_v = processed_v; + packed_g_tensor = g; + packed_beta_tensor = beta; + } else { + std::vector packed_q; + std::vector packed_k; + std::vector packed_v; + std::vector packed_g; + std::vector packed_beta; + packed_q.reserve(batch_size); + packed_k.reserve(batch_size); + packed_v.reserve(batch_size); + packed_g.reserve(batch_size); + 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]; + 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)); + } + 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); + } + + xllm::kernel::MegaChunkGdnParams mega_chunk_gdn_params; + mega_chunk_gdn_params.q = packed_processed_q; + mega_chunk_gdn_params.k = packed_processed_k; + mega_chunk_gdn_params.v = packed_processed_v; + mega_chunk_gdn_params.g = packed_g_tensor; + mega_chunk_gdn_params.beta = packed_beta_tensor; + // Get initial state from ssm_cache for sequences with previous state + // Shape: [batch_size, num_heads, head_k_dim, head_v_dim] + torch::Tensor initial_state_tensor = + torch::index_select(ssm_cache, 0, linear_state_base_indices); + CHECK_EQ(input_params.linear_state_validity_mask.size(), + input_params.embedding.linear_state_ids.size()) + << "linear state validity mask must be sequence-scoped."; + for (size_t i = 0; i < input_params.linear_state_validity_mask.size(); + ++i) { + if (input_params.linear_state_validity_mask[i] == 0) { + initial_state_tensor.select(0, static_cast(i)).fill_(0.0); + } + } + if (!fla_ssm_state_layout && attn_metadata.is_chunked_prefill) { + initial_state_tensor = + initial_state_tensor.transpose(-1, -2).contiguous(); + } + mega_chunk_gdn_params.initial_state = initial_state_tensor; + mega_chunk_gdn_params.output_final_state = true; + 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 = !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); + if (use_single_prefill_pack) { + core_attn_out = packed_core_attn_out; + if (core_attn_out.scalar_type() != processed_v.scalar_type()) { + core_attn_out = core_attn_out.to(processed_v.scalar_type()); + } + } else { + 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]; + core_attn_out[batch_idx] + .narrow(/*dim=*/0, /*start=*/0, valid_len) + .copy_(packed_core_attn_out[0].narrow( + /*dim=*/0, packed_offset, valid_len)); + packed_offset += valid_len; + } + } + torch::Tensor state_to_store = fla_ssm_state_layout + ? last_recurrent_state + : last_recurrent_state.transpose(-1, -2); + ssm_cache.index_put_({linear_state_base_indices}, + state_to_store.to(ssm_cache.dtype())); + } else if (checkpoint_stride > 1) { + auto ssm_state = + torch::index_select(ssm_cache, 0, linear_state_base_indices); + if (!fla_ssm_state_layout) { + ssm_state = ssm_state.transpose(-1, -2); + } + ssm_state = ssm_state.contiguous(); + std::tie(core_attn_out, last_recurrent_state) = + torch_recurrent_gated_delta_rule( + processed_q, processed_k, processed_v, g, beta, ssm_state); + torch::Tensor state_to_store = fla_ssm_state_layout + ? last_recurrent_state + : last_recurrent_state.transpose(-1, -2); + ssm_cache.index_put_({linear_state_base_indices}, + state_to_store.to(ssm_cache.dtype())); + } else { + double scale = 1.0 / std::sqrt(static_cast(processed_q.size(-1))); + if (fla_ssm_state_layout) { + xllm::kernel::FusedSigmoidGatingDeltaRuleUpdateParams params; + params.A_log = A_log_.contiguous(); + params.a = a.contiguous(); + params.dt_bias = dt_bias_.contiguous(); + params.q = processed_q.contiguous(); + params.k = processed_k.contiguous(); + params.v = processed_v.contiguous(); + params.b = b.contiguous(); + params.initial_state_source = ssm_cache; + params.initial_state_indices = linear_state_base_indices.contiguous(); + params.cu_seqlens = attn_metadata.q_cu_seq_lens.contiguous(); + params.scale = static_cast(scale); + params.use_qk_l2norm_in_kernel = true; + params.softplus_beta = 1.0f; + params.softplus_threshold = 20.0f; + core_attn_out = + xllm::kernel::fused_sigmoid_gating_delta_rule_update(params); + } else { + processed_q = xllm::kernel::l2_norm(processed_q, /*eps=*/1e-6); + processed_k = xllm::kernel::l2_norm(processed_k, /*eps=*/1e-6); + auto zero = torch::zeros({1}, attn_metadata.q_seq_lens.options()); + torch::Tensor actual_seq_lengths = + torch::cat({zero, attn_metadata.q_seq_lens}, 0); + core_attn_out = xllm::kernel::recurrent_gated_delta_rule( + processed_q.reshape( + {-1, processed_q.size(-2), processed_q.size(-1)}), + processed_k.reshape( + {-1, processed_k.size(-2), processed_k.size(-1)}), + processed_v.reshape( + {-1, processed_v.size(-2), processed_v.size(-1)}), + ssm_cache, + beta.squeeze(0).contiguous(), + scale, + actual_seq_lengths, + logical_state_indices, + c10::nullopt, + g.squeeze(0).contiguous(), + c10::nullopt) + .unsqueeze(0) + .contiguous(); + } + } + auto z_reshaped = z.view({-1, z.size(-1)}); + auto core_attn_out_reshaped = + core_attn_out.view({-1, core_attn_out.size(-1)}); + auto norm_out = norm_->forward(core_attn_out_reshaped, z_reshaped); + auto z_shape_og = z.sizes().vec(); + norm_out = norm_out.view(z_shape_og); + norm_out = norm_out.view({-1, norm_out.size(2), norm_out.size(3)}); + + // Project the normalized attention output back to hidden size. + auto rearranged_norm = + norm_out.reshape({norm_out.size(0), norm_out.size(1) * norm_out.size(2)}); + rearranged_norm = reshape_qkvz_unpad(attn_metadata, rearranged_norm); + // For chunked prefill or spec verify, reshape_projected_tokens_with_pad may + // pad each batch to max_len, causing output tokens > original_num_tokens. We + // need to slice back to original_num_tokens to match the residual shape. + if (rearranged_norm.size(0) > original_num_tokens) { + // Slice excess padding tokens + rearranged_norm = + rearranged_norm.slice(0, 0, original_num_tokens).contiguous(); + } + if (fc1_ctx && is_sequence_sharded(*fc1_ctx)) { + return o_proj_->forward(rearranged_norm, + row_parallel_reduce_mode_for_fc1(*fc1_ctx)); + } + return o_proj_->forward(rearranged_norm); +} + +torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_qkvz_unpad( + const AttentionMetadata& attn_metadata, + const torch::Tensor& padded_qkvz) const { + const bool has_padded_queries = + attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; + if (!has_padded_queries) { + return padded_qkvz; + } + std::vector valid_batches; + const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty(); + int64_t bs = has_host_lens + ? static_cast(attn_metadata.q_seq_lens_vec.size()) + : attn_metadata.q_seq_lens.size(0); + valid_batches.reserve(bs); + int64_t max_len = attn_metadata.max_query_len; + const auto& ori_seq_lens = attn_metadata.q_seq_lens; + auto reshaped_qkvz = padded_qkvz.view({bs, max_len, -1}); + for (int64_t b = 0; b < bs; ++b) { + int64_t ori_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b] + : ori_seq_lens[b].template item(); + torch::Tensor valid_batch = + reshaped_qkvz[b].slice(/*dim=*/0, /*start=*/0, ori_len); + valid_batches.emplace_back(valid_batch); + } + if (valid_batches.size() == 1) { + return valid_batches[0].contiguous(); + } + return torch::cat(valid_batches, 0).contiguous(); +} + +torch::Tensor Qwen3GatedDeltaNetBaseImpl::get_linear_state_indices( + const ModelInputParams& input_params, + const torch::Device& device) const { + CHECK(!input_params.embedding.linear_state_ids.empty()) + << "linear_state_ids must be populated for gated delta net"; + if (input_params.embedding.linear_state_indices.defined()) { + auto indices = input_params.embedding.linear_state_indices; + if (indices.device() != device || indices.scalar_type() != torch::kInt) { + indices = + indices.to(torch::TensorOptions().dtype(torch::kInt).device(device), + /*non_blocking=*/true, + /*copy=*/true); + } + return indices.contiguous(); + } + return torch::tensor( + input_params.embedding.linear_state_ids, + torch::TensorOptions().dtype(torch::kInt).device(device)); +} + +torch::Tensor Qwen3GatedDeltaNetBaseImpl::reshape_projected_tokens_with_pad( + const AttentionMetadata& attn_metadata, + const torch::Tensor& projected_tokens) const { + const bool has_host_lens = !attn_metadata.q_seq_lens_vec.empty(); + int64_t bs = has_host_lens + ? static_cast(attn_metadata.q_seq_lens_vec.size()) + : attn_metadata.q_seq_lens.size(0); + int64_t max_len = attn_metadata.max_query_len; + const auto& start_loc = attn_metadata.q_seq_lens; + const bool need_padding = + attn_metadata.is_prefill || attn_metadata.is_chunked_prefill; + if (!need_padding) { + return projected_tokens.view({bs, -1, projected_tokens.size(-1)}); + } + if (has_host_lens && bs == 1 && attn_metadata.q_seq_lens_vec[0] == max_len && + projected_tokens.dim() == 2 && projected_tokens.size(0) == max_len) { + return projected_tokens.view({1, max_len, projected_tokens.size(-1)}); + } + std::vector batches; + batches.reserve(bs); + int64_t idx = 0; + for (int64_t b = 0; b < bs; ++b) { + int64_t cur_len = has_host_lens ? attn_metadata.q_seq_lens_vec[b] + : start_loc[b].template item(); + torch::Tensor batch = + projected_tokens.slice(/*dim=*/0, idx, idx + cur_len).contiguous(); + idx = idx + cur_len; + if (batch.size(0) != max_len) { + batch = batch.size(0) > max_len + ? batch.slice(/*dim=*/0, /*start=*/0, max_len).contiguous() + : torch::nn::functional::pad( + batch, + torch::nn::functional::PadFuncOptions( + {0, 0, 0, max_len - batch.size(0)})) + .contiguous(); + } + batches.emplace_back(batch); + } + auto ret = torch::stack(batches, 0).contiguous(); + return ret; +} + +std::tuple +Qwen3GatedDeltaNetBaseImpl::process_mixed_qkv(torch::Tensor& mixed_qkv) const { + mixed_qkv = mixed_qkv.transpose(1, 2); + int64_t batch_size = mixed_qkv.size(0); + int64_t seq_len = mixed_qkv.size(1); + std::vector split_sizes = { + k_size_ / tp_size_, k_size_ / tp_size_, v_size_ / tp_size_}; + auto processed_qkv = torch::split(mixed_qkv, split_sizes, 2); + auto processed_q = processed_qkv[0]; + auto processed_k = processed_qkv[1]; + auto processed_v = processed_qkv[2]; + processed_q = processed_q.view( + {batch_size, seq_len, num_k_heads_ / tp_size_, head_k_dim_}); + processed_k = processed_k.view( + {batch_size, seq_len, num_k_heads_ / tp_size_, head_k_dim_}); + processed_v = processed_v.view( + {batch_size, seq_len, num_v_heads_ / tp_size_, head_v_dim_}); + return std::make_tuple(processed_q, processed_k, processed_v); +} + +} // namespace layer +} // namespace xllm diff --git a/ex_engine/csrc/qwen3_gated_delta_net_base.h b/ex_engine/csrc/qwen3_gated_delta_net_base.h new file mode 100644 index 00000000..fdc82b4d --- /dev/null +++ b/ex_engine/csrc/qwen3_gated_delta_net_base.h @@ -0,0 +1,112 @@ +/* Copyright 2025-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. +==============================================================================*/ + +#pragma once + +#include + +#include +#include +#include +#include + +#include "attention.h" +#include "framework/kv_cache/kv_cache.h" +#include "framework/model/model_args.h" +#include "framework/parallel_state/parallel_args.h" +#include "framework/quant_args.h" +#include "framework/state_dict/state_dict.h" +#include "framework/state_dict/utils.h" +#include "layers/common/linear.h" +#include "layers/common/rms_norm_gated.h" + +namespace xllm { +namespace layer { + +class Qwen3GatedDeltaNetBaseImpl : public torch::nn::Module { + public: + Qwen3GatedDeltaNetBaseImpl() = default; + Qwen3GatedDeltaNetBaseImpl(const ModelArgs& args, + const QuantArgs& quant_args, + const ParallelArgs& parallel_args, + const torch::TensorOptions& options); + + virtual void load_state_dict(const StateDict& state_dict) = 0; + virtual void verify_loaded_weights(const std::string& prefix) const = 0; + + torch::Tensor forward(const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata, + KVCache& kv_cache, + const ModelInputParams& input_params); + + protected: + virtual std::pair project_decode_inputs( + const torch::Tensor& hidden_states) = 0; + virtual std::pair project_flat_inputs( + const torch::Tensor& hidden_states) = 0; + // Qwen3.5 overrides this to project and reshape its separate qkv/z/b/a + // weights in every forward mode. Qwen3Next keeps qkvz/ba packed and returns + // nullopt to select the fused-split fallback. + virtual std::optional< + std::tuple> + project_split_inputs(const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata) { + return std::nullopt; + } + virtual bool use_fla_ssm_state_layout() const { return false; } + + void load_common_state_dict(const StateDict& state_dict); + void verify_common_loaded_weights(const std::string& prefix) const; + + torch::Tensor get_linear_state_indices(const ModelInputParams& input_params, + const torch::Device& device) const; + + std::pair project_padded_inputs( + const torch::Tensor& hidden_states, + const AttentionMetadata& attn_metadata); + + torch::Tensor reshape_qkvz_unpad(const AttentionMetadata& attn_metadata, + const torch::Tensor& padded_qkvz) const; + + // Projection outputs are packed as [total_tokens, dim], while GDN kernels + // consume dense [batch, max_query_len, dim] tensors. Split the packed tokens + // by query length and pad each sequence before entering the kernels. + torch::Tensor reshape_projected_tokens_with_pad( + const AttentionMetadata& attn_metadata, + const torch::Tensor& projected_tokens) const; + + std::tuple process_mixed_qkv( + torch::Tensor& mixed_qkv) const; + + int64_t num_k_heads_ = 0; + int64_t num_v_heads_ = 0; + int64_t head_k_dim_ = 0; + int64_t head_v_dim_ = 0; + int64_t k_size_ = 0; + int64_t v_size_ = 0; + int64_t tp_size_ = 1; + int64_t rank_ = 0; + int32_t conv_kernel_size_ = 0; + + ColumnParallelLinear conv1d_{nullptr}; + RowParallelLinear o_proj_{nullptr}; + RmsNormGated norm_{nullptr}; + + DEFINE_WEIGHT(dt_bias); + DEFINE_WEIGHT(A_log); +}; + +} // namespace layer +} // namespace xllm diff --git a/qwen3_6_scripts/patch_ops.sh b/qwen3_6_scripts/patch_ops.sh index 47cf1e89..a7fc5714 100755 --- a/qwen3_6_scripts/patch_ops.sh +++ b/qwen3_6_scripts/patch_ops.sh @@ -1,34 +1,41 @@ #!/bin/bash # ========================================================================== -# PATCH_OPS.SH — Deploy our engine fixes + serving layer +# PATCH_OPS.SH v2 — Align with comp 168 strategy # -# BASE IMAGE HAS BUGS (proven by NaN when using base-only): -# - GDN layers produce NaN (base corex_gdn.py interface mismatch) -# - corex_fa2.py missing from model_executor/models/ -# - No multimodal support in model → engine death on image request +# COMP 168 PROOF (dockerrizhi.txt 07-23 lines 310-397): +# corex_gdn.py:56 → dlopen libcorex_gdn.so ✅ +# corex_gdn.py:228 → GDN prefill fused ✅ +# corex_gdn.py:138 → GDN decode fused ✅ +# corex_moe.py:339 → MoE prefill: expert-grouped-wmma ✅ +# corex_moe.py:249 → MoE decode fused ✅ +# corex_fa2.py:333 → FA2 packed prefill ✅ +# corex_fa2.py:507 → FA2 paged chunked prefill ✅ +# corex_fa2.py:225 → FA2 paged decode ✅ # -# COMP 168 DEPLOYED CUSTOM CODE on top of base image to fix these → 48/52 pass -# We must do the same. +# ALL 3 corex modules are IN THE BASE IMAGE and work correctly. +# Our Sub508 failed because we OVERWROTE qwen3_5.py, breaking the call chain. +# +# STRATEGY: DO NOT TOUCH model layer. Only deploy: +# 1. transformers config (Qwen3_5Config) +# 2. serving layer (protocol/serving_chat/api_server/chat_utils/tool_parser/reasoning) +# 3. ix_bridge.so (fills ixf_F.vllm_moe_topk_softmax gap if base _custom_ops hits it) +# 4. _custom_ops.py patch (make topk_softmax use ix_bridge instead of crashing) # ========================================================================== cd "$(dirname "$0")" -echo "[patch_ops] START" +echo "[patch_ops.v2] START — comp 168 aligned strategy" VLLM="" for P in /usr/local/corex/lib/python3/dist-packages/vllm \ /usr/local/corex/lib64/python3/dist-packages/vllm; do - if [ -d "$P" ]; then - VLLM="$P" - echo "[patch_ops] Found vllm at: $VLLM" - break - fi + [ -d "$P" ] && VLLM="$P" && echo "[patch_ops] Found vllm at: $VLLM" && break done [ -z "$VLLM" ] && echo "[patch_ops] ERROR: vllm not found" && exit 1 # ---- PROBE ---- echo "[probe] === Base image state ===" _QW="$VLLM/model_executor/models/qwen3_5.py" -[ -f "$_QW" ] && echo "[probe] qwen3_5.py: $(wc -c < "$_QW") bytes" || echo "[probe] qwen3_5.py: MISSING" +[ -f "$_QW" ] && echo "[probe] qwen3_5.py: $(wc -c < "$_QW") bytes, $(wc -l < "$_QW") lines" || echo "[probe] qwen3_5.py: MISSING" for m in corex_gdn.py corex_moe.py corex_fa2.py; do _F="$VLLM/model_executor/models/$m" [ -f "$_F" ] && echo "[probe] $m: $(wc -c < "$_F") bytes" || echo "[probe] $m: MISSING" @@ -36,7 +43,23 @@ done ls -la /usr/local/corex/lib64/libcorex_*.so 2>/dev/null || echo "[probe] no libcorex_*.so" echo "[probe] ===========================" -# ---- 1. Transformers config ---- +# Find secondary vllm path for mirroring +VLLM2="" +for P in /usr/local/corex/lib/python3/dist-packages/vllm \ + /usr/local/corex/lib64/python3/dist-packages/vllm; do + [ -d "$P" ] && [ "$P" != "$VLLM" ] && VLLM2="$P" && break +done + +# Helper: deploy to both vllm paths +deploy_both() { + local src="$1" dst="$2" + cp "$src" "$VLLM/$dst" 2>/dev/null || true + [ -n "$VLLM2" ] && cp "$src" "$VLLM2/$dst" 2>/dev/null || true +} + +# =========================================================== +# 1. Transformers config (Qwen3_5Config support) +# =========================================================== TMODELS="" for P in /usr/local/lib/python3.10/site-packages/transformers/models \ /usr/local/corex/lib/python3/dist-packages/transformers/models; do @@ -51,194 +74,118 @@ if [ -n "$TMODELS" ]; then echo "[patch_ops] transformers config deployed" fi -# ---- 2. Model layer — deploy OUR fixes over base image ---- -# 2a. qwen3_5.py — ALWAYS deploy ours (base image has NaN + no multimodal) -cp ./qwen3_5.py "$VLLM/model_executor/models/qwen3_5.py" && \ - echo "[patch_ops] qwen3_5.py deployed (fixes NaN + adds multimodal handling)" +# =========================================================== +# 2. MODEL LAYER — CONDITIONAL deployment +# If base has qwen3_5.py > 1000 bytes → DO NOT OVERWRITE +# This is the comp 168 strategy. +# =========================================================== +_QW_SIZE=0 +[ -f "$_QW" ] && _QW_SIZE=$(wc -c < "$_QW") -# 2b. corex modules — ALWAYS deploy ours (base interface mismatch causes fallback) -cp /workspace/ex_engine/python/corex_gdn.py "$VLLM/model_executor/models/corex_gdn.py" && \ - echo "[patch_ops] corex_gdn.py deployed (interface matches qwen3_5.py)" -cp /workspace/ex_engine/python/corex_moe.py "$VLLM/model_executor/models/corex_moe.py" && \ - echo "[patch_ops] corex_moe.py deployed" -cp /workspace/ex_engine/python/corex_fa2.py "$VLLM/model_executor/models/corex_fa2.py" && \ - echo "[patch_ops] corex_fa2.py deployed (was MISSING from base)" - -# 2c. Registry -if grep -q "Qwen3_5ForCausalLM" "$VLLM/model_executor/models/registry.py" 2>/dev/null; then - echo "[patch_ops] registry already has Qwen3_5" +if [ "$_QW_SIZE" -gt 1000 ]; then + echo "[patch_ops] *** BASE IMAGE HAS qwen3_5.py (${_QW_SIZE} bytes) — KEEPING IT ***" + echo "[patch_ops] *** This is the comp 168 strategy: don't break corex_* call chain ***" + + # Only add registry entry if missing + if ! grep -q "Qwen3_5ForCausalLM" "$VLLM/model_executor/models/registry.py" 2>/dev/null; then + cp ./registry.py "$VLLM/model_executor/models/registry.py" 2>/dev/null && \ + echo "[patch_ops] registry.py deployed (was missing Qwen3_5)" + [ -n "$VLLM2" ] && cp ./registry.py "$VLLM2/model_executor/models/registry.py" 2>/dev/null || true + fi else - cp ./registry.py "$VLLM/model_executor/models/registry.py" 2>/dev/null && \ - echo "[patch_ops] registry.py deployed" + echo "[patch_ops] *** BASE IMAGE MISSING qwen3_5.py — deploying ours ***" + deploy_both ./qwen3_5.py "model_executor/models/qwen3_5.py" + deploy_both ./registry.py "model_executor/models/registry.py" + deploy_both ./mamba_cache.py "model_executor/models/mamba_cache.py" + + # Only deploy corex modules if base doesn't have them + for m in corex_gdn.py corex_moe.py corex_fa2.py; do + if [ ! -f "$VLLM/model_executor/models/$m" ]; then + deploy_both "/workspace/ex_engine/python/$m" "model_executor/models/$m" + echo "[patch_ops] deployed $m (was MISSING)" + fi + done + + # flash_qla_sm70 (only if we deployed our qwen3_5.py) + _FLASH_SRC="/workspace/qwen3_6_scripts/flash_qla_sm70" + if [ -d "$_FLASH_SRC" ]; then + for _VPATH in "$VLLM" "$VLLM2"; do + [ -z "$_VPATH" ] && continue + cp -r "$_FLASH_SRC" "$_VPATH/model_executor/models/flash_qla_sm70" 2>/dev/null || true + done + echo "[patch_ops] flash_qla_sm70 deployed" + fi fi -# 2d. XFormers patches (head_dim=256 bypass) +# =========================================================== +# 3. SERVING LAYER — always deploy (comp 168 also used custom serving) +# =========================================================== +mkdir -p "$VLLM/entrypoints/openai/tool_parsers" 2>/dev/null || true +[ -n "$VLLM2" ] && mkdir -p "$VLLM2/entrypoints/openai/tool_parsers" 2>/dev/null || true + +deploy_both ./protocol.py "entrypoints/openai/protocol.py" +deploy_both ./cli_args.py "entrypoints/openai/cli_args.py" +deploy_both ./serving_chat.py "entrypoints/openai/serving_chat.py" +deploy_both ./api_server.py "entrypoints/openai/api_server.py" +deploy_both ./chat_utils.py "entrypoints/chat_utils.py" +deploy_both ./qwen3coder_tool_parser.py "entrypoints/openai/tool_parsers/qwen3coder_tool_parser.py" +deploy_both ./tool_parsers_init.py "entrypoints/openai/tool_parsers/__init__.py" +python3 ./patch_vllm_tool_parser.py 2>&1 || true +cp -r ./reasoning "$VLLM/" 2>/dev/null || true +[ -n "$VLLM2" ] && cp -r ./reasoning "$VLLM2/" 2>/dev/null || true +echo "[patch_ops] serving layer deployed" + +# =========================================================== +# 4. ix_bridge.so — ONLY PURPOSE: fill ixf_F.vllm_moe_topk_softmax gap +# Even comp 168 had this issue — the base _custom_ops.py tries to call +# ixf_F.vllm_moe_topk_softmax which doesn't exist. +# BUT comp 168's corex_moe.py bypasses _custom_ops entirely. +# So ix_bridge is only needed if base qwen3_5.py path hits _custom_ops. +# =========================================================== +_SITE="/usr/local/corex/lib/python3/dist-packages" +if [ -d "$_SITE" ]; then + _EX_DST="$_SITE/ex_engine" + mkdir -p "$_EX_DST/python" "$_EX_DST/build" "$_EX_DST/csrc" + cp /workspace/ex_engine/python/*.py "$_EX_DST/python/" 2>/dev/null || true + touch "$_EX_DST/__init__.py" "$_EX_DST/python/__init__.py" + + # Deploy pre-built .so + if [ -d "/workspace/ex_engine/build" ]; then + cp /workspace/ex_engine/build/*.so "$_EX_DST/build/" 2>/dev/null || true + cp /workspace/ex_engine/build/*.so "$_EX_DST/" 2>/dev/null || true + echo "[patch_ops] ex_engine .so deployed: $(ls /workspace/ex_engine/build/*.so 2>/dev/null | wc -l) files" + fi + + # C++ sources for JIT + cp /workspace/ex_engine/csrc/ix_full_bridge.cpp "$_EX_DST/csrc/" 2>/dev/null || true + cp /workspace/ex_engine/csrc/ix_moe_bridge.cpp "$_EX_DST/csrc/" 2>/dev/null || true + + echo "[patch_ops] ex_engine package deployed to $_SITE" +fi + +# =========================================================== +# 5. XFormers patches — head_dim=256 bypass for BI-V100 +# Comp 168 also had xformers patches (base uses xformers for attention) +# =========================================================== python3 ./patch_xformers_sdpa_seq.py 2>&1 || true python3 ./patch_xformers_sdpa_batch.py 2>&1 || true echo "[patch_ops] xformers patches applied" -# 2e. paged_attn.py — CRITICAL: base image uses Triton context_attention_fwd which hangs BI-V100 -cp ./paged_attn.py "$VLLM/attention/ops/paged_attn.py" && \ - echo "[patch_ops] paged_attn.py deployed (replaces Triton context_attention_fwd with PyTorch)" -[ -n "$VLLM2" ] && cp ./paged_attn.py "$VLLM2/attention/ops/paged_attn.py" 2>/dev/null || true - -# 2f. prefix_prefill.py — provides context_attention_fwd if anything still imports it -if [ -f "./prefix_prefill.py" ]; then - cp ./prefix_prefill.py "$VLLM/attention/ops/prefix_prefill.py" && \ - echo "[patch_ops] prefix_prefill.py deployed" - [ -n "$VLLM2" ] && cp ./prefix_prefill.py "$VLLM2/attention/ops/prefix_prefill.py" 2>/dev/null || true -fi - -# 2g. model_runner prefix_cache_hit fix +# =========================================================== +# 6. model_runner patch (prefix_cache_hit fix) +# =========================================================== python3 ./patch_model_runner.py 2>&1 || true +echo "[patch_ops] model_runner patched" -# 2h. mamba_cache (GDN state management) -cp ./mamba_cache.py "$VLLM/model_executor/models/mamba_cache.py" 2>/dev/null && \ - echo "[patch_ops] mamba_cache.py deployed" - -# 2i. sequence.py (token count fix) -cp ./sequence.py "$VLLM/sequence.py" 2>/dev/null && \ - echo "[patch_ops] sequence.py deployed" - -# 2j. scheduler.py (cache metrics) -cp ./scheduler.py "$VLLM/core/scheduler.py" 2>/dev/null && \ - echo "[patch_ops] scheduler.py deployed" - -# ---- 3. Serving layer ---- -mkdir -p "$VLLM/entrypoints/openai/tool_parsers" 2>/dev/null || true -cp ./qwen3coder_tool_parser.py "$VLLM/entrypoints/openai/tool_parsers/" 2>/dev/null || true -cp ./tool_parsers_init.py "$VLLM/entrypoints/openai/tool_parsers/__init__.py" 2>/dev/null || true -python3 ./patch_vllm_tool_parser.py 2>&1 || true -echo "[patch_ops] tool parser deployed" - -cp -r ./reasoning "$VLLM/" 2>/dev/null || true -echo "[patch_ops] reasoning parser deployed" - -cp ./protocol.py "$VLLM/entrypoints/openai/protocol.py" 2>/dev/null || true -cp ./cli_args.py "$VLLM/entrypoints/openai/cli_args.py" 2>/dev/null || true -cp ./serving_chat.py "$VLLM/entrypoints/openai/serving_chat.py" 2>/dev/null || true -cp ./api_server.py "$VLLM/entrypoints/openai/api_server.py" 2>/dev/null || true -cp ./chat_utils.py "$VLLM/entrypoints/chat_utils.py" 2>/dev/null || true -echo "[patch_ops] serving layer deployed" - -# ---- 4. Mirror to VLLM2 ---- -VLLM2="" -for P in /usr/local/corex/lib/python3/dist-packages/vllm \ - /usr/local/corex/lib64/python3/dist-packages/vllm; do - [ -d "$P" ] && [ "$P" != "$VLLM" ] && VLLM2="$P" && break -done -if [ -n "$VLLM2" ]; then - echo "[patch_ops] Mirroring to $VLLM2" - cp ./qwen3_5.py "$VLLM2/model_executor/models/qwen3_5.py" 2>/dev/null || true - cp /workspace/ex_engine/python/corex_gdn.py "$VLLM2/model_executor/models/corex_gdn.py" 2>/dev/null || true - cp /workspace/ex_engine/python/corex_moe.py "$VLLM2/model_executor/models/corex_moe.py" 2>/dev/null || true - cp /workspace/ex_engine/python/corex_fa2.py "$VLLM2/model_executor/models/corex_fa2.py" 2>/dev/null || true - if ! grep -q "Qwen3_5ForCausalLM" "$VLLM2/model_executor/models/registry.py" 2>/dev/null; then - cp ./registry.py "$VLLM2/model_executor/models/registry.py" 2>/dev/null || true - fi - cp ./mamba_cache.py "$VLLM2/model_executor/models/mamba_cache.py" 2>/dev/null || true - cp ./sequence.py "$VLLM2/sequence.py" 2>/dev/null || true - cp ./scheduler.py "$VLLM2/core/scheduler.py" 2>/dev/null || true - mkdir -p "$VLLM2/entrypoints/openai/tool_parsers" 2>/dev/null || true - cp ./qwen3coder_tool_parser.py "$VLLM2/entrypoints/openai/tool_parsers/" 2>/dev/null || true - cp ./tool_parsers_init.py "$VLLM2/entrypoints/openai/tool_parsers/__init__.py" 2>/dev/null || true - cp -r ./reasoning "$VLLM2/" 2>/dev/null || true - cp ./protocol.py "$VLLM2/entrypoints/openai/protocol.py" 2>/dev/null || true - cp ./cli_args.py "$VLLM2/entrypoints/openai/cli_args.py" 2>/dev/null || true - cp ./serving_chat.py "$VLLM2/entrypoints/openai/serving_chat.py" 2>/dev/null || true - cp ./api_server.py "$VLLM2/entrypoints/openai/api_server.py" 2>/dev/null || true - cp ./chat_utils.py "$VLLM2/entrypoints/chat_utils.py" 2>/dev/null || true -fi - -# ---- 5. _custom_ops.py (topk_softmax fallback) ---- -cp ./_custom_ops.py "$VLLM/_custom_ops.py" 2>/dev/null && \ - echo "[patch_ops] _custom_ops.py deployed" || true -[ -n "$VLLM2" ] && cp ./_custom_ops.py "$VLLM2/_custom_ops.py" 2>/dev/null || true - -# ---- 6. ex_engine.python subpackage (qwen3_5.py does "from ex_engine.python.ix_bridge") ---- -# The flat ex_engine package has ix_bridge.py at top level, but qwen3_5.py imports from .python subdir -_EX_PKG=$(python3 -c "import ex_engine; import os; print(os.path.dirname(ex_engine.__file__))" 2>/dev/null) -if [ -n "$_EX_PKG" ] && [ -d "$_EX_PKG" ]; then - mkdir -p "$_EX_PKG/python" - touch "$_EX_PKG/python/__init__.py" - for f in ix_bridge.py corex_moe.py corex_gdn.py corex_fa2.py; do - [ -f "$_EX_PKG/$f" ] && ln -sf "$_EX_PKG/$f" "$_EX_PKG/python/$f" - done - echo "[patch_ops] ex_engine.python subpackage linked" -fi - -# ---- 7. flash_qla_sm70 deployment to BOTH vllm paths ---- -_FLASH_SRC="/workspace/qwen3_6_scripts/flash_qla_sm70" -if [ -d "$_FLASH_SRC" ]; then - for _VPATH in "$VLLM" "$VLLM2"; do - [ -z "$_VPATH" ] && continue - _FLASH_DST="$_VPATH/model_executor/models/flash_qla_sm70" - cp -r "$_FLASH_SRC" "$_FLASH_DST" 2>/dev/null || true - done - echo "[patch_ops] flash_qla_sm70 deployed to vllm model dirs" -fi - -echo "[patch_ops] DONE" - -# ---- 8. Deploy ex_engine package + compiled .so to Python path ---- -_SITE="/usr/local/corex/lib/python3/dist-packages" -if [ -d "$_SITE" ]; then - # Deploy ex_engine as importable package - _EX_DST="$_SITE/ex_engine" - mkdir -p "$_EX_DST/python" "$_EX_DST/build" "$_EX_DST/csrc" - - # Python files - cp /workspace/ex_engine/python/*.py "$_EX_DST/python/" 2>/dev/null || true - touch "$_EX_DST/__init__.py" - touch "$_EX_DST/python/__init__.py" - - # Compiled .so files from build.sh - if [ -d "/workspace/ex_engine/build" ]; then - cp /workspace/ex_engine/build/*.so "$_EX_DST/build/" 2>/dev/null || true - # Also copy to package root for easy loading - cp /workspace/ex_engine/build/*.so "$_EX_DST/" 2>/dev/null || true - echo "[patch_ops] ex_engine .so files deployed: $(ls /workspace/ex_engine/build/*.so 2>/dev/null | wc -l) files" - fi - - # C++ sources for JIT compilation at runtime - cp /workspace/ex_engine/csrc/ix_full_bridge.cpp "$_EX_DST/csrc/" 2>/dev/null || true - cp /workspace/ex_engine/csrc/moe_topk_softmax_v3.cu "$_EX_DST/csrc/" 2>/dev/null || true - if [ -d "/workspace/ex_engine/csrc/moe_v055" ]; then - cp -r /workspace/ex_engine/csrc/moe_v055 "$_EX_DST/csrc/" 2>/dev/null || true - fi - - # Also deploy to vllm models dir for import compatibility - _EX_VLLM="$VLLM/model_executor/models/ex_engine" - mkdir -p "$_EX_VLLM/python" "$_EX_VLLM/csrc" - cp /workspace/ex_engine/python/*.py "$_EX_VLLM/python/" 2>/dev/null || true - touch "$_EX_VLLM/__init__.py" - touch "$_EX_VLLM/python/__init__.py" - cp /workspace/ex_engine/csrc/ix_full_bridge.cpp "$_EX_VLLM/csrc/" 2>/dev/null || true - if [ -d "/workspace/ex_engine/build" ]; then - cp /workspace/ex_engine/build/*.so "$_EX_VLLM/" 2>/dev/null || true - fi - - echo "[patch_ops] ex_engine deployed to $_SITE and $VLLM" -fi - -# ---- 9. Deploy precompiled MoE .so ---- -# moe_topk_softmax_v3.so (from precompile_moe_topk.py) +# =========================================================== +# 7. Deploy precompiled .so files +# =========================================================== for _SO in /workspace/ex_engine/moe_topk_softmax_v3*.so /tmp/torch_extensions/*/moe_topk_softmax_v3*.so; do - if [ -f "$_SO" ]; then - cp "$_SO" "$_SITE/" 2>/dev/null || true - echo "[patch_ops] MoE topk .so deployed: $(basename $_SO)" - break - fi + [ -f "$_SO" ] && cp "$_SO" "$_SITE/" 2>/dev/null && echo "[patch_ops] MoE topk .so: $(basename $_SO)" && break done - -# moe_v055 kernels .so (from precompile_moe_kernels.py) for _SO in /workspace/ex_engine/moe_ops_v055*.so /tmp/torch_extensions/*/moe_ops_v055*.so; do - if [ -f "$_SO" ]; then - cp "$_SO" "$_SITE/" 2>/dev/null || true - echo "[patch_ops] MoE v055 .so deployed: $(basename $_SO)" - break - fi + [ -f "$_SO" ] && cp "$_SO" "$_SITE/" 2>/dev/null && echo "[patch_ops] MoE v055 .so: $(basename $_SO)" && break done -echo "[patch_ops] FINAL: all .so and Python packages deployed" -ls -la "$_EX_DST/build/"*.so 2>/dev/null || echo "[patch_ops] WARNING: no .so in ex_engine/build/" +echo "[patch_ops.v2] DONE — comp 168 aligned" +echo "[patch_ops.v2] KEY: base qwen3_5.py $([ "$_QW_SIZE" -gt 1000 ] && echo "KEPT" || echo "REPLACED"), serving layer deployed"