[fix] baseline4 docker build move ex_engine into qwen3_6_scripts, remove COPY ex_engine from Dockerfile
This commit is contained in:
root
163
qwen3_6_scripts/ex_engine/include/ex_engine.h
Normal file
163
qwen3_6_scripts/ex_engine/include/ex_engine.h
Normal file
@@ -0,0 +1,163 @@
|
||||
// ex_engine/include/ex_engine.h — EX Engine: Algorithm Factor Replacement via dlopen
|
||||
//
|
||||
// Architecture mirrors CCCL's dispatch pattern:
|
||||
// CCCL: compute_capability → policy_selector → {threads, items, vec_size} → kernel
|
||||
// EX: hardware_id → factor_table → {op_fn_ptr, tuning_params} → dlopen .so
|
||||
//
|
||||
// The base image (BI-V100 corex SDK) has ixformer with gaps:
|
||||
// PRESENT in ixformer.functions:
|
||||
// silu_and_mul, gelu_and_mul, rms_norm, fused_add_rms_norm,
|
||||
// vllm_rotary_embedding_neox, vllm_single_query_cached_kv_attention (v1/v2),
|
||||
// vllm_cache_ops_reshape_and_cache, vllm_swap_blocks, vllm_copy_cache
|
||||
//
|
||||
// MISSING from ixformer.functions (every call falls back to slow PyTorch):
|
||||
// vllm_moe_topk_softmax — MoE routing, called 36× per token per layer
|
||||
// vllm_moe_align_block_size — MoE block alignment
|
||||
// vllm_invoke_fused_moe_kernel — MoE expert GEMM fusion
|
||||
// gelu_tanh_and_mul — activation variant
|
||||
// batched_rotary_embedding — batch RoPE
|
||||
//
|
||||
// This engine provides .so replacements for each missing factor, compiled for
|
||||
// BI-V100's SM70-class architecture using the corex clang/16 toolchain.
|
||||
|
||||
#ifndef EX_ENGINE_H
|
||||
#define EX_ENGINE_H
|
||||
|
||||
#ifdef __cplusplus
|
||||
extern "C" {
|
||||
#endif
|
||||
|
||||
#include <stdint.h>
|
||||
#include <stddef.h>
|
||||
|
||||
// ============================================================================
|
||||
// Hardware descriptor (CCCL compute_capability equivalent)
|
||||
// ============================================================================
|
||||
typedef struct {
|
||||
int sm_major; // SM version major (BI-V100 = 7)
|
||||
int sm_minor; // SM version minor (BI-V100 = 0)
|
||||
int sm_count; // Number of SMs (BI-V100 = 16)
|
||||
int max_threads_per_sm; // Max resident threads per SM
|
||||
int shared_mem_per_sm; // Shared memory per SM in bytes (49152)
|
||||
int l2_cache_size; // L2 cache size in bytes
|
||||
int memory_bus_width; // Memory bus width in bits
|
||||
float memory_bandwidth; // GB/s (BI-V100 ≈ 56 GB/s per SM)
|
||||
} ex_hardware_t;
|
||||
|
||||
// ============================================================================
|
||||
// Tuning policy (CCCL ReducePassPolicy / ScanPolicy equivalent)
|
||||
// ============================================================================
|
||||
typedef struct {
|
||||
int threads_per_block;
|
||||
int items_per_thread;
|
||||
int vec_size;
|
||||
int shared_mem_bytes; // SMEM budget (BI-V100 max 49152)
|
||||
int num_warps;
|
||||
int num_stages; // Pipeline stages (1 = no async, 2 = SW pipeline)
|
||||
} ex_tuning_t;
|
||||
|
||||
// ============================================================================
|
||||
// Factor IDs — each represents one algorithm factor to replace
|
||||
// Maps directly to the missing ixformer.functions ops
|
||||
// ============================================================================
|
||||
typedef enum {
|
||||
// MoE factors (P0 — called 36× per layer, 64 layers)
|
||||
EX_FACTOR_MOE_TOPK_SOFTMAX = 0, // topk + softmax routing
|
||||
EX_FACTOR_MOE_ALIGN_BLOCK = 1, // block alignment for scatter
|
||||
EX_FACTOR_MOE_FUSED_GEMM = 2, // fused expert GEMM
|
||||
|
||||
// Activation factors (P1)
|
||||
EX_FACTOR_GELU_TANH_MUL = 3, // gelu_tanh_and_mul
|
||||
|
||||
// RoPE factors (P1)
|
||||
EX_FACTOR_BATCHED_ROTARY = 4, // batched rotary embedding
|
||||
|
||||
// GDN factors (P0 — 4 GDN layers produce NaN without proper kernel)
|
||||
EX_FACTOR_GDN_CHUNK_FWD = 5, // GatedDeltaNet chunked prefill
|
||||
EX_FACTOR_GDN_RECURRENT = 6, // GatedDeltaNet single-step decode
|
||||
|
||||
// Cache factors (P2)
|
||||
EX_FACTOR_CACHE_APPEND = 7, // paged_attention_cache_appended
|
||||
EX_FACTOR_RESHAPE_CACHE_FLASH = 8, // reshape_and_cache_flash
|
||||
|
||||
EX_FACTOR_COUNT = 9
|
||||
} ex_factor_id_t;
|
||||
|
||||
// ============================================================================
|
||||
// Factor entry point — each .so exports this struct
|
||||
// ============================================================================
|
||||
|
||||
// Generic function pointer for the kernel dispatch
|
||||
typedef int (*ex_kernel_fn_t)(
|
||||
void* output, // output tensor data_ptr
|
||||
const void* input, // primary input tensor data_ptr
|
||||
const void* aux_inputs[], // auxiliary inputs (weights, etc.)
|
||||
int n_aux, // number of auxiliary inputs
|
||||
const int64_t dims[], // tensor dimensions
|
||||
int n_dims, // number of dimensions
|
||||
void* stream // CUDA stream
|
||||
);
|
||||
|
||||
// Each .so exports exactly one of these
|
||||
typedef struct {
|
||||
ex_factor_id_t factor_id;
|
||||
const char* name; // human-readable name
|
||||
const char* version; // semver string
|
||||
ex_tuning_t tuning; // tuned parameters for this hardware
|
||||
ex_kernel_fn_t kernel; // the replacement kernel
|
||||
ex_kernel_fn_t kernel_fallback; // PyTorch reference (NULL = no fallback)
|
||||
} ex_factor_t;
|
||||
|
||||
// Standard entry point name for dlopen: "ex_get_factor"
|
||||
typedef ex_factor_t* (*ex_get_factor_fn_t)(const ex_hardware_t* hw);
|
||||
|
||||
// ============================================================================
|
||||
// Factor registry — manages loaded .so factors
|
||||
// ============================================================================
|
||||
typedef struct {
|
||||
ex_factor_t* factors[EX_FACTOR_COUNT];
|
||||
void* handles[EX_FACTOR_COUNT]; // dlopen handles
|
||||
ex_hardware_t hardware;
|
||||
int loaded_count;
|
||||
} ex_registry_t;
|
||||
|
||||
// Initialize registry with hardware info
|
||||
int ex_registry_init(ex_registry_t* reg, const ex_hardware_t* hw);
|
||||
|
||||
// Load a single factor .so
|
||||
int ex_registry_load(ex_registry_t* reg, ex_factor_id_t id, const char* so_path);
|
||||
|
||||
// Load all .so files from a directory
|
||||
int ex_registry_load_dir(ex_registry_t* reg, const char* dir_path);
|
||||
|
||||
// Dispatch: call the loaded factor kernel, or return -1 if not loaded
|
||||
int ex_dispatch(const ex_registry_t* reg, ex_factor_id_t id,
|
||||
void* output, const void* input,
|
||||
const void* aux_inputs[], int n_aux,
|
||||
const int64_t dims[], int n_dims,
|
||||
void* stream);
|
||||
|
||||
// Cleanup
|
||||
void ex_registry_destroy(ex_registry_t* reg);
|
||||
|
||||
// ============================================================================
|
||||
// BI-V100 default hardware descriptor
|
||||
// ============================================================================
|
||||
static inline ex_hardware_t ex_bi_v100_hardware(void) {
|
||||
return (ex_hardware_t){
|
||||
.sm_major = 7,
|
||||
.sm_minor = 0,
|
||||
.sm_count = 16,
|
||||
.max_threads_per_sm = 2048,
|
||||
.shared_mem_per_sm = 49152,
|
||||
.l2_cache_size = 6 * 1024 * 1024, // 6MB
|
||||
.memory_bus_width = 4096,
|
||||
.memory_bandwidth = 900.0f // ~900 GB/s total
|
||||
};
|
||||
}
|
||||
|
||||
#ifdef __cplusplus
|
||||
}
|
||||
#endif
|
||||
|
||||
#endif // EX_ENGINE_H
|
||||
82
qwen3_6_scripts/ex_engine/include/ilu_layer_attention.h
Normal file
82
qwen3_6_scripts/ex_engine/include/ilu_layer_attention.h
Normal file
@@ -0,0 +1,82 @@
|
||||
/* Copyright 2025 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.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include <tuple>
|
||||
|
||||
#include "framework/kv_cache/kv_cache.h"
|
||||
#include "framework/model/model_input_params.h"
|
||||
#include "layers/common/attention_metadata.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
class AttentionImpl : public torch::nn::Module {
|
||||
public:
|
||||
AttentionImpl() = default;
|
||||
|
||||
AttentionImpl(int64_t num_heads,
|
||||
int64_t head_size,
|
||||
float scale,
|
||||
int64_t num_kv_heads,
|
||||
int64_t sliding_window);
|
||||
AttentionImpl(int64_t num_heads,
|
||||
int64_t head_size,
|
||||
int64_t num_kv_heads,
|
||||
int64_t v_head_dim,
|
||||
int64_t sliding_window,
|
||||
float scale,
|
||||
bool use_fused_mla_qkv,
|
||||
bool enable_lighting_indexer,
|
||||
bool enable_mla);
|
||||
|
||||
std::tuple<torch::Tensor, std::optional<torch::Tensor>> forward(
|
||||
const AttentionMetadata& attn_metadata,
|
||||
torch::Tensor& query,
|
||||
torch::Tensor& key,
|
||||
torch::Tensor& value,
|
||||
KVCache& kv_cache);
|
||||
|
||||
void prefill_forward(torch::Tensor& query,
|
||||
torch::Tensor& key,
|
||||
torch::Tensor& value,
|
||||
torch::Tensor& output,
|
||||
const torch::Tensor& k_cache,
|
||||
const std::optional<torch::Tensor>& v_cache,
|
||||
const AttentionMetadata& attn_metadata);
|
||||
|
||||
void decoder_forward(torch::Tensor& query,
|
||||
torch::Tensor& output,
|
||||
const torch::Tensor& k_cache,
|
||||
const std::optional<torch::Tensor>& v_cache,
|
||||
const AttentionMetadata& attn_metadata);
|
||||
|
||||
private:
|
||||
int64_t num_heads_;
|
||||
int64_t head_size_;
|
||||
float scale_;
|
||||
int64_t num_kv_heads_;
|
||||
int64_t v_head_dim_;
|
||||
bool use_fused_mla_qkv_;
|
||||
bool enable_lighting_indexer_;
|
||||
bool enable_mla_;
|
||||
int64_t sliding_window_;
|
||||
};
|
||||
TORCH_MODULE(Attention);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
131
qwen3_6_scripts/ex_engine/include/ilu_layer_fused_moe.h
Normal file
131
qwen3_6_scripts/ex_engine/include/ilu_layer_fused_moe.h
Normal file
@@ -0,0 +1,131 @@
|
||||
/* Copyright 2025 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.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/torch.h>
|
||||
|
||||
#include "framework/model/model_args.h"
|
||||
#include "framework/model/model_input_params.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/deep_ep.h"
|
||||
#include "layers/common/dense_mlp.h"
|
||||
#include "layers/common/fused_moe_base.h"
|
||||
#include "layers/common/linear.h"
|
||||
#include "platform/device.h"
|
||||
#include "util/tensor_helper.h"
|
||||
|
||||
namespace xllm {
|
||||
namespace layer {
|
||||
|
||||
class FusedMoEImpl : public torch::nn::Module {
|
||||
public:
|
||||
FusedMoEImpl() = default;
|
||||
FusedMoEImpl(const ModelArgs& model_args,
|
||||
const FusedMoEArgs& moe_args,
|
||||
const QuantArgs& quant_args,
|
||||
const ParallelArgs& parallel_args,
|
||||
const torch::TensorOptions& options);
|
||||
|
||||
torch::Tensor forward_experts(const torch::Tensor& hidden_states,
|
||||
const torch::Tensor& router_logits,
|
||||
bool enable_all2all_communication);
|
||||
torch::Tensor forward(const torch::Tensor& hidden_states,
|
||||
const ModelInputParams& input_params);
|
||||
void load_state_dict(const StateDict& state_dict);
|
||||
|
||||
private:
|
||||
// struct to store the selected expert info
|
||||
struct SelectedExpertInfo {
|
||||
torch::Tensor reduce_weight;
|
||||
torch::Tensor combine_idx;
|
||||
torch::Tensor token_count_slice;
|
||||
std::optional<torch::Tensor> cusum_token_count;
|
||||
std::optional<torch::Tensor> input_scale;
|
||||
};
|
||||
|
||||
// initial steps for MoE computation, select the experts for each token
|
||||
torch::Tensor select_experts(const torch::Tensor& hidden_states_2d,
|
||||
const torch::Tensor& router_logits_2d,
|
||||
SelectedExpertInfo& selected_expert_info,
|
||||
bool enable_all2all_communication);
|
||||
|
||||
private:
|
||||
int64_t num_total_experts_;
|
||||
int64_t topk_;
|
||||
int64_t num_expert_group_;
|
||||
int64_t topk_group_;
|
||||
double route_scale_;
|
||||
int64_t hidden_size_;
|
||||
int64_t n_shared_experts_;
|
||||
bool is_gated_;
|
||||
int64_t renormalize_;
|
||||
std::string hidden_act_;
|
||||
std::string scoring_func_;
|
||||
bool is_smoothquant_;
|
||||
|
||||
int64_t num_experts_per_rank_;
|
||||
int64_t start_expert_id_;
|
||||
|
||||
// Deep EP related parameters
|
||||
bool enable_deep_ep_;
|
||||
DeepEPBuffer deep_ep_buffer_;
|
||||
DeepEPParams deep_ep_params_;
|
||||
torch::Tensor dispatch_recv_token_tensor_head_;
|
||||
torch::Tensor dispatch_recv_token_tensor_tail_;
|
||||
|
||||
// steams for parallel shared experts
|
||||
std::unique_ptr<Stream> shared_stream_;
|
||||
std::unique_ptr<Stream> routed_stream_;
|
||||
xllm::Device device_;
|
||||
bool stream_initialized_ = false;
|
||||
|
||||
ReplicatedLinear gate_{nullptr};
|
||||
DenseMLP shared_experts_{nullptr};
|
||||
DeepEP deep_ep_{nullptr};
|
||||
|
||||
QuantArgs quant_args_;
|
||||
ParallelArgs parallel_args_;
|
||||
torch::TensorOptions options_;
|
||||
ProcessGroup* tp_pg_;
|
||||
|
||||
DEFINE_WEIGHT(w13);
|
||||
DEFINE_FUSED_WEIGHT(w1);
|
||||
DEFINE_FUSED_WEIGHT(w3);
|
||||
DEFINE_FUSED_WEIGHT(w2);
|
||||
DEFINE_WEIGHT(e_score_correction_bias);
|
||||
DEFINE_WEIGHT(w13_scale);
|
||||
DEFINE_FUSED_WEIGHT(w1_scale);
|
||||
DEFINE_FUSED_WEIGHT(w3_scale);
|
||||
DEFINE_FUSED_WEIGHT(w2_scale);
|
||||
DEFINE_FUSED_WEIGHT(input_smooth);
|
||||
DEFINE_FUSED_WEIGHT(act_smooth);
|
||||
|
||||
void load_e_score_correction_bias(const StateDict& state_dict);
|
||||
void load_experts(const StateDict& state_dict);
|
||||
// create the group gemm output tensor with the workspace
|
||||
torch::Tensor create_group_gemm_output(const torch::Tensor& a,
|
||||
const torch::Tensor& b,
|
||||
const torch::Tensor& group_list,
|
||||
torch::ScalarType dtype,
|
||||
torch::Tensor& workspace);
|
||||
};
|
||||
TORCH_MODULE(FusedMoE);
|
||||
|
||||
} // namespace layer
|
||||
} // namespace xllm
|
||||
153
qwen3_6_scripts/ex_engine/include/ilu_ops_api.h
Normal file
153
qwen3_6_scripts/ex_engine/include/ilu_ops_api.h
Normal file
@@ -0,0 +1,153 @@
|
||||
/* Copyright 2025 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.
|
||||
==============================================================================*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <ATen/DynamicLibrary.h>
|
||||
#include <ATen/core/dispatch/Dispatcher.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <glog/logging.h>
|
||||
#include <torch/all.h>
|
||||
|
||||
#include <optional>
|
||||
|
||||
#include "ATen/Tensor.h"
|
||||
#include "ATen/cuda/CUDAEvent.h"
|
||||
#include "c10/core/Device.h"
|
||||
#include "c10/core/DeviceGuard.h"
|
||||
#include "c10/core/GradMode.h"
|
||||
#include "c10/core/InferenceMode.h"
|
||||
#include "c10/core/MemoryFormat.h"
|
||||
#include "c10/core/ScalarType.h"
|
||||
#include "c10/core/TensorOptions.h"
|
||||
#include "c10/cuda/CUDAFunctions.h"
|
||||
#include "c10/cuda/CUDAGuard.h"
|
||||
#include "c10/cuda/CUDAStream.h"
|
||||
#include "ixformer.h"
|
||||
#include "kernels/kernels.h"
|
||||
|
||||
// #include "utils.h"
|
||||
using namespace ixformer;
|
||||
|
||||
namespace xllm::kernel::ilu {
|
||||
|
||||
void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query,
|
||||
torch::Tensor& key,
|
||||
torch::Tensor& cos_sin_cache,
|
||||
torch::Tensor& positions,
|
||||
bool interleave);
|
||||
|
||||
// act_mode only support silu, gelu, gelu_tanh
|
||||
void act_and_mul(torch::Tensor out,
|
||||
torch::Tensor input,
|
||||
const std::string& act_mode);
|
||||
|
||||
void reshape_paged_cache(
|
||||
torch::Tensor& key, // (num_tokens, num_heads, head_size)
|
||||
std::optional<torch::Tensor>& value, // (num_tokens, num_heads, head_size)
|
||||
torch::Tensor& key_cache, // (num_blocks, num_heads, block_size, head_size)
|
||||
std::optional<torch::Tensor>&
|
||||
value_cache, // (num_blocks, num_heads, block_size, head_size)
|
||||
torch::Tensor& slot_mapping); //(num_tokens)
|
||||
|
||||
void batch_prefill(torch::Tensor& query,
|
||||
const torch::Tensor& key,
|
||||
const std::optional<torch::Tensor>& value,
|
||||
torch::Tensor& output,
|
||||
std::optional<torch::Tensor>& output_lse,
|
||||
const std::optional<torch::Tensor>& q_cu_seq_lens,
|
||||
const std::optional<torch::Tensor>& kv_cu_seq_lens,
|
||||
const std::optional<torch::Tensor>& alibi_slope,
|
||||
const std::optional<torch::Tensor>& attn_bias,
|
||||
const std::optional<torch::Tensor>& q_quant_scale,
|
||||
const std::optional<torch::Tensor>& k_quant_scale,
|
||||
const std::optional<torch::Tensor>& v_quant_scale,
|
||||
const torch::Tensor& block_tables,
|
||||
int64_t max_query_len,
|
||||
int64_t max_seq_len,
|
||||
float scale,
|
||||
bool is_causal,
|
||||
int64_t window_size_left,
|
||||
int64_t window_size_right,
|
||||
const std::string& compute_dtype,
|
||||
bool return_lse);
|
||||
|
||||
void batch_decode(torch::Tensor& query,
|
||||
const torch::Tensor& k_cache,
|
||||
torch::Tensor& output,
|
||||
const torch::Tensor& block_table,
|
||||
const torch::Tensor& seq_lens,
|
||||
const std::optional<torch::Tensor>& v_cache,
|
||||
std::optional<torch::Tensor>& output_lse,
|
||||
const std::optional<torch::Tensor>& q_quant_scale,
|
||||
const std::optional<torch::Tensor>& k_cache_quant_scale,
|
||||
const std::optional<torch::Tensor>& v_cache_quant_scale,
|
||||
const std::optional<torch::Tensor>& out_quant_scale,
|
||||
const std::optional<torch::Tensor>& alibi_slope,
|
||||
const std::optional<torch::Tensor>& mask,
|
||||
const std::string& compute_dtype,
|
||||
int64_t max_seq_len,
|
||||
int64_t window_size_left,
|
||||
int64_t window_size_right,
|
||||
float scale,
|
||||
bool return_lse,
|
||||
bool is_causal,
|
||||
int64_t kv_cache_quant_bit_size);
|
||||
|
||||
void residual_layer_norm(torch::Tensor& input,
|
||||
torch::Tensor& output,
|
||||
std::optional<torch::Tensor>& residual,
|
||||
torch::Tensor& weight,
|
||||
std::optional<torch::Tensor>& bias,
|
||||
std::optional<torch::Tensor>& residual_out,
|
||||
double eps);
|
||||
|
||||
void rms_norm(torch::Tensor& output,
|
||||
torch::Tensor& input,
|
||||
torch::Tensor& weight,
|
||||
double eps);
|
||||
|
||||
torch::Tensor matmul(torch::Tensor a,
|
||||
torch::Tensor b,
|
||||
std::optional<torch::Tensor> bias);
|
||||
|
||||
std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
|
||||
const torch::Tensor& input,
|
||||
int64_t topk,
|
||||
int64_t num_expert_group,
|
||||
int64_t topk_group,
|
||||
bool normalize,
|
||||
const std::optional<torch::Tensor>& mask,
|
||||
const std::string& normed_by,
|
||||
const std::string& scoring_func,
|
||||
double route_scale,
|
||||
const std::optional<torch::Tensor>& e_score_correction_bias);
|
||||
|
||||
std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
|
||||
int64_t expert_num);
|
||||
|
||||
torch::Tensor moe_expand_input(const torch::Tensor& input,
|
||||
const torch::Tensor& gather_index,
|
||||
const torch::Tensor& combine_idx,
|
||||
int64_t topk);
|
||||
|
||||
torch::Tensor group_gemm(torch::Tensor& input,
|
||||
torch::Tensor& weight,
|
||||
torch::Tensor& tokens_per_experts,
|
||||
const std::optional<torch::Tensor>& dst_to_src,
|
||||
torch::Tensor& output);
|
||||
|
||||
torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight);
|
||||
} // namespace xllm::kernel::ilu
|
||||
63
qwen3_6_scripts/ex_engine/include/ilu_utils.h
Normal file
63
qwen3_6_scripts/ex_engine/include/ilu_utils.h
Normal file
@@ -0,0 +1,63 @@
|
||||
/* Copyright 2025 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.
|
||||
==============================================================================*/
|
||||
#pragma once
|
||||
namespace xllm::kernel::ilu {
|
||||
#undef check_tensor_contiguous
|
||||
#define check_tensor_contiguous(x, type) \
|
||||
TORCH_CHECK(x.scalar_type() == type); \
|
||||
TORCH_CHECK(x.is_cuda()); \
|
||||
TORCH_CHECK(x.is_contiguous());
|
||||
|
||||
#undef check_tensor_half_bf_float
|
||||
#define check_tensor_half_bf_float(x) \
|
||||
TORCH_CHECK(x.scalar_type() == at::ScalarType::Half || \
|
||||
x.scalar_type() == at::ScalarType::Float || \
|
||||
x.scalar_type() == at::ScalarType::BFloat16); \
|
||||
TORCH_CHECK(x.is_cuda());
|
||||
|
||||
// from torchCheckMsgImpl
|
||||
inline const char* ixformer_check_msg_impl(const char* msg) { return msg; }
|
||||
// // If there is just 1 user-provided C-string argument, use it.
|
||||
|
||||
#define IXFORMER_CHECK_MSG(cond, type, ...) \
|
||||
(ixformer_check_msg_impl( \
|
||||
"Expected " #cond \
|
||||
" to be true, but got false. " \
|
||||
"(Could this error message be improved? If so, " \
|
||||
"please report an enhancement request to ixformer.)", \
|
||||
##__VA_ARGS__))
|
||||
|
||||
#define IXFORMER_CHECK(cond, ...) \
|
||||
{ \
|
||||
if (!(cond)) { \
|
||||
std::cerr << __FILE__ << " (" << __LINE__ << ")" \
|
||||
<< "-" << __FUNCTION__ << " : " \
|
||||
<< IXFORMER_CHECK_MSG(cond, "", ##__VA_ARGS__) << std::endl; \
|
||||
throw std::runtime_error("IXFORMER_CHECK ERROR"); \
|
||||
} \
|
||||
}
|
||||
|
||||
#undef CUINFER_CHECK
|
||||
#define CUINFER_CHECK(func) \
|
||||
do { \
|
||||
cuinferStatus_t status = (func); \
|
||||
if (status != CUINFER_STATUS_SUCCESS) { \
|
||||
std::cerr << "Error in file " << __FILE__ << " on line " << __LINE__ \
|
||||
<< ": " << cuinferGetErrorString(status) << std::endl; \
|
||||
throw std::runtime_error("CUINFER_CHECK ERROR"); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
} // namespace xllm::kernel::ilu
|
||||
147
qwen3_6_scripts/ex_engine/include/ixformer.h
Normal file
147
qwen3_6_scripts/ex_engine/include/ixformer.h
Normal file
@@ -0,0 +1,147 @@
|
||||
/* Copyright 2025 The xLLM Authors. All Rights Reserved.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
==============================================================================*/
|
||||
#include <torch/all.h>
|
||||
|
||||
#include "ATen/Tensor.h"
|
||||
#include "utils.h"
|
||||
|
||||
namespace ixformer::infer {
|
||||
torch::Tensor ixinfer_flash_attn_unpad_with_block_tables(
|
||||
torch::Tensor& query,
|
||||
torch::Tensor& key_cache,
|
||||
torch::Tensor& value_cache,
|
||||
torch::Tensor& out,
|
||||
torch::Tensor& block_tables,
|
||||
torch::Tensor& cu_seq_q,
|
||||
torch::Tensor& cu_seq_k,
|
||||
int64_t max_seq_q,
|
||||
int64_t max_seq_k,
|
||||
bool is_causal,
|
||||
int64_t window_left,
|
||||
int64_t window_right,
|
||||
double scale,
|
||||
double softcap,
|
||||
bool sqrt_alibi,
|
||||
const std::optional<torch::Tensor>& alibi_slopes,
|
||||
const std::optional<torch::Tensor>& sinks,
|
||||
std::optional<torch::Tensor>& lse);
|
||||
|
||||
void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
|
||||
|
||||
torch::Tensor xllm_paged_attention(
|
||||
torch::Tensor& out,
|
||||
torch::Tensor& query,
|
||||
torch::Tensor& key_cache,
|
||||
torch::Tensor& value_cache,
|
||||
int64_t num_kv_heads,
|
||||
double scale,
|
||||
torch::Tensor& block_tables,
|
||||
torch::Tensor& context_lens,
|
||||
int64_t block_size,
|
||||
int64_t max_context_len,
|
||||
const std::optional<torch::Tensor>& alibi_slopes,
|
||||
bool causal,
|
||||
int32_t window_left,
|
||||
int32_t window_right,
|
||||
double softcap,
|
||||
bool enable_cuda_graph,
|
||||
bool use_sqrt_alibi,
|
||||
const std::optional<torch::Tensor>& sinks);
|
||||
|
||||
torch::Tensor ixformer_linear(torch::Tensor& input,
|
||||
torch::Tensor& weight,
|
||||
int64_t act_type,
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const std::optional<torch::Tensor>& out,
|
||||
const std::optional<bool> persistent);
|
||||
|
||||
torch::Tensor ixformer_linear_ex(torch::Tensor& input,
|
||||
torch::Tensor& weight,
|
||||
const c10::optional<torch::Tensor>& bias,
|
||||
const c10::optional<torch::Tensor>& out);
|
||||
|
||||
void xllm_reshape_and_cache(torch::Tensor& key,
|
||||
torch::Tensor& value,
|
||||
torch::Tensor& key_cache,
|
||||
torch::Tensor& value_cache,
|
||||
torch::Tensor& slot_mapping,
|
||||
int64_t key_token_stride,
|
||||
int64_t value_token_stride);
|
||||
|
||||
void xllm_rotary_embedding(torch::Tensor& positions,
|
||||
torch::Tensor& query,
|
||||
torch::Tensor& key,
|
||||
int64_t head_size,
|
||||
torch::Tensor& cos_sin_cache,
|
||||
bool is_neox);
|
||||
|
||||
void residual_rms_norm(torch::Tensor& input,
|
||||
torch::Tensor& residual,
|
||||
torch::Tensor& weight,
|
||||
torch::Tensor& output,
|
||||
torch::Tensor& residual_output,
|
||||
const std::optional<torch::Tensor>& fused_bias,
|
||||
double alpha,
|
||||
double eps,
|
||||
bool is_post);
|
||||
|
||||
void rms_norm(torch::Tensor& input,
|
||||
torch::Tensor& weight,
|
||||
torch::Tensor& output,
|
||||
const std::optional<torch::Tensor>& fused_bias,
|
||||
double eps);
|
||||
|
||||
void topk_softmax(torch::Tensor& topk_weights,
|
||||
torch::Tensor& topk_indices,
|
||||
torch::Tensor& token_expert_indices,
|
||||
torch::Tensor& gating_output,
|
||||
bool renormalize);
|
||||
|
||||
void moe_compute_token_index_api(
|
||||
torch::Tensor& topk_ids,
|
||||
torch::Tensor& src_dst,
|
||||
torch::Tensor& dst_src,
|
||||
torch::Tensor& expert_sizes_gpu,
|
||||
const c10::optional<torch::Tensor>& expert_mask,
|
||||
const c10::optional<torch::Tensor>& expert_sizes_cpu,
|
||||
const c10::optional<torch::Tensor>& expand_tokens_gpu,
|
||||
int64_t start_expert_id,
|
||||
int64_t end_expert_id,
|
||||
int64_t num_experts);
|
||||
|
||||
void moe_expand_input(torch::Tensor outputs,
|
||||
torch::Tensor inputs,
|
||||
torch::Tensor dst_to_src,
|
||||
const c10::optional<torch::Tensor>& src_to_dst,
|
||||
int64_t dst_tokens,
|
||||
int64_t expand_factor);
|
||||
|
||||
void moe_w16a16_group_gemm(torch::Tensor output,
|
||||
torch::Tensor inputs,
|
||||
torch::Tensor weights,
|
||||
torch::Tensor tokens_per_experts,
|
||||
const c10::optional<torch::Tensor>& dst_to_src,
|
||||
const c10::optional<torch::Tensor>& bias,
|
||||
std::string format,
|
||||
int64_t persistent,
|
||||
int64_t output_n);
|
||||
|
||||
void moe_output_reduce_sum(torch::Tensor outputs,
|
||||
torch::Tensor inputs,
|
||||
const c10::optional<torch::Tensor>& mul_weight,
|
||||
const c10::optional<torch::Tensor>& mask,
|
||||
const c10::optional<torch::Tensor>& extra_residual,
|
||||
double scaling_factor);
|
||||
} // namespace ixformer::infer
|
||||
Reference in New Issue
Block a user