fix(CRITICAL): align patch_ops.sh with comp 168 — keep base qwen3_5.py + upstream搬运
patch_ops.sh v2: conditional model layer deployment
搬运: moe_combine.cu, moe_compute_index.cu, fused_moe_xllm.cpp,
qwen3_gated_delta_net_base.cpp/.h, ilu_layer_fused_moe.h, ilu_layer_attention.h
This commit is contained in:
project6-dev
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
124
ex_engine/csrc/moe/fused_moe_xllm.cpp
Normal file
124
ex_engine/csrc/moe/fused_moe_xllm.cpp
Normal file
@@ -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<torch::Tensor>& 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<torch::Tensor>& fc1_expert_biases,
|
||||
const std::optional<torch::Tensor>& fc2_expert_biases,
|
||||
const std::optional<torch::Tensor>& input_sf,
|
||||
const std::optional<torch::Tensor>& swiglu_alpha,
|
||||
const std::optional<torch::Tensor>& swiglu_beta,
|
||||
const std::optional<torch::Tensor>& swiglu_limit,
|
||||
const std::optional<torch::Tensor>& 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<int64_t> 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<ffi::Module>();
|
||||
|
||||
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<ffi::Array<int64_t>>(), // TODO: support
|
||||
// auto tuning
|
||||
// profile ids
|
||||
support_pdl(),
|
||||
activation_type);
|
||||
|
||||
return result_output;
|
||||
}
|
||||
} // namespace xllm::kernel::cuda
|
||||
105
ex_engine/csrc/moe/moe_combine.cu
Executable file
105
ex_engine/csrc/moe/moe_combine.cu
Executable file
@@ -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 <c10/cuda/CUDAGuard.h>
|
||||
|
||||
#include "device_utils.cuh"
|
||||
#include "kernels/cuda/cuda_ops_api.h"
|
||||
|
||||
namespace xllm::kernel::cuda {
|
||||
|
||||
constexpr int32_t kCombineBlockSize = 256;
|
||||
|
||||
template <typename scalar_t>
|
||||
__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<float>(gemm2[flat_idx * H + h]);
|
||||
}
|
||||
output[token_id * H + h] = static_cast<scalar_t>(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<c10::Half>
|
||||
<<<N, kCombineBlockSize, 0, stream>>>(gemm2.data_ptr<c10::Half>(),
|
||||
rw.data_ptr<float>(),
|
||||
output.data_ptr<c10::Half>(),
|
||||
N,
|
||||
topk,
|
||||
H);
|
||||
} else if (dtype == torch::kBFloat16) {
|
||||
moe_combine_kernel<c10::BFloat16>
|
||||
<<<N, kCombineBlockSize, 0, stream>>>(gemm2.data_ptr<c10::BFloat16>(),
|
||||
rw.data_ptr<float>(),
|
||||
output.data_ptr<c10::BFloat16>(),
|
||||
N,
|
||||
topk,
|
||||
H);
|
||||
} else {
|
||||
moe_combine_kernel<float>
|
||||
<<<N, kCombineBlockSize, 0, stream>>>(gemm2.data_ptr<float>(),
|
||||
rw.data_ptr<float>(),
|
||||
output.data_ptr<float>(),
|
||||
N,
|
||||
topk,
|
||||
H);
|
||||
}
|
||||
|
||||
return output;
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::cuda
|
||||
155
ex_engine/csrc/moe/moe_compute_index.cu
Normal file
155
ex_engine/csrc/moe/moe_compute_index.cu
Normal file
@@ -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 <c10/cuda/CUDAGuard.h>
|
||||
|
||||
#include <cub/block/block_scan.cuh>
|
||||
|
||||
#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<int32_t, kMoeIndexBlock>;
|
||||
__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<int32_t>(flat_idx);
|
||||
src_dst[flat_idx] = pos;
|
||||
}
|
||||
|
||||
// ---- Host-side orchestrator ----
|
||||
// Returns {src_dst, dst_src, expert_sizes}
|
||||
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> 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<int32_t>(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<<<grid, kMoeIndexBlock, 0, stream>>>(
|
||||
expert_id_i32.data_ptr<int32_t>(),
|
||||
expert_sizes.data_ptr<int32_t>(),
|
||||
N,
|
||||
E);
|
||||
|
||||
// Phase 2: prefix sum (1 block)
|
||||
moe_prefix_sum_kernel<<<1, kMoeIndexBlock, 0, stream>>>(
|
||||
expert_sizes.data_ptr<int32_t>(),
|
||||
expert_offsets.data_ptr<int32_t>(),
|
||||
E,
|
||||
nullptr);
|
||||
|
||||
// Phase 3: place indices
|
||||
moe_place_indices_kernel<<<grid, kMoeIndexBlock, 0, stream>>>(
|
||||
expert_id_i32.data_ptr<int32_t>(),
|
||||
expert_offsets.data_ptr<int32_t>(),
|
||||
dst_src.data_ptr<int32_t>(),
|
||||
src_dst.data_ptr<int32_t>(),
|
||||
N,
|
||||
E);
|
||||
|
||||
return std::make_tuple(src_dst, dst_src, expert_sizes);
|
||||
}
|
||||
|
||||
} // namespace xllm::kernel::cuda
|
||||
1164
ex_engine/csrc/qwen3_gated_delta_net_base.cpp
Normal file
1164
ex_engine/csrc/qwen3_gated_delta_net_base.cpp
Normal file
File diff suppressed because it is too large
Load Diff
112
ex_engine/csrc/qwen3_gated_delta_net_base.h
Normal file
112
ex_engine/csrc/qwen3_gated_delta_net_base.h
Normal file
@@ -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 <torch/torch.h>
|
||||
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <tuple>
|
||||
#include <utility>
|
||||
|
||||
#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<torch::Tensor, torch::Tensor> project_decode_inputs(
|
||||
const torch::Tensor& hidden_states) = 0;
|
||||
virtual std::pair<torch::Tensor, torch::Tensor> 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<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>>
|
||||
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<torch::Tensor, torch::Tensor> 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<torch::Tensor, torch::Tensor, torch::Tensor> 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
|
||||
@@ -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"
|
||||
|
||||
Reference in New Issue
Block a user