From b342eb6b98e533ed43547e37a46f31476a8462d4 Mon Sep 17 00:00:00 2001 From: root Date: Mon, 17 Aug 2026 08:26:54 +0000 Subject: [PATCH] =?UTF-8?q?[fix]=20baseline4=20fix=20prebuilt=20ix=5Ffull?= =?UTF-8?q?=5Fbridge.so=20=E5=A4=A7=E6=A6=82=E7=8E=87=E6=98=AF=E4=BB=8E=20?= =?UTF-8?q?v2=20=E7=BC=96=E8=AF=91=E7=9A=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- ex_engine/csrc/ix_full_bridge_v2.cpp | 9 ++++++--- ex_engine/python/patch_vllm_ops.py | 17 +++++++++++------ 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/ex_engine/csrc/ix_full_bridge_v2.cpp b/ex_engine/csrc/ix_full_bridge_v2.cpp index d81cef68..22b94ead 100644 --- a/ex_engine/csrc/ix_full_bridge_v2.cpp +++ b/ex_engine/csrc/ix_full_bridge_v2.cpp @@ -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 --- diff --git a/ex_engine/python/patch_vllm_ops.py b/ex_engine/python/patch_vllm_ops.py index 6a23fcbb..3403c73d 100644 --- a/ex_engine/python/patch_vllm_ops.py +++ b/ex_engine/python/patch_vllm_ops.py @@ -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) \ No newline at end of file