3581dd5435
Layer 1: hw_config.h (245 lines) — BI-V100 hardware descriptor + tuning tables Layer 2: moe_pipeline.py (461 lines) — MoE 7-step pipeline orchestrator Layer 3: attn_dispatch.py (270 lines) — Attention prefill/decode dispatch Layer 4: ilu_ops_api.h (182 lines) — Dispatch signature contract Layer 5: kernel_moe_ops.cpp (155 lines) — MoE kernel-level ops wrappers Layer 6: kernel_elem_ops.cpp (210 lines) — Element-wise kernel wrappers Layer 7: ixformer_infer.h (246 lines) — ixformer::infer namespace contract Layer 8: factor_topk_softmax.cu (456 lines) — MoE routing CUDA kernel Layer 9: factor_moe_compute_index.cu (174 lines) — Token index CUDA kernel Layer 10: factor_moe_combine.cu (154 lines) — Weighted combine CUDA kernel Total: 2553 lines across 10 files (h/cpp/cu/py) Upstream reference: 2787 lines across corresponding 10 xllm AST layers Each file follows the read-read-read-write pattern from upstream xllm, ds_vllm, and fla repos. No hand-written inference code — all kernel logic is cat-migrated from the upstream references.
211 lines
7.3 KiB
C++
211 lines
7.3 KiB
C++
// ex_engine/factors/kernel_elem_ops.cpp
|
||
//
|
||
// Layer 6: Element-wise kernel dispatch wrappers
|
||
//
|
||
// Upstream parallel: kernels/ilu/activation.cpp (30 lines)
|
||
// + kernels/ilu/norm.cpp (45 lines)
|
||
// + kernels/ilu/rope.cpp (20 lines)
|
||
// + kernels/ilu/group_gemm.cpp (25 lines)
|
||
// + kernels/ilu/matmul.cpp (~15 lines)
|
||
//
|
||
// Total upstream: ~135 lines across 5 files.
|
||
// Each function is a 3-5 line dispatch wrapper that calls ixformer::infer.
|
||
//
|
||
// These ops are PRESENT in the base image's ixformer — they don't need
|
||
// EX factor replacement. But they must be in the call chain because:
|
||
// - Activation is Step 5 of the MoE pipeline (called between GEMM1 and GEMM2)
|
||
// - RMSNorm is called before/after every decoder layer (2× per layer × 64 layers)
|
||
// - RoPE is called once per attention layer (32 full attention + 4 GDN = 36)
|
||
// - Group GEMM is Steps 4+6 (called twice per MoE layer × 64 layers)
|
||
//
|
||
// The presence in ixformer is confirmed by:
|
||
// cat_files/symbol_dumps (nm -D output from real device)
|
||
// SYSTEM_DESIGN.md PRESENT list
|
||
|
||
#include "ilu_ops_api.h"
|
||
#include "ixformer.h"
|
||
|
||
using namespace ixformer;
|
||
|
||
namespace xllm {
|
||
namespace kernel {
|
||
namespace ilu {
|
||
|
||
// =====================================================================
|
||
// Activation: silu_and_mul
|
||
// =====================================================================
|
||
// Upstream: kernels/ilu/activation.cpp::act_and_mul
|
||
// BI-V100 ixformer: PRESENT (silu_and_mul confirmed in symbol dump)
|
||
//
|
||
// Input: (tokens, 2 × intermediate_size) — gate + up projections concatenated
|
||
// Output: (tokens, intermediate_size) — silu(gate) × up
|
||
//
|
||
// For Qwen3.5: intermediate_size = 18944 / TP4 = 4736
|
||
// Each call processes 4736 × 2 = 9472 half values per token.
|
||
// At 200 tokens/batch decode: 200 × 9472 × 2 bytes = 3.6 MB bandwidth.
|
||
|
||
void act_and_mul(
|
||
torch::Tensor out,
|
||
torch::Tensor input,
|
||
const std::string& act_mode) {
|
||
|
||
if (act_mode == "silu") {
|
||
infer::silu_and_mul(input, out);
|
||
} else {
|
||
// gelu_tanh_and_mul is MISSING from ixformer on BI-V100.
|
||
// The EX factor system provides this as EX_FACTOR_GELU_TANH_MUL (id=3).
|
||
// For now, fallback to PyTorch.
|
||
LOG(FATAL) << "Unsupported act mode: " << act_mode
|
||
<< ", only silu is available via ixformer on BI-V100. "
|
||
<< "Use EX factor 3 for gelu_tanh.";
|
||
}
|
||
}
|
||
|
||
// =====================================================================
|
||
// RMSNorm: rms_norm + residual_layer_norm
|
||
// =====================================================================
|
||
// Upstream: kernels/ilu/norm.cpp
|
||
// BI-V100 ixformer: PRESENT (rms_norm, fused_add_rms_norm confirmed)
|
||
//
|
||
// rms_norm: out = x × rsqrt(mean(x²) + eps) × weight
|
||
// residual_layer_norm: fused residual add + rms_norm
|
||
//
|
||
// Called 2× per decoder layer (pre-attention + post-attention norm).
|
||
// 64 layers × 2 = 128 calls per forward pass.
|
||
|
||
void rms_norm(
|
||
torch::Tensor& output,
|
||
torch::Tensor& input,
|
||
torch::Tensor& weight,
|
||
double eps) {
|
||
|
||
std::optional<torch::Tensor> fused_bias = std::nullopt;
|
||
infer::rms_norm(input, weight, output, fused_bias, eps);
|
||
}
|
||
|
||
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) {
|
||
|
||
auto residual_ = residual.value_or(torch::zeros_like(input));
|
||
torch::Tensor residual_out_ = residual_out.value_or(torch::zeros_like(input));
|
||
infer::residual_rms_norm(
|
||
input, residual_, weight, output, residual_out_,
|
||
bias, /*alpha=*/1.0, eps, /*is_post=*/false);
|
||
}
|
||
|
||
// =====================================================================
|
||
// RoPE: Rotary Position Embedding
|
||
// =====================================================================
|
||
// Upstream: kernels/ilu/rope.cpp
|
||
// BI-V100 ixformer: PRESENT (vllm_rotary_embedding_neox confirmed)
|
||
//
|
||
// Applies cosine-sine rotation to query and key tensors.
|
||
// Called once per attention layer per forward pass.
|
||
// Qwen3.5: 36 attention layers (32 full + 4 GDN).
|
||
|
||
void apply_rope_pos_ids_cos_sin_cache(
|
||
torch::Tensor& query,
|
||
torch::Tensor& key,
|
||
torch::Tensor& cos_sin_cache,
|
||
torch::Tensor& positions,
|
||
bool interleave) {
|
||
|
||
const int64_t head_size = cos_sin_cache.size(-1);
|
||
// is_neox = !interleave (NeoX-style = non-interleaved)
|
||
infer::xllm_rotary_embedding(
|
||
positions, query, key, head_size, cos_sin_cache, !interleave);
|
||
}
|
||
|
||
// =====================================================================
|
||
// Group GEMM: Batched matrix multiplication for MoE experts
|
||
// =====================================================================
|
||
// Upstream: kernels/ilu/group_gemm.cpp
|
||
// BI-V100 ixformer: PRESENT (moe_w16a16_group_gemm confirmed)
|
||
//
|
||
// Performs A × B^T for each expert group simultaneously.
|
||
// tokens_per_experts defines the row count per group.
|
||
// Called twice per MoE layer: once for w13 (gate+up), once for w2 (down).
|
||
//
|
||
// For Qwen3.5 with TP4:
|
||
// w13: (16 experts_local, 9472, 3584) — 16 experts × [inter*2, hidden]
|
||
// w2: (16 experts_local, 3584, 4736) — 16 experts × [hidden, inter]
|
||
//
|
||
// This is the primary compute bottleneck on BI-V100.
|
||
// sub694 shows prompt_tok >50K requests drop to 1-3 TPS — GEMM bound.
|
||
|
||
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) {
|
||
|
||
infer::moe_w16a16_group_gemm(
|
||
output,
|
||
input,
|
||
weight,
|
||
tokens_per_experts,
|
||
dst_to_src,
|
||
/*bias=*/std::nullopt,
|
||
/*format=*/"TN",
|
||
/*persistent=*/0,
|
||
/*output_n=*/tokens_per_experts.sum().item<int64_t>());
|
||
|
||
return output;
|
||
}
|
||
|
||
// =====================================================================
|
||
// Reshape and cache: KV cache management
|
||
// =====================================================================
|
||
// Upstream: kernels/ilu/attention.cpp::reshape_paged_cache
|
||
// BI-V100 ixformer: PRESENT (vllm_cache_ops_reshape_and_cache)
|
||
//
|
||
// Writes new KV pairs into paged cache at the positions specified
|
||
// by slot_mapping.
|
||
|
||
void reshape_paged_cache(
|
||
torch::Tensor& key,
|
||
std::optional<torch::Tensor>& value,
|
||
torch::Tensor& key_cache,
|
||
std::optional<torch::Tensor>& value_cache,
|
||
torch::Tensor& slot_mapping) {
|
||
|
||
auto value_ = value.value_or(torch::Tensor());
|
||
auto value_cache_ = value_cache.value_or(torch::Tensor());
|
||
|
||
int64_t key_token_stride = key.stride(0);
|
||
int64_t value_token_stride = 0;
|
||
if (value_.defined()) {
|
||
value_token_stride = value_.stride(0);
|
||
}
|
||
slot_mapping = slot_mapping.to(at::kLong);
|
||
|
||
infer::xllm_reshape_and_cache(
|
||
key, value_, key_cache, value_cache_,
|
||
slot_mapping, key_token_stride, value_token_stride);
|
||
}
|
||
|
||
// =====================================================================
|
||
// Matmul: General matrix multiplication
|
||
// =====================================================================
|
||
// Upstream: kernels/ilu/matmul.cpp
|
||
// Used for linear projections (q/k/v proj, out proj, gate proj)
|
||
|
||
torch::Tensor matmul(
|
||
torch::Tensor a,
|
||
torch::Tensor b,
|
||
std::optional<torch::Tensor> bias) {
|
||
|
||
return infer::ixformer_linear_ex(a, b, bias, std::nullopt);
|
||
}
|
||
|
||
} // namespace ilu
|
||
} // namespace kernel
|
||
} // namespace xllm
|