From 9a5206853362ea6a55a08d823f2d833d3ad429e0 Mon Sep 17 00:00:00 2001 From: root Date: Wed, 19 Aug 2026 09:57:51 +0000 Subject: [PATCH] =?UTF-8?q?[fix]=20baseline5=20fused=5Fadd=5Frms=5Fnorm=20?= =?UTF-8?q?=E7=A1=AE=E5=AE=9E=E6=98=AF=E5=8F=82=E6=95=B0=E6=95=B0=E9=87=8F?= =?UTF-8?q?=E4=B8=8D=E5=8C=B9=E9=85=8D=E3=80=82rms=5Fnorm=20=E5=88=99?= =?UTF-8?q?=E6=98=AF=E5=8F=82=E6=95=B0=E6=95=B0=E9=87=8F=E5=AF=B9=E4=BA=86?= =?UTF-8?q?=E4=BD=86=E7=BB=93=E6=9E=9C=E6=B2=A1=E5=86=99=E5=9B=9E=E5=8E=BB?= =?UTF-8?q?=E3=80=82=E4=B8=A4=E4=B8=AA=20bug=20=E5=8F=A0=E5=8A=A0=E5=AF=BC?= =?UTF-8?q?=E8=87=B4=E6=A8=A1=E5=9E=8B=E8=BE=93=E5=87=BA=E5=85=A8=E6=8D=9F?= =?UTF-8?q?=E5=9D=8F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../ex_engine/python/patch_vllm_hot_path.py | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/qwen3_6_scripts/ex_engine/python/patch_vllm_hot_path.py b/qwen3_6_scripts/ex_engine/python/patch_vllm_hot_path.py index 43b72474..ebbe780e 100644 --- a/qwen3_6_scripts/ex_engine/python/patch_vllm_hot_path.py +++ b/qwen3_6_scripts/ex_engine/python/patch_vllm_hot_path.py @@ -93,14 +93,18 @@ def apply(strict=True): # vllm uses ops.rms_norm / ops.fused_add_rms_norm import vllm._custom_ops as ops - def patched_rms_norm(output, input, weight, epsilon): - # vllm: ops.rms_norm(output, input, weight, eps) - # xllm_norm.so: rms_norm(output, input, weight, eps) — same - xllm_ops._get("xllm_norm").rms_norm(output, input, weight, epsilon) + # Cache the C++ module once (avoids dict lookup per call) + _norm_mod = xllm_ops._get("xllm_norm") - def patched_fused_add_rms_norm(input, residual, weight, epsilon): - # vllm: ops.fused_add_rms_norm(input, residual, weight, eps) - # xllm_norm.so: same signature, in-place + def patched_rms_norm(output, input, weight, epsilon): + # C++ signature: rms_norm(output, input, weight, eps) + # Write directly into caller's output tensor (zero-copy). + _norm_mod.rms_norm(output, input, weight, epsilon) + + def patched_fused_add_rms_norm(input, residual, weight, epsilon, + residual_alpha=1.0): + if residual_alpha != 1.0: + residual.mul_(residual_alpha) xllm_ops.residual_rms_norm(input, residual, weight, epsilon) if hasattr(ops, 'rms_norm'): @@ -126,9 +130,7 @@ def apply(strict=True): import vllm._custom_ops as ops def patched_silu_and_mul(output, input): - # vllm: ops.silu_and_mul(output, input) - # xllm_activation.so: silu_and_mul(out, input) — same order - xllm_ops._get("xllm_activation").silu_and_mul(output, input) + xllm_ops.silu_and_mul(input, output) if hasattr(ops, 'silu_and_mul'): ops.silu_and_mul = patched_silu_and_mul @@ -170,8 +172,7 @@ def apply(strict=True): import vllm._custom_ops as ops def patched_reshape_and_cache(key, value, key_cache, value_cache, - slot_mapping, kv_cache_dtype, - k_scale, v_scale): + slot_mapping, kv_cache_dtype, kv_scale): xllm_ops.reshape_and_cache(key, value, key_cache, value_cache, slot_mapping)