7127d18491
Symbol probe revealed ixformer::infer namespace does NOT exist in base image.
That namespace is xllm's own compiled wrapper layer.
Actual available symbols in base image:
_ixformer_torch.so: silu_and_mul_forward, rms_norm_forward,
fused_add_rms_norm_forward, ixformer_linear, ixformer_linear_ex
libixformer.so: ixinfer_flash_attn_unpad_fwd (different signature)
MoE functions (topk_softmax, group_gemm, moe_expand_input, etc.)
are NOT in any base image .so — MoE must use Python path.
Bridge now only wraps: silu_and_mul, rms_norm, fused_add_rms_norm, linear
These accelerate the per-layer ops that run 200x per token.
91 lines
3.9 KiB
C++
91 lines
3.9 KiB
C++
// ix_full_bridge.cpp — Bridge to ixformer C++ functions available in base image
|
|
//
|
|
// Based on symbol probe of the actual BI-V100 base image:
|
|
// _ixformer_torch.so has: silu_and_mul_forward, rms_norm_forward,
|
|
// fused_add_rms_norm_forward, ixformer_linear, ixformer_linear_ex
|
|
// libixformer.so has: ixinfer_flash_attn_unpad_fwd
|
|
//
|
|
// MoE functions (topk_softmax, group_gemm, etc.) are NOT in base image.
|
|
// They exist only in xllm's compiled library. MoE must use Python fallback.
|
|
|
|
#include <torch/extension.h>
|
|
#include <optional>
|
|
#include <tuple>
|
|
#include <vector>
|
|
|
|
// ============================================================================
|
|
// Forward declarations — ACTUAL symbols from base image .so files
|
|
// Namespace: ixformer_torch_ext (in _ixformer_torch.cpython-310.so)
|
|
// ============================================================================
|
|
namespace ixformer_torch_ext {
|
|
|
|
// silu_and_mul: _ZN18ixformer_torch_ext20silu_and_mul_forwardERN2at6TensorES2_
|
|
void silu_and_mul_forward(at::Tensor& input, at::Tensor& output);
|
|
|
|
// rms_norm: _ZN18ixformer_torch_ext16rms_norm_forwardERN2at6TensorES2_S2_d
|
|
void rms_norm_forward(at::Tensor& input, at::Tensor& weight, at::Tensor& output, double eps);
|
|
|
|
// fused_add_rms_norm: _ZN18ixformer_torch_ext26fused_add_rms_norm_forwardERN2at6TensorES2_S2_dd
|
|
void fused_add_rms_norm_forward(at::Tensor& input, at::Tensor& residual,
|
|
at::Tensor& weight, double eps, double alpha);
|
|
|
|
// ixformer_linear: _ZN18ixformer_torch_ext15ixformer_linearERN2at6TensorES2_RKN3c108optionalIS1_EES7_
|
|
at::Tensor ixformer_linear(at::Tensor& input, at::Tensor& weight,
|
|
const c10::optional<at::Tensor>& bias,
|
|
const c10::optional<at::Tensor>& out);
|
|
|
|
// ixformer_linear_ex: _ZN18ixformer_torch_ext18ixformer_linear_exERN2at6TensorES2_RKN3c108optionalIS1_EE
|
|
at::Tensor ixformer_linear_ex(at::Tensor& input, at::Tensor& weight,
|
|
const c10::optional<at::Tensor>& bias);
|
|
|
|
} // namespace ixformer_torch_ext
|
|
|
|
|
|
// ============================================================================
|
|
// Python wrappers
|
|
// ============================================================================
|
|
|
|
// --- silu_and_mul ---
|
|
torch::Tensor ix_silu_and_mul(torch::Tensor input) {
|
|
int64_t half_dim = input.size(-1) / 2;
|
|
auto output = input.new_empty({input.size(0), half_dim});
|
|
ixformer_torch_ext::silu_and_mul_forward(input, output);
|
|
return output;
|
|
}
|
|
|
|
// --- rms_norm ---
|
|
void ix_rms_norm(torch::Tensor output, torch::Tensor input,
|
|
torch::Tensor weight, double eps) {
|
|
ixformer_torch_ext::rms_norm_forward(input, weight, output, eps);
|
|
}
|
|
|
|
// --- fused_add_rms_norm ---
|
|
void ix_fused_add_rms_norm(torch::Tensor input, torch::Tensor residual,
|
|
torch::Tensor weight, double eps) {
|
|
ixformer_torch_ext::fused_add_rms_norm_forward(input, residual, weight, eps, 1.0);
|
|
}
|
|
|
|
// --- linear ---
|
|
torch::Tensor ix_linear(torch::Tensor input, torch::Tensor weight,
|
|
const c10::optional<torch::Tensor>& bias) {
|
|
// Use linear_ex for decode (m<=1), linear for prefill
|
|
auto input_2d = input.view({-1, input.size(-1)});
|
|
int64_t m = input_2d.size(0);
|
|
if (m <= 1 && !bias.has_value()) {
|
|
return ixformer_torch_ext::ixformer_linear_ex(input, weight, bias);
|
|
}
|
|
return ixformer_torch_ext::ixformer_linear(input, weight, bias,
|
|
c10::optional<at::Tensor>());
|
|
}
|
|
|
|
|
|
// ============================================================================
|
|
// Module registration
|
|
// ============================================================================
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("silu_and_mul", &ix_silu_and_mul, "Fused SiLU+mul activation");
|
|
m.def("rms_norm", &ix_rms_norm, "RMSNorm");
|
|
m.def("fused_add_rms_norm", &ix_fused_add_rms_norm, "Fused residual + RMSNorm");
|
|
m.def("linear", &ix_linear, "ixformer GEMM (linear/linear_ex)");
|
|
}
|