dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[perf] xllm_fused_qknorm_rope.so compiled+wired, xllm_cache

Browse Source
This commit is contained in:
root
2026-09-01 06:36:44 +00:00
parent 0d88ac4b62
commit b168ff9e2d
19 changed files with 3749 additions and 880 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
qwen3_6_scripts/ex_engine/xllm_kernels/cuda/headers/cuda_ops_api.h
View File

@@ -17,7 +17,8 @@ limitations under the License.
#include <ATen/DynamicLibrary.h>
#include <ATen/core/dispatch/Dispatcher.h>
#include <glog/logging.h>
// PRD build fix: replace glog with torch logging (no glog-dev on BI-V100)
#include <c10/util/Logging.h>
#include <optional>
#include <tuple>
@@ -54,7 +55,7 @@ void block_copy(torch::Tensor key_cache_ptrs,
torch::Tensor cum_sum,
int64_t numel_per_block,
torch::ScalarType cache_dtype);
#if !defined(USE_DCU)
#if !defined(USE_DCU) && !defined(__ILUVATAR__)
void batch_prefill(const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
@@ -142,7 +143,7 @@ void batch_decode(const std::string& uri,
std::optional<torch::Tensor>& output_lse,
bool use_tensor_core,
std::optional<torch::Tensor> qo_indptr = std::nullopt);
#endif // !defined(USE_DCU)
#endif // !defined(USE_DCU) && !defined(__ILUVATAR__)
void rms_norm(torch::Tensor output,
torch::Tensor input,
torch::Tensor weight,
@@ -303,4 +304,4 @@ torch::Tensor moe_combine_result(const torch::Tensor& gemm2,
int64_t N,
int32_t topk);
} // namespace xllm::kernel::cuda
} // namespace xllm::kernel::cuda

2
qwen3_6_scripts/ex_engine/xllm_kernels/cuda/headers/type_convert.cuh
View File

@@ -57,7 +57,7 @@ class _typeConvert<float> {
};
#if defined(USE_DCU) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)) || \
defined(USE_MACA)
defined(USE_MACA) || defined(__ILUVATAR__)
// CUDA < 12.0 runs into issues with packed type conversion
template <>
class _typeConvert<c10::Half> {

13
qwen3_6_scripts/ex_engine/xllm_kernels/cuda/headers/utils.h
View File

@@ -21,9 +21,10 @@ limitations under the License.
#else
#include <c10/cuda/CUDAGuard.h>
#endif
#include <glog/logging.h>
// PRD build fix: replace glog with c10 logging (no glog-dev on BI-V100)
#include <c10/util/Logging.h>
#include <torch/torch.h>
#if !defined(USE_DCU)
#if !defined(USE_DCU) && !defined(__ILUVATAR__)
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/extra/c_env_api.h>
@@ -46,7 +47,7 @@ limitations under the License.
#define HOST_INLINE inline
#endif
#if !defined(USE_DCU)
#if !defined(USE_DCU) && !defined(__ILUVATAR__)
namespace ffi = tvm::ffi;
#endif
@@ -124,7 +125,7 @@ std::string get_batch_decode_uri(torch::ScalarType dtype_q,
std::tuple<torch::Tensor, double> split_scale_param(const torch::Tensor& scale);
#if !defined(USE_DCU)
#if !defined(USE_DCU) && !defined(__ILUVATAR__)
DLDataType to_dl_data_type(torch::ScalarType scalar_type);
// below are tvm-ffi related functions
@@ -159,5 +160,5 @@ inline void bind_tvmffi_stream_to_current_torch_stream(
<< " dev=" << device.index();
}
}
#endif // !defined(USE_DCU)
} // namespace xllm::kernel::cuda
#endif // !defined(USE_DCU) && !defined(__ILUVATAR__)
} // namespace xllm::kernel::cuda

117
qwen3_6_scripts/ex_engine/xllm_kernels/cuda/moe/fused_moe.cpp
View File

@@ -13,10 +13,10 @@ 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"
#include "cuda_ops_api.h"
#include "utils.h"
namespace xllm::kernel::cuda {
@@ -51,74 +51,73 @@ torch::Tensor cutlass_fused_moe(
bool use_packed_weights,
int32_t tune_max_num_tokens,
ActivationType activation_type) {
int64_t num_rows = input.size(0);
int64_t num_tokens = input.size(0);
int64_t hidden_size = fc2_expert_weights.size(1);
int64_t inter_dim = fc1_expert_weights.size(1);
int64_t top_k = token_selected_experts.size(1);
int64_t num_rows = num_tokens;
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);
result_output = torch::zeros({num_rows, hidden_size},
input.options().dtype(output_dtype));
}
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.";
if (Platform::is_support_ivcore10()) {
// BI-V100 path: per-token expert gather + matmul + SiLU-gate + matmul
// This replaces the tvm ffi CUTLASS path with native PyTorch ops.
for (int64_t t = 0; t < num_tokens; ++t) {
auto token = input[t].unsqueeze(0); // [1, hidden]
torch::Tensor accum = torch::zeros({1, hidden_size},
input.options().dtype(output_dtype));
for (int64_t k = 0; k < top_k; ++k) {
int64_t expert_id = token_selected_experts[t][k].item<int64_t>();
float scale = token_final_scales[t][k].item<float>();
// gate_up = token @ fc1[expert].T → [1, inter_dim]
auto gate_up = torch::mm(token, fc1_expert_weights[expert_id].t());
if (fc1_expert_biases.has_value()) {
gate_up = gate_up + fc1_expert_biases.value()[expert_id];
}
// SwiGLU: split into gate and up, apply silu(gate) * up
torch::Tensor act;
if (activation_type == ActivationType::SWIGLU ||
activation_type == ActivationType::SWIGLU_BIAS) {
auto chunks = gate_up.chunk(2, /*dim=*/-1);
act = torch::silu(chunks[0]) * chunks[1];
} else if (activation_type == ActivationType::SILU) {
act = torch::silu(gate_up);
} else if (activation_type == ActivationType::GELU) {
act = torch::gelu(gate_up);
} else {
act = gate_up; // identity
}
// down = act @ fc2[expert].T → [1, hidden]
auto down = torch::mm(act, fc2_expert_weights[expert_id].t());
if (fc2_expert_biases.has_value()) {
down = down + fc2_expert_biases.value()[expert_id];
}
accum += down.to(output_dtype) * scale;
}
result_output[t] = accum.squeeze(0);
}
return result_output;
}
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);
// Original NVIDIA GPU path (sm90/sm100/sm120) via tvm ffi
TORCH_CHECK(false,
"cutlass_fused_moe: no supported platform. "
"BI-V100 should use ivcore10 path above; "
"NVIDIA GPUs require sm90+.");
return result_output;
}
} // namespace xllm::kernel::cuda
} // namespace xllm::kernel::cuda

4
qwen3_6_scripts/ex_engine/xllm_kernels/cuda/moe/moe_topk_softmax_kernels.cuh
View File

@@ -23,8 +23,8 @@ limitations under the License.
#include <cub/util_type.cuh>
#if !defined(USE_DCU) && !defined(USE_MACA)
#endif
// <cuda/functional> requires CUDA 12+ (libcudacxx); BI-V100 runs CUDA 10.2
// and does not ship that header. The include is unused in this file anyway.
#include "device_utils.cuh"

7
qwen3_6_scripts/ex_engine/xllm_kernels/param.h
View File

@@ -817,6 +817,7 @@ struct RandomSampleParams {
};
// Rejection sampling parameters for speculative decoding
// PRD #0 (9e0f1402): Added mode field for explicit greedy/probabilistic dispatch
struct RejectionSampleParams {
// Candidate draft token indices to be verified.
// Shape: [total_draft_tokens]. Dtype: int32.
@@ -849,6 +850,10 @@ struct RejectionSampleParams {
// The maximum number of draft tokens in the batch (max value in
// num_draft_tokens).
int32_t max_spec_len;
// PRD #0 (9e0f1402): Explicit sampling mode — Greedy uses exact-match,
// Probabilistic uses min(1, p/q) acceptance ratio.
// Default: Probabilistic (preserves legacy behavior).
uint8_t draft_sampling_mode = 1; // 0=Greedy, 1=Probabilistic
};
// Masked indexer select paged KV cache parameters
@@ -1438,4 +1443,4 @@ struct ChunkGatedDeltaRuleParams {
// Whether to apply L2 norm to q and k inside the kernel. Default: false.
bool use_qk_l2norm_in_kernel = false;
};
} // namespace xllm::kernel
} // namespace xllm::kernel