152 lines
5.5 KiB
Python
152 lines
5.5 KiB
Python
"""
|
|
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()
|