d9ffc5159d
nm -D confirmed real symbols are in ixformer_torch_ext:: not ixformer::infer:: Bridges: silu_and_mul, rms_norm, fused_add_rms_norm, rotary_embedding, reshape_and_cache
86 lines
3.9 KiB
C++
86 lines
3.9 KiB
C++
// ix_full_bridge.cpp — Bridge to ixformer_torch_ext C++ functions
|
|
//
|
|
// Real namespace: ixformer_torch_ext (from nm -D _ixformer_torch.cpython-310.so)
|
|
// NOT ixformer::infer (that namespace doesn't exist in BI-V100 ixformer 3.2.3)
|
|
//
|
|
// Symbols confirmed on real machine:
|
|
// ixformer_torch_ext::silu_and_mul_forward(Tensor&, Tensor&)
|
|
// ixformer_torch_ext::rms_norm_forward(Tensor&, Tensor&, Tensor&, double)
|
|
// ixformer_torch_ext::fused_add_rms_norm_forward(Tensor&, Tensor&, Tensor&, double, double)
|
|
// ixformer_torch_ext::vllm_rotary_embedding_neox(Tensor&, Tensor&, Tensor&, long, Tensor&, long, bool)
|
|
// ixformer_torch_ext::vllm_cache_ops_reshape_and_cache(Tensor&, Tensor&, Tensor&, Tensor&, Tensor&, long, long)
|
|
|
|
#include <torch/extension.h>
|
|
|
|
// Forward declarations — exact signatures from nm -D on real machine
|
|
namespace ixformer_torch_ext {
|
|
|
|
void silu_and_mul_forward(at::Tensor& input, at::Tensor& output);
|
|
|
|
void rms_norm_forward(at::Tensor& output, at::Tensor& input,
|
|
at::Tensor& weight, double eps);
|
|
|
|
void fused_add_rms_norm_forward(at::Tensor& input, at::Tensor& residual,
|
|
at::Tensor& weight, double eps, double dropout);
|
|
|
|
void vllm_rotary_embedding_neox(at::Tensor& positions, at::Tensor& query,
|
|
at::Tensor& key, long head_size,
|
|
at::Tensor& cos_sin_cache, long is_neox, bool interleaved);
|
|
|
|
void vllm_cache_ops_reshape_and_cache(at::Tensor& key, at::Tensor& value,
|
|
at::Tensor& key_cache, at::Tensor& value_cache,
|
|
at::Tensor& slot_mapping, long block_size,
|
|
long x);
|
|
|
|
void rms_norm_quant(at::Tensor& output, at::Tensor& input,
|
|
at::Tensor& weight, double eps);
|
|
|
|
} // namespace ixformer_torch_ext
|
|
|
|
// ============================================================================
|
|
// Python-facing wrappers
|
|
// ============================================================================
|
|
|
|
void ix_silu_and_mul(torch::Tensor& input, torch::Tensor& output) {
|
|
ixformer_torch_ext::silu_and_mul_forward(input, output);
|
|
}
|
|
|
|
void ix_rms_norm(torch::Tensor& output, torch::Tensor& input,
|
|
torch::Tensor& weight, double eps) {
|
|
ixformer_torch_ext::rms_norm_forward(output, input, weight, eps);
|
|
}
|
|
|
|
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, 0.0);
|
|
}
|
|
|
|
void ix_rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
|
|
torch::Tensor& key, long head_size,
|
|
torch::Tensor& cos_sin_cache, bool is_neox) {
|
|
ixformer_torch_ext::vllm_rotary_embedding_neox(
|
|
positions, query, key, head_size, cos_sin_cache, is_neox ? 1 : 0, false);
|
|
}
|
|
|
|
void ix_reshape_and_cache(torch::Tensor& key, torch::Tensor& value,
|
|
torch::Tensor& key_cache, torch::Tensor& value_cache,
|
|
torch::Tensor& slot_mapping) {
|
|
long block_size = key_cache.size(1);
|
|
long x = key_cache.size(3);
|
|
ixformer_torch_ext::vllm_cache_ops_reshape_and_cache(
|
|
key, value, key_cache, value_cache, slot_mapping, block_size, x);
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
|
m.def("silu_and_mul", &ix_silu_and_mul,
|
|
"ixformer_torch_ext::silu_and_mul_forward bridge");
|
|
m.def("rms_norm", &ix_rms_norm,
|
|
"ixformer_torch_ext::rms_norm_forward bridge");
|
|
m.def("fused_add_rms_norm", &ix_fused_add_rms_norm,
|
|
"ixformer_torch_ext::fused_add_rms_norm_forward bridge");
|
|
m.def("rotary_embedding", &ix_rotary_embedding,
|
|
"ixformer_torch_ext::vllm_rotary_embedding_neox bridge");
|
|
m.def("reshape_and_cache", &ix_reshape_and_cache,
|
|
"ixformer_torch_ext::vllm_cache_ops_reshape_and_cache bridge");
|
|
}
|