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

merge modelhub: bridge_v2 + patch_vllm_ops fix

Browse Source
This commit is contained in:
root
2026-08-17 08:34:14 +00:00
parent c044d04509 b342eb6b98
commit b81a071f84
2 changed files with 17 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
ex_engine/csrc/ix_full_bridge_v2.cpp
View File

@@ -36,8 +36,9 @@ namespace ixformer_torch_ext {
void silu_and_mul_forward(at::Tensor& input, at::Tensor& output);
// rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double)
void rms_norm_forward(at::Tensor& output, at::Tensor& input,
at::Tensor& weight, double eps);
// Real ixformer signature order: (input, weight, output, eps)
void rms_norm_forward(at::Tensor& input, at::Tensor& weight,
at::Tensor& output, double eps);
// fused_add_rms_norm_forward(at::Tensor&, at::Tensor&, at::Tensor&, double, double)
void fused_add_rms_norm_forward(at::Tensor& input, at::Tensor& residual,
@@ -148,7 +149,9 @@ torch::Tensor ix_silu_and_mul(torch::Tensor input) {
// --- rms_norm ---
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);
// pybind receives (output, input, weight, eps)
// ixformer expects (input, weight, output, eps)
ixformer_torch_ext::rms_norm_forward(input, weight, output, eps);
}
// --- fused_add_rms_norm ---

17
ex_engine/python/patch_vllm_ops.py
View File

@@ -84,25 +84,30 @@ def _patch_layernorm() -> int:
_orig_forward = GemmaRMSNorm.forward
def _patched_forward(self, x, residual=None):
# GemmaRMSNorm: output = rms_norm(x) * (1 + weight)
# ixformer rms_norm: output = rms_norm(x) * weight
# Pass (1 + weight) to ixformer to match GemmaRMSNorm semantics.
w = self.weight
if w.dim() != 1 or w.shape[0] != x.shape[-1]:
return _orig_forward(self, x, residual)
w_adjusted = 1.0 + w
if residual is not None:
# fused_add_rms_norm: norm(x + residual) → (normed, new_residual)
if ix_ops.has_fused_add_rms_norm():
out = torch.empty_like(x)
residual_out = torch.empty_like(x)
ix_ops.fused_add_rms_norm(
x, residual, self.weight, out, residual_out,
x, residual, w_adjusted, out, residual_out,
self.variance_epsilon)
return out, residual_out
else:
# Two-step fallback using just rms_norm
new_residual = x + residual
out = torch.empty_like(x)
ix_ops.rms_norm(out, new_residual, self.weight,
ix_ops.rms_norm(out, new_residual, w_adjusted,
self.variance_epsilon)
return out, new_residual
else:
out = torch.empty_like(x)
ix_ops.rms_norm(out, x, self.weight, self.variance_epsilon)
ix_ops.rms_norm(out, x, w_adjusted, self.variance_epsilon)
return out
GemmaRMSNorm.forward = _patched_forward
@@ -198,4 +203,4 @@ if os.environ.get("IX_OPS_AUTO_PATCH", "0") == "1":
try:
apply_all_patches()
except Exception as e:
logger.warning("ix_ops auto-patch failed: %s", e)
logger.warning("ix_ops auto-patch failed: %s", e)