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.
247 lines
7.8 KiB
C++
247 lines
7.8 KiB
C++
// ex_engine/factors/ixformer_infer.h
|
|
//
|
|
// Layer 7: ixformer::infer namespace contract
|
|
//
|
|
// Upstream parallel: kernels/ilu/ixformer.h (~140 lines)
|
|
// Declares every function in ixformer::infer that the ilu kernel
|
|
// wrappers (Layer 5-6) call through to.
|
|
//
|
|
// On BI-V100, this namespace is implemented by two sources:
|
|
//
|
|
// 1. BASE IMAGE (libixformer.so from corex SDK 3.2.3):
|
|
// PRESENT — these symbols exist in nm -D of the .so:
|
|
// silu_and_mul
|
|
// rms_norm
|
|
// residual_rms_norm
|
|
// xllm_rotary_embedding (aka vllm_rotary_embedding_neox)
|
|
// xllm_reshape_and_cache
|
|
// xllm_paged_attention (v1/v2)
|
|
// ixinfer_flash_attn_unpad_with_block_tables
|
|
// ixformer_linear / ixformer_linear_ex
|
|
// topk_softmax
|
|
// moe_compute_token_index_api
|
|
// moe_expand_input
|
|
// moe_w16a16_group_gemm
|
|
// moe_output_reduce_sum
|
|
//
|
|
// 2. EX FACTORS (.so files from ex_engine/build):
|
|
// For MISSING ops that may not be in all base image versions.
|
|
// The EX factor .so exports the same symbol → dlopen replaces it.
|
|
//
|
|
// Signature source: Verbatim from upstream xllm ixformer.h + utils.h,
|
|
// cross-referenced with cat_files/symbol_dumps nm -D output.
|
|
|
|
#ifndef EX_FACTORS_IXFORMER_INFER_H
|
|
#define EX_FACTORS_IXFORMER_INFER_H
|
|
|
|
#include <torch/all.h>
|
|
#include <ATen/Tensor.h>
|
|
#include <optional>
|
|
#include <string>
|
|
|
|
namespace ixformer {
|
|
namespace infer {
|
|
|
|
// =====================================================================
|
|
// Attention kernels
|
|
// =====================================================================
|
|
|
|
// Flash attention with block tables (prefill path)
|
|
// Source: ixinfer flash attention unpadded variant
|
|
// BI-V100 status: PRESENT in base image
|
|
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);
|
|
|
|
// Paged attention (decode path, v1 or v2 selected internally)
|
|
// BI-V100 status: PRESENT
|
|
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);
|
|
|
|
// =====================================================================
|
|
// Activation kernels
|
|
// =====================================================================
|
|
|
|
// SiLU-and-mul: out = silu(input[:half]) * input[half:]
|
|
// BI-V100 status: PRESENT
|
|
void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
|
|
|
|
// =====================================================================
|
|
// Linear / GEMM kernels
|
|
// =====================================================================
|
|
|
|
// ixformer linear: fused matmul with optional activation
|
|
// BI-V100 status: PRESENT
|
|
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);
|
|
|
|
// ixformer linear extended: simplified interface
|
|
// BI-V100 status: PRESENT
|
|
torch::Tensor ixformer_linear_ex(
|
|
torch::Tensor& input,
|
|
torch::Tensor& weight,
|
|
const c10::optional<torch::Tensor>& bias,
|
|
const c10::optional<torch::Tensor>& out);
|
|
|
|
// =====================================================================
|
|
// Cache management kernels
|
|
// =====================================================================
|
|
|
|
// Write KV into paged cache
|
|
// BI-V100 status: PRESENT
|
|
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);
|
|
|
|
// =====================================================================
|
|
// Rotary embedding kernels
|
|
// =====================================================================
|
|
|
|
// Apply rotary position encoding
|
|
// BI-V100 status: PRESENT
|
|
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);
|
|
|
|
// =====================================================================
|
|
// Normalization kernels
|
|
// =====================================================================
|
|
|
|
// Fused residual + RMS normalization
|
|
// BI-V100 status: PRESENT
|
|
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);
|
|
|
|
// RMS normalization
|
|
// BI-V100 status: PRESENT
|
|
void rms_norm(
|
|
torch::Tensor& input,
|
|
torch::Tensor& weight,
|
|
torch::Tensor& output,
|
|
const std::optional<torch::Tensor>& fused_bias,
|
|
double eps);
|
|
|
|
// =====================================================================
|
|
// MoE kernels
|
|
// =====================================================================
|
|
|
|
// Fused softmax + top-k for MoE routing
|
|
// BI-V100 status: PRESENT (confirmed in base image symbol dump)
|
|
// Calls CUDA kernel: moe_topk_softmax_kernels.cuh (Layer 8)
|
|
void topk_softmax(
|
|
torch::Tensor& topk_weights,
|
|
torch::Tensor& topk_indices,
|
|
torch::Tensor& token_expert_indices,
|
|
torch::Tensor& gating_output,
|
|
bool renormalize);
|
|
|
|
// 3-phase permutation index computation for MoE token dispatch
|
|
// BI-V100 status: PRESENT
|
|
// Calls CUDA kernel: moe_compute_index.cu (Layer 9)
|
|
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);
|
|
|
|
// Gather tokens from natural order to expert-sorted order
|
|
// BI-V100 status: PRESENT
|
|
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);
|
|
|
|
// Group GEMM for MoE expert computation (half-precision)
|
|
// BI-V100 status: PRESENT
|
|
// This is the primary compute bottleneck.
|
|
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);
|
|
|
|
// Weighted combine of expert outputs
|
|
// BI-V100 status: PRESENT
|
|
// Calls CUDA kernel: moe_combine.cu (Layer 10)
|
|
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 infer
|
|
} // namespace ixformer
|
|
|
|
#endif // EX_FACTORS_IXFORMER_INFER_H
|