完整部署快照

qwen3_6_scripts/patch_ixformer_infer.py

@@ -0,0 +1,151 @@
+"""
+Bridge ixformer SDK 0.6.0 ops.infer.* API to CoreX 3.2.3 _C._functions.* API.
+
+CoreX 3.2.3 ships ixformer._C with _functions submodule (v0.3.0 API).
+SDK 0.6.0 inference/functions/*.py expects _C.infer.* namespace.
+This shim creates _C.infer and maps each call to _C._functions or
+a torch-native fallback where _functions lacks the operation.
+
+Must be imported BEFORE ixformer.inference.functions.
+Idempotent.
+"""
+import types
+import functools
+
+
+def patch():
+    try:
+        import ixformer._C as _C
+    except ImportError:
+        return False
+
+    if hasattr(_C, 'infer'):
+        return False  # already patched or native
+
+    _fn = _C._functions
+    infer = types.ModuleType('ixformer._C.infer')
+
+    # --- Direct mappings (same semantics, slightly different signatures) ---
+
+    def _linear(input, weight, act_type=-1, bias=None, output=None,
+                persistent=False):
+        """Map infer.linear -> _functions.linear_forward + activation."""
+        import torch, torch.nn.functional as F
+        if output is None:
+            out_shape = list(input.shape)
+            out_shape[-1] = weight.shape[0]
+            output = input.new_empty(out_shape)
+        # _functions.linear_forward expects (input, weight) or (input, weight, bias)
+        # and writes into a pre-allocated output via linear_forward_
+        try:
+            if bias is not None:
+                _fn.linear_forward_(input, weight, bias, output)
+            else:
+                _fn.linear_forward_(input, weight, output)
+        except Exception:
+            # Signature mismatch fallback to torch
+            output = F.linear(input, weight, bias)
+
+        # Apply activation
+        if act_type == 3:
+            output = F.gelu(output)
+        elif act_type == 4:
+            output = F.relu(output)
+        elif act_type == 12:
+            output = F.silu(output)
+        return output
+
+    def _linear_ex(input, weight, bias=None, output=None):
+        """Map infer.linear_ex -> _functions.linear_forward (gemv path)."""
+        import torch, torch.nn.functional as F
+        if output is None:
+            out_shape = list(input.shape)
+            out_shape[-1] = weight.shape[0]
+            output = input.new_empty(out_shape)
+        try:
+            if bias is not None:
+                _fn.linear_forward_(input, weight, bias, output)
+            else:
+                _fn.linear_forward_(input, weight, output)
+        except Exception:
+            output = F.linear(input, weight, bias)
+        return output
+
+    def _mixed_type_linear(input, weight, bias=None, output=None):
+        import torch.nn.functional as F
+        result = F.linear(input.to(weight.dtype), weight, bias)
+        if output is not None:
+            output.copy_(result)
+            return output
+        return result
+
+    # --- Functions with direct name mapping ---
+    def _map(infer_name, fn_name):
+        """Create a passthrough if _functions has the function."""
+        if hasattr(_fn, fn_name):
+            setattr(infer, infer_name, getattr(_fn, fn_name))
+            return True
+        return False
+
+    # Register mappings
+    infer.linear = _linear
+    infer.linear_ex = _linear_ex
+    infer.mixed_type_linear = _mixed_type_linear
+
+    # Direct name mappings (infer.X -> _functions.X_forward or _functions.X)
+    direct_maps = {
+        'rms_norm': 'rms_norm_forward',
+        'silu_and_mul': 'silu_and_mul_forward',
+        'softmax': 'softmax_forward',
+        'layernorm': 'layernorm_forward',
+        'layer_norm': 'layernorm_forward',
+        'groupnorm': 'groupnorm_forward',
+        'conv': 'conv2d_forward',
+        'matmul': 'matmul',
+        'gelu_forward': 'gelu_forward',
+        'act_bias_mm': 'act_bias_mm_forward',
+        'bnb_dequant': 'bnb_dequant_forward',
+        'bnb_qgemm': 'bnb_qgemm_forward',
+        'bnb_quant': 'bnb_quant_forward',
+        'bnb_rowcol_absmax': 'bnb_rowcol_absmax_forward',
+        'bnb_doubleRowColQuant': 'bnb_doubleRowColQuant',
+        'bnb_getColRowStats': 'bnb_getColRowStats',
+        'bnb_mm_dequant': 'bnb_mm_dequant',
+        'ixinfer_flash_attn_pad_fwd': 'ixinfer_flash_attn_pad_fwd',
+        'ixinfer_flash_attn_pad_fwd_nomask': 'ixinfer_flash_attn_pad_fwd_nomask',
+        'ixinfer_flash_attn_unpad': 'ixinfer_flash_attn_unpad_fwd',
+        'vllm_cache_ops_reshape_and_cache': 'vllm_cache_ops_reshape_and_cache',
+        'vllm_rotary_embedding': 'vllm_rotary_embedding_neox',
+    }
+    for infer_name, fn_name in direct_maps.items():
+        _map(infer_name, fn_name)
+
+    # For any remaining ops.infer.xxx call that we haven't mapped,
+    # provide a __getattr__ that gives a clear error
+    class _InferModule(types.ModuleType):
+        def __getattr__(self, name):
+            # Check _functions first
+            fn_name = name + '_forward'
+            if hasattr(_fn, fn_name):
+                return getattr(_fn, fn_name)
+            if hasattr(_fn, name):
+                return getattr(_fn, name)
+            raise AttributeError(
+                f"ixformer._C.infer has no attribute '{name}' "
+                f"(CoreX 3.2.3 _functions does not provide it)")
+
+    # Transfer all explicit mappings to the smart module
+    smart_infer = _InferModule('ixformer._C.infer')
+    for attr in dir(infer):
+        if not attr.startswith('_'):
+            setattr(smart_infer, attr, getattr(infer, attr))
+
+    _C.infer = smart_infer
+    import sys
+    sys.modules['ixformer._C.infer'] = smart_infer
+    print('[ixformer_infer] bridged ops.infer -> _C._functions '
+          f'({len(direct_maps)+3} mappings)')
+    return True
+
+
+_applied = patch()