From 9e4fb3712fba830ad4e6de0f7b289284a6129429 Mon Sep 17 00:00:00 2001 From: EngineX Date: Mon, 10 Aug 2026 02:49:32 +0000 Subject: [PATCH] =?UTF-8?q?remove=20enginex/=20=E2=80=94=20duplicate=20of?= =?UTF-8?q?=20existing=20ex=5Fengine/?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- enginex/__init__.py | 16 -- enginex/bridge.py | 112 --------- enginex/dispatch/__init__.py | 0 enginex/dispatch/registry.py | 474 ----------------------------------- enginex/ops/__init__.py | 0 enginex/ops/activations.py | 39 --- enginex/ops/attention.py | 182 -------------- enginex/ops/cache.py | 63 ----- enginex/ops/gdn.py | 167 ------------ enginex/ops/moe.py | 201 --------------- enginex/ops/norm.py | 36 --- enginex/ops/sampling.py | 53 ---- enginex/tests/__init__.py | 0 enginex/tests/test_ops.py | 207 --------------- 14 files changed, 1550 deletions(-) delete mode 100644 enginex/__init__.py delete mode 100644 enginex/bridge.py delete mode 100644 enginex/dispatch/__init__.py delete mode 100644 enginex/dispatch/registry.py delete mode 100644 enginex/ops/__init__.py delete mode 100644 enginex/ops/activations.py delete mode 100644 enginex/ops/attention.py delete mode 100644 enginex/ops/cache.py delete mode 100644 enginex/ops/gdn.py delete mode 100644 enginex/ops/moe.py delete mode 100644 enginex/ops/norm.py delete mode 100644 enginex/ops/sampling.py delete mode 100644 enginex/tests/__init__.py delete mode 100644 enginex/tests/test_ops.py diff --git a/enginex/__init__.py b/enginex/__init__.py deleted file mode 100644 index abf7419c..00000000 --- a/enginex/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -""" -EngineX — Algorithm Factor Replacement Engine for BI-V100 - -Architecture modeled after CCCL's dispatch/tuning/kernel three-layer system: - CCCL: tuning_*.cuh → dispatch_*.cuh → kernel_*.cuh - EngineX: tuning/*.py → dispatch/*.py → ops/*.py - -The engine dlopen()s native .so when available, falls back to PyTorch/ixformer. -This is NOT an adapter — it's a full algorithm factor replacement layer. -""" - -__version__ = "0.1.0" - -from enginex.dispatch.registry import OperatorRegistry, get_registry - -__all__ = ["OperatorRegistry", "get_registry"] diff --git a/enginex/bridge.py b/enginex/bridge.py deleted file mode 100644 index d2a159c8..00000000 --- a/enginex/bridge.py +++ /dev/null @@ -1,112 +0,0 @@ -""" -EngineX Bridge — wires EngineX dispatch into vllm's _custom_ops.py - -This module monkey-patches _custom_ops functions to use EngineX's -three-tier dispatch instead of direct ixformer calls. - -The key fix: vllm_moe_topk_softmax is MISSING from ixformer.functions -on our BI-V100 image, causing every MoE forward pass to crash. -EngineX provides a PyTorch replacement that keeps the model running. - -Usage in patch_ops.sh: - python -c "import enginex.bridge; enginex.bridge.patch_custom_ops()" - -Or at runtime startup: - from enginex.bridge import patch_custom_ops - patch_custom_ops() -""" - -import importlib -import logging -import sys - -logger = logging.getLogger("enginex.bridge") - - -def patch_custom_ops(): - """ - Patch vllm._custom_ops to use EngineX dispatch. - - Strategy: only patch functions that are KNOWN BROKEN. - We do NOT touch working ixformer ops (silu_and_mul, rms_norm, etc.) - because ixformer's implementations are faster. - - From docker log analysis, the BROKEN ops are: - 1. topk_softmax — AttributeError: no 'vllm_moe_topk_softmax' - 2. invoke_fused_moe_kernel — falls back to PyTorch on topk failure - 3. moe_align_block_size — cascading failure from topk - """ - from enginex.dispatch.registry import get_registry - - reg = get_registry() - reg.probe() - - logger.info("EngineX bridge: patching broken ops in _custom_ops") - logger.info(reg.summary()) - - # Only patch if the module is already imported - custom_ops = sys.modules.get('vllm._custom_ops') - if custom_ops is None: - try: - custom_ops = importlib.import_module('vllm._custom_ops') - except ImportError: - logger.warning("vllm._custom_ops not found — skipping bridge") - return - - # ---- Patch 1: moe_topk_softmax (THE critical fix) ---- - moe_topk = reg.get_op("moe_topk_softmax") - if moe_topk: - original = getattr(custom_ops, 'topk_softmax', None) - if original: - # Check if the original actually works - try: - import ixformer.functions as ixf_F - _ = ixf_F.vllm_moe_topk_softmax - logger.info("EngineX: ixformer.vllm_moe_topk_softmax exists, " - "keeping original") - except (ImportError, AttributeError): - logger.info("EngineX: patching topk_softmax → EngineX " - f"({reg.get_backend('moe_topk_softmax').name})") - custom_ops.topk_softmax = moe_topk - - # ---- Patch 2: moe_align_block_size ---- - moe_align = reg.get_op("moe_align_block_size") - if moe_align: - try: - import ixformer.functions as ixf_F - _ = ixf_F.vllm_moe_align_block_size - except (ImportError, AttributeError): - logger.info("EngineX: patching moe_align_block_size → EngineX") - custom_ops.moe_align_block_size = moe_align - - # ---- Report status ---- - n_patched = 0 - for op_name in reg.ops: - backend = reg.get_backend(op_name) - if backend is not None: - n_patched += 1 - logger.info(f"EngineX bridge: {n_patched} operators registered, " - f"patched broken ops") - - -def patch_qwen3_5_moe(): - """ - Patch the MoE dispatch in qwen3_5.py to use EngineX. - - The model code tries: - 1st: corex_moe.py (not in our image) - 2nd: ixformer fused_moe (crashes on topk_softmax) - 3rd: PyTorch loop (works but slow) - - With EngineX, the topk_softmax fallback prevents the crash, - so tier 2 (ixformer) works for the GEMM even though topk - is handled by our PyTorch replacement. - """ - from enginex.dispatch.registry import get_registry - reg = get_registry() - reg.probe() - - # The actual patching happens through _custom_ops - # Since qwen3_5.py calls ops.topk_softmax(), which calls _custom_ops, - # patching _custom_ops is sufficient. - logger.info("EngineX: MoE dispatch chain patched via _custom_ops bridge") diff --git a/enginex/dispatch/__init__.py b/enginex/dispatch/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/enginex/dispatch/registry.py b/enginex/dispatch/registry.py deleted file mode 100644 index 2be3ab68..00000000 --- a/enginex/dispatch/registry.py +++ /dev/null @@ -1,474 +0,0 @@ -""" -EngineX Operator Registry — CCCL-style policy_selector for BI-V100 - -Maps each operator to its best available implementation: - Tier 1: Native .so via ctypes/dlopen (libcorex_gdn.so, libixattn.so, etc.) - Tier 2: ixformer Python ops (ixf_F.silu_and_mul, etc.) - Tier 3: PyTorch fallback (torch.nn.functional, manual loops) - -Modeled after CCCL dispatch_reduce.cuh → PolicySelector → tuning_reduce.cuh chain: - CCCL picks {threads, items, algorithm} per SM arch - We pick {backend, tile_size, num_warps} per BI-V100 hardware constraints -""" - -import ctypes -import logging -import os -from dataclasses import dataclass, field -from enum import IntEnum -from pathlib import Path -from typing import Callable, Dict, List, Optional, Tuple - -import torch - -logger = logging.getLogger("enginex.registry") - - -class Backend(IntEnum): - """Dispatch tiers — same ordering as CCCL's dispatch priority.""" - NATIVE_SO = 0 # dlopen .so — fastest, hardware-fused - IXFORMER = 1 # ixformer Python ops — vendor-provided - PYTORCH = 2 # torch fallback — slowest but always works - - -@dataclass -class HardwareProfile: - """BI-V100 hardware constants (from HARDWARE_PROBE_20260808.md).""" - sm_count: int = 16 - smem_per_sm: int = 49152 # 48KB confirmed - max_threads_per_block: int = 1024 - warp_size: int = 32 - mem_bandwidth_gbps: float = 900.0 - per_sm_bandwidth_gbps: float = 56.25 # 900/16 - compute_capability: str = "bi_v100" - cuda_version: str = "10.2" - driver_version: str = "3.2.1" - - -@dataclass -class OperatorImpl: - """A single implementation of an operator.""" - name: str - backend: Backend - fn: Optional[Callable] = None - so_path: Optional[str] = None - available: bool = False - load_error: Optional[str] = None - - -@dataclass -class OperatorEntry: - """ - One logical operator with multiple implementations. - Mirrors CCCL's policy_selector: each entry has a chain of candidates - sorted by priority. At dispatch time, we pick the first available. - """ - op_name: str - impls: List[OperatorImpl] = field(default_factory=list) - active: Optional[OperatorImpl] = None - - def select_best(self) -> Optional[OperatorImpl]: - """Pick first available impl (lowest Backend enum = highest priority).""" - for impl in sorted(self.impls, key=lambda x: x.backend): - if impl.available: - self.active = impl - return impl - return None - - -# --------------------------------------------------------------------------- -# .so probe paths — where to look for native kernels -# --------------------------------------------------------------------------- -_SO_SEARCH_PATHS = [ - "/usr/local/corex/lib64", - "/usr/local/corex/lib", - "/workspace/enginex/lib", - "/home/claude/project_6/enginex/lib", -] - - -def _probe_so(name: str) -> Optional[str]: - """Try to find a .so file by name in known search paths.""" - for d in _SO_SEARCH_PATHS: - p = os.path.join(d, name) - if os.path.isfile(p): - return p - return None - - -def _try_dlopen(path: str) -> Tuple[Optional[ctypes.CDLL], Optional[str]]: - """Attempt dlopen, return (handle, error_or_None).""" - try: - handle = ctypes.CDLL(path) - return handle, None - except OSError as e: - return None, str(e) - - -def _try_import_ixformer(): - """Probe ixformer availability.""" - try: - import ixformer.functions as ixf_F - return ixf_F, None - except ImportError as e: - return None, str(e) - - -# --------------------------------------------------------------------------- -# The global registry -# --------------------------------------------------------------------------- -class OperatorRegistry: - """ - Central operator registry — the EngineX equivalent of CCCL's - DeviceReducePolicy / DeviceScanPolicy / DeviceTopkPolicy system. - - Usage: - reg = get_registry() - moe_topk = reg.get_op("moe_topk_softmax") - if moe_topk: - moe_topk(topk_weights, topk_ids, token_expert_indices, gating_output) - """ - - def __init__(self): - self.hw = HardwareProfile() - self.ops: Dict[str, OperatorEntry] = {} - self.ixf_F = None - self._probed = False - - def probe(self): - """ - One-time hardware + library probe. - Called automatically on first get_op(). - """ - if self._probed: - return - self._probed = True - - logger.info(f"EngineX probe: SM={self.hw.sm_count}, " - f"SMEM={self.hw.smem_per_sm}, " - f"BW={self.hw.mem_bandwidth_gbps} GB/s") - - # Probe ixformer - self.ixf_F, ixf_err = _try_import_ixformer() - if self.ixf_F: - logger.info("EngineX: ixformer.functions available") - else: - logger.warning(f"EngineX: ixformer not available: {ixf_err}") - - # Register all operators - self._register_gdn_ops() - self._register_moe_ops() - self._register_fa2_ops() - self._register_activation_ops() - self._register_norm_ops() - self._register_attention_ops() - self._register_cache_ops() - self._register_sampling_ops() - - # Select best impl for each - for name, entry in self.ops.items(): - best = entry.select_best() - if best: - logger.info(f"EngineX [{name}]: using {best.backend.name} " - f"({best.name})") - else: - logger.error(f"EngineX [{name}]: NO IMPLEMENTATION AVAILABLE") - - def _add_op(self, op_name: str, impl: OperatorImpl): - if op_name not in self.ops: - self.ops[op_name] = OperatorEntry(op_name=op_name) - self.ops[op_name].impls.append(impl) - - def get_op(self, name: str) -> Optional[Callable]: - """Get the best available implementation for an operator.""" - self.probe() - entry = self.ops.get(name) - if entry and entry.active and entry.active.fn: - return entry.active.fn - return None - - def get_backend(self, name: str) -> Optional[Backend]: - """Get which backend is active for an operator.""" - self.probe() - entry = self.ops.get(name) - if entry and entry.active: - return entry.active.backend - return None - - # ------------------------------------------------------------------ - # GDN (GatedDeltaNet) operators - # From log: corex_gdn.py:56 loads /usr/local/corex/lib64/libcorex_gdn.so - # ------------------------------------------------------------------ - def _register_gdn_ops(self): - # Tier 1: native .so - so_path = _probe_so("libcorex_gdn.so") - if so_path: - handle, err = _try_dlopen(so_path) - from enginex.ops.gdn import make_native_gdn_decode, make_native_gdn_prefill - self._add_op("gdn_decode", OperatorImpl( - name="corex_gdn_decode", backend=Backend.NATIVE_SO, - fn=make_native_gdn_decode(handle) if handle else None, - so_path=so_path, available=handle is not None, - load_error=err)) - self._add_op("gdn_prefill", OperatorImpl( - name="corex_gdn_prefill", backend=Backend.NATIVE_SO, - fn=make_native_gdn_prefill(handle) if handle else None, - so_path=so_path, available=handle is not None, - load_error=err)) - # Tier 2: our FlashQLA SM70 kernel (.so compiled from gdn_forward.cu) - flash_so = _probe_so("flash_qla_sm70_gdn_strided.so") - if not flash_so: - # Check build directory - for d in ["/workspace/qwen3_6_scripts/flash_qla_sm70/build", - "/home/claude/project_6/qwen3_6_scripts/flash_qla_sm70/build"]: - candidate = os.path.join(d, "flash_qla_sm70_gdn_strided.so") - if os.path.isfile(candidate): - flash_so = candidate - break - if flash_so: - from enginex.ops.gdn import make_flashqla_gdn_prefill - handle, err = _try_dlopen(flash_so) - self._add_op("gdn_prefill", OperatorImpl( - name="flashqla_sm70_prefill", backend=Backend.NATIVE_SO, - fn=make_flashqla_gdn_prefill(flash_so) if handle else None, - so_path=flash_so, available=handle is not None, - load_error=err)) - - # Tier 3: PyTorch fallback - from enginex.ops.gdn import gdn_decode_pytorch, gdn_prefill_pytorch - self._add_op("gdn_decode", OperatorImpl( - name="pytorch_gdn_decode", backend=Backend.PYTORCH, - fn=gdn_decode_pytorch, available=True)) - self._add_op("gdn_prefill", OperatorImpl( - name="pytorch_gdn_prefill", backend=Backend.PYTORCH, - fn=gdn_prefill_pytorch, available=True)) - - # ------------------------------------------------------------------ - # MoE operators - # From log: vllm_moe_topk_softmax missing from ixformer.functions - # corex_moe.py:339 uses expert-grouped-wmma kernel - # ------------------------------------------------------------------ - def _register_moe_ops(self): - # Tier 2: ixformer (but topk_softmax is KNOWN MISSING) - if self.ixf_F: - has_topk = hasattr(self.ixf_F, 'vllm_moe_topk_softmax') - has_fused = hasattr(self.ixf_F, 'vllm_invoke_fused_moe_kernel') - has_align = hasattr(self.ixf_F, 'vllm_moe_align_block_size') - - if has_topk: - self._add_op("moe_topk_softmax", OperatorImpl( - name="ixf_moe_topk_softmax", backend=Backend.IXFORMER, - fn=self.ixf_F.vllm_moe_topk_softmax, available=True)) - if has_fused: - self._add_op("moe_fused_kernel", OperatorImpl( - name="ixf_fused_moe", backend=Backend.IXFORMER, - fn=self.ixf_F.vllm_invoke_fused_moe_kernel, available=True)) - if has_align: - self._add_op("moe_align_block_size", OperatorImpl( - name="ixf_moe_align", backend=Backend.IXFORMER, - fn=self.ixf_F.vllm_moe_align_block_size, available=True)) - - # Tier 3: PyTorch fallback (THE FIX for the topk_softmax crash) - from enginex.ops.moe import (moe_topk_softmax_pytorch, - moe_fused_kernel_pytorch, - moe_align_block_size_pytorch) - self._add_op("moe_topk_softmax", OperatorImpl( - name="pytorch_moe_topk_softmax", backend=Backend.PYTORCH, - fn=moe_topk_softmax_pytorch, available=True)) - self._add_op("moe_fused_kernel", OperatorImpl( - name="pytorch_fused_moe", backend=Backend.PYTORCH, - fn=moe_fused_kernel_pytorch, available=True)) - self._add_op("moe_align_block_size", OperatorImpl( - name="pytorch_moe_align", backend=Backend.PYTORCH, - fn=moe_align_block_size_pytorch, available=True)) - - # ------------------------------------------------------------------ - # FA2 (Flash Attention 2) operators - # From log: corex_fa2.py:333 "CoreX FA2 packed prefill" B=2 Hq=4 Hkv=1 D=256 - # corex_fa2.py:507 "CoreX paged FA2 chunked prefill" - # ------------------------------------------------------------------ - def _register_fa2_ops(self): - # Tier 2: ixformer flash_attn - if self.ixf_F: - import ixformer - has_fa = hasattr(ixformer, 'flash_attn_varlen_func') - has_fa_pad = hasattr(ixformer, 'flash_attn_func') - if has_fa: - self._add_op("fa2_varlen", OperatorImpl( - name="ixf_flash_attn_varlen", backend=Backend.IXFORMER, - fn=ixformer.flash_attn_varlen_func, available=True)) - if has_fa_pad: - self._add_op("fa2_padded", OperatorImpl( - name="ixf_flash_attn_padded", backend=Backend.IXFORMER, - fn=ixformer.flash_attn_func, available=True)) - - # Tier 2: libixattn.so (confirmed present in hardware probe) - ixattn_so = _probe_so("libixattn.so") - if ixattn_so: - handle, err = _try_dlopen(ixattn_so) - self._add_op("fa2_native", OperatorImpl( - name="libixattn", backend=Backend.NATIVE_SO, - so_path=ixattn_so, available=handle is not None, - load_error=err)) - - # Tier 3: xformers SDPA fallback (what we currently use) - from enginex.ops.attention import fa2_xformers_fallback - self._add_op("fa2_varlen", OperatorImpl( - name="xformers_sdpa_fallback", backend=Backend.PYTORCH, - fn=fa2_xformers_fallback, available=True)) - self._add_op("fa2_padded", OperatorImpl( - name="xformers_sdpa_fallback", backend=Backend.PYTORCH, - fn=fa2_xformers_fallback, available=True)) - - # ------------------------------------------------------------------ - # Activation ops (silu_and_mul, gelu, etc.) - # These work via ixformer — confirmed in hardware probe - # ------------------------------------------------------------------ - def _register_activation_ops(self): - if self.ixf_F: - for op_name, ixf_name in [ - ("silu_and_mul", "silu_and_mul"), - ("gelu_and_mul", "gelu_and_mul"), - ("gelu_tanh_and_mul", "gelu_tanh_and_mul"), - ]: - fn = getattr(self.ixf_F, ixf_name, None) - if fn: - self._add_op(op_name, OperatorImpl( - name=f"ixf_{ixf_name}", backend=Backend.IXFORMER, - fn=fn, available=True)) - - # PyTorch fallbacks - from enginex.ops.activations import (silu_and_mul_pytorch, - gelu_and_mul_pytorch, - gelu_tanh_and_mul_pytorch) - self._add_op("silu_and_mul", OperatorImpl( - name="pytorch_silu_and_mul", backend=Backend.PYTORCH, - fn=silu_and_mul_pytorch, available=True)) - self._add_op("gelu_and_mul", OperatorImpl( - name="pytorch_gelu_and_mul", backend=Backend.PYTORCH, - fn=gelu_and_mul_pytorch, available=True)) - self._add_op("gelu_tanh_and_mul", OperatorImpl( - name="pytorch_gelu_tanh_and_mul", backend=Backend.PYTORCH, - fn=gelu_tanh_and_mul_pytorch, available=True)) - - # ------------------------------------------------------------------ - # Norm ops (rms_norm, fused_add_rms_norm) - # ------------------------------------------------------------------ - def _register_norm_ops(self): - if self.ixf_F: - for op_name, ixf_name in [ - ("rms_norm", "rms_norm"), - ("fused_add_rms_norm", "fused_add_rms_norm"), - ]: - fn = getattr(self.ixf_F, ixf_name, None) - if fn: - self._add_op(op_name, OperatorImpl( - name=f"ixf_{ixf_name}", backend=Backend.IXFORMER, - fn=fn, available=True)) - - from enginex.ops.norm import rms_norm_pytorch, fused_add_rms_norm_pytorch - self._add_op("rms_norm", OperatorImpl( - name="pytorch_rms_norm", backend=Backend.PYTORCH, - fn=rms_norm_pytorch, available=True)) - self._add_op("fused_add_rms_norm", OperatorImpl( - name="pytorch_fused_add_rms_norm", backend=Backend.PYTORCH, - fn=fused_add_rms_norm_pytorch, available=True)) - - # ------------------------------------------------------------------ - # Paged attention ops - # ------------------------------------------------------------------ - def _register_attention_ops(self): - if self.ixf_F: - fn_v1 = getattr(self.ixf_F, - 'vllm_single_query_cached_kv_attention', None) - fn_v2 = getattr(self.ixf_F, - 'vllm_single_query_cached_kv_attention_v2', None) - if fn_v1: - self._add_op("paged_attention_v1", OperatorImpl( - name="ixf_paged_attn_v1", backend=Backend.IXFORMER, - fn=fn_v1, available=True)) - if fn_v2: - self._add_op("paged_attention_v2", OperatorImpl( - name="ixf_paged_attn_v2", backend=Backend.IXFORMER, - fn=fn_v2, available=True)) - - from enginex.ops.attention import (paged_attention_v1_pytorch, - paged_attention_v2_pytorch) - self._add_op("paged_attention_v1", OperatorImpl( - name="pytorch_paged_attn_v1", backend=Backend.PYTORCH, - fn=paged_attention_v1_pytorch, available=True)) - self._add_op("paged_attention_v2", OperatorImpl( - name="pytorch_paged_attn_v2", backend=Backend.PYTORCH, - fn=paged_attention_v2_pytorch, available=True)) - - # ------------------------------------------------------------------ - # Cache ops (reshape_and_cache, copy_blocks, swap_blocks) - # ------------------------------------------------------------------ - def _register_cache_ops(self): - if self.ixf_F: - for op_name, ixf_name in [ - ("reshape_and_cache", "vllm_cache_ops_reshape_and_cache"), - ("copy_blocks", "copy_blocks"), - ("swap_blocks", "swap_blocks"), - ]: - fn = getattr(self.ixf_F, ixf_name, None) - if fn: - self._add_op(op_name, OperatorImpl( - name=f"ixf_{ixf_name}", backend=Backend.IXFORMER, - fn=fn, available=True)) - - from enginex.ops.cache import (reshape_and_cache_pytorch, - copy_blocks_pytorch, - swap_blocks_pytorch) - self._add_op("reshape_and_cache", OperatorImpl( - name="pytorch_reshape_cache", backend=Backend.PYTORCH, - fn=reshape_and_cache_pytorch, available=True)) - self._add_op("copy_blocks", OperatorImpl( - name="pytorch_copy_blocks", backend=Backend.PYTORCH, - fn=copy_blocks_pytorch, available=True)) - self._add_op("swap_blocks", OperatorImpl( - name="pytorch_swap_blocks", backend=Backend.PYTORCH, - fn=swap_blocks_pytorch, available=True)) - - # ------------------------------------------------------------------ - # Sampling ops (rotary_embedding, topk) - # ------------------------------------------------------------------ - def _register_sampling_ops(self): - if self.ixf_F: - fn = getattr(self.ixf_F, 'vllm_rotary_embedding_neox', None) - if fn: - self._add_op("rotary_embedding", OperatorImpl( - name="ixf_rotary", backend=Backend.IXFORMER, - fn=fn, available=True)) - - from enginex.ops.sampling import rotary_embedding_pytorch - self._add_op("rotary_embedding", OperatorImpl( - name="pytorch_rotary", backend=Backend.PYTORCH, - fn=rotary_embedding_pytorch, available=True)) - - def summary(self) -> str: - """Print a summary of all operators and their active backends.""" - self.probe() - lines = ["EngineX Operator Registry Summary", - "=" * 50] - for name, entry in sorted(self.ops.items()): - active = entry.active - if active: - lines.append( - f" {name:30s} → {active.backend.name:12s} ({active.name})") - else: - lines.append(f" {name:30s} → MISSING") - return "\n".join(lines) - - -# --------------------------------------------------------------------------- -# Singleton -# --------------------------------------------------------------------------- -_global_registry: Optional[OperatorRegistry] = None - - -def get_registry() -> OperatorRegistry: - global _global_registry - if _global_registry is None: - _global_registry = OperatorRegistry() - return _global_registry diff --git a/enginex/ops/__init__.py b/enginex/ops/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/enginex/ops/activations.py b/enginex/ops/activations.py deleted file mode 100644 index 775bd1b3..00000000 --- a/enginex/ops/activations.py +++ /dev/null @@ -1,39 +0,0 @@ -""" -EngineX activation operators. - -These map to CCCL's dispatch_transform pattern — element-wise kernels -that fuse activation + multiply in a single pass. - -ixformer provides these natively (confirmed working in hardware probe). -PyTorch fallbacks here for completeness. - -CCCL tuning: tuning_transform.cuh bytes_in_flight = 64KB on BI-V100 - (56 GB/s per-SM × 1100ns latency, 16 SMs) -""" - -import torch -import torch.nn.functional as F - - -def silu_and_mul_pytorch(x: torch.Tensor, out: torch.Tensor) -> None: - """Fused SiLU(x[..., :d]) * x[..., d:]""" - d = x.shape[-1] // 2 - gate = x[..., :d] - up = x[..., d:] - out.copy_(F.silu(gate) * up) - - -def gelu_and_mul_pytorch(x: torch.Tensor, out: torch.Tensor) -> None: - """Fused GELU(x[..., :d]) * x[..., d:]""" - d = x.shape[-1] // 2 - gate = x[..., :d] - up = x[..., d:] - out.copy_(F.gelu(gate) * up) - - -def gelu_tanh_and_mul_pytorch(x: torch.Tensor, out: torch.Tensor) -> None: - """Fused GELU_tanh(x[..., :d]) * x[..., d:]""" - d = x.shape[-1] // 2 - gate = x[..., :d] - up = x[..., d:] - out.copy_(F.gelu(gate, approximate='tanh') * up) diff --git a/enginex/ops/attention.py b/enginex/ops/attention.py deleted file mode 100644 index cdc911ab..00000000 --- a/enginex/ops/attention.py +++ /dev/null @@ -1,182 +0,0 @@ -""" -EngineX Attention operators. - -Sub168 log shows three attention paths: - 1. CoreX FA2 packed prefill: B=2 Hq=4 Hkv=1 D=256 (full attention layers) - 2. CoreX paged FA2 chunked prefill: B=1 Hq=4 Hkv=1 D=256 cache_blocks=2 - 3. CoreX GDN (handled in gdn.py, 4 of 36 layers) - -Our image has: - - libixattn.so (present but not wired) - - ixformer.flash_attn_varlen_func (available) - - xformers SDPA (current fallback, patched for head_dim=256) - -CCCL parallel: - paged_attention_v1 = dispatch_reduce (reduce over KV blocks) - paged_attention_v2 = dispatch_reduce two-pass (partition-level reduce + final reduce) -""" - -import math -from typing import List, Optional - -import torch -import torch.nn.functional as F - - -def fa2_xformers_fallback( - query: torch.Tensor, - key: torch.Tensor, - value: torch.Tensor, - cu_seqlens_q: Optional[torch.Tensor] = None, - cu_seqlens_k: Optional[torch.Tensor] = None, - max_seqlen_q: int = 0, - max_seqlen_k: int = 0, - dropout_p: float = 0.0, - softmax_scale: Optional[float] = None, - causal: bool = False, -) -> torch.Tensor: - """ - xformers SDPA fallback for FA2. - This is what we currently use — works but slower than native FA2. - Head_dim=256 bypass already applied in patch_xformers_sdpa_*.py. - """ - if softmax_scale is None: - softmax_scale = 1.0 / math.sqrt(query.shape[-1]) - - # Standard scaled dot product attention - attn_weights = torch.matmul(query, key.transpose(-2, -1)) * softmax_scale - - if causal and attn_weights.shape[-2] > 1: - L = attn_weights.shape[-2] - S = attn_weights.shape[-1] - mask = torch.triu( - torch.full((L, S), float('-inf'), device=query.device), - diagonal=S - L + 1 - ) - attn_weights = attn_weights + mask - - attn_weights = F.softmax(attn_weights, dim=-1) - output = torch.matmul(attn_weights, value) - return output - - -def paged_attention_v1_pytorch( - output: torch.Tensor, # [num_seqs, num_heads, head_size] - query: torch.Tensor, # [num_seqs, num_heads, head_size] - key_cache: torch.Tensor, # [num_blocks, num_kv_heads, block_size, head_size] - value_cache: torch.Tensor, # [num_blocks, num_kv_heads, block_size, head_size] - num_kv_heads: int, - scale: float, - block_tables: torch.Tensor, # [num_seqs, max_blocks_per_seq] - seq_lens: torch.Tensor, # [num_seqs] - block_size: int, - max_seq_len: int, - alibi_slopes: Optional[torch.Tensor] = None, - kv_cache_dtype: str = "auto", - k_scale: float = 1.0, - v_scale: float = 1.0, - tp_rank: int = 0, - blocksparse_local_blocks: int = 0, - blocksparse_vert_stride: int = 0, - blocksparse_block_size: int = 64, - blocksparse_head_sliding_step: int = 0, -) -> None: - """ - Paged attention v1 — single-pass reduce over all KV blocks. - - CCCL parallel: dispatch_reduce single-tile kernel. - For short sequences (< 2 × sm_count × partition_size), v1 is faster - because it avoids the two-pass overhead. - - BI-V100 with 16 SMs: threshold ≈ 16 × 2 × 512 = 16384 tokens. - """ - num_seqs = query.shape[0] - num_heads = query.shape[1] - head_size = query.shape[2] - num_queries_per_kv = num_heads // num_kv_heads - - for seq_idx in range(num_seqs): - seq_len = seq_lens[seq_idx].item() - if seq_len == 0: - continue - - q = query[seq_idx] # [num_heads, head_size] - - num_blocks = (seq_len + block_size - 1) // block_size - keys_list = [] - values_list = [] - - for block_idx in range(num_blocks): - physical_block = block_tables[seq_idx, block_idx].item() - if block_idx == num_blocks - 1: - # Last block may be partial - tokens_in_block = seq_len - block_idx * block_size - else: - tokens_in_block = block_size - - k_block = key_cache[physical_block, :, :tokens_in_block, :] - v_block = value_cache[physical_block, :, :tokens_in_block, :] - keys_list.append(k_block) - values_list.append(v_block) - - # Concatenate all KV - all_keys = torch.cat(keys_list, dim=1) # [num_kv_heads, seq_len, head_size] - all_values = torch.cat(values_list, dim=1) - - # GQA: repeat KV heads - if num_queries_per_kv > 1: - all_keys = all_keys.repeat_interleave(num_queries_per_kv, dim=0) - all_values = all_values.repeat_interleave(num_queries_per_kv, dim=0) - - # Attention: q @ k^T → softmax → @ v - attn = torch.einsum('hd,hsd->hs', q, all_keys) * scale - attn = F.softmax(attn, dim=-1) - out = torch.einsum('hs,hsd->hd', attn, all_values) - - output[seq_idx].copy_(out) - - -def paged_attention_v2_pytorch( - output: torch.Tensor, - exp_sums: torch.Tensor, # [num_seqs, num_heads, max_partitions] - max_logits: torch.Tensor, # [num_seqs, num_heads, max_partitions] - tmp_output: torch.Tensor, # [num_seqs, num_heads, max_partitions, head_size] - query: torch.Tensor, - key_cache: torch.Tensor, - value_cache: torch.Tensor, - num_kv_heads: int, - scale: float, - block_tables: torch.Tensor, - seq_lens: torch.Tensor, - block_size: int, - max_seq_len: int, - alibi_slopes: Optional[torch.Tensor] = None, - kv_cache_dtype: str = "auto", - k_scale: float = 1.0, - v_scale: float = 1.0, - tp_rank: int = 0, - blocksparse_local_blocks: int = 0, - blocksparse_vert_stride: int = 0, - blocksparse_block_size: int = 64, - blocksparse_head_sliding_step: int = 0, -) -> None: - """ - Paged attention v2 — two-pass reduce with partitioning. - - CCCL parallel: dispatch_reduce two-pass pattern. - Pass 1: per-partition reduce (each partition = PARTITION_SIZE KV tokens) - Pass 2: reduce across partitions (log-sum-exp correction) - - For BI-V100 with 16 SMs, v2 is better when seq_len > 8192 (multiple - waves of partitions keep all SMs busy). - """ - # For correctness, delegate to v1 — the two-pass optimization - # only matters for perf on long sequences - paged_attention_v1_pytorch( - output, query, key_cache, value_cache, - num_kv_heads, scale, block_tables, seq_lens, - block_size, max_seq_len, alibi_slopes, kv_cache_dtype, - k_scale, v_scale, tp_rank, - blocksparse_local_blocks, blocksparse_vert_stride, - blocksparse_block_size, blocksparse_head_sliding_step, - ) diff --git a/enginex/ops/cache.py b/enginex/ops/cache.py deleted file mode 100644 index 80905962..00000000 --- a/enginex/ops/cache.py +++ /dev/null @@ -1,63 +0,0 @@ -""" -EngineX cache operators. - -KV cache management for paged attention. -CCCL parallel: dispatch_batch_memcpy (block copies between cache slots). -""" - -from typing import Dict, List - -import torch - - -def reshape_and_cache_pytorch( - key: torch.Tensor, # [num_tokens, num_kv_heads, head_size] - value: torch.Tensor, # [num_tokens, num_kv_heads, head_size] - key_cache: torch.Tensor, # [num_blocks, num_kv_heads, block_size, head_size] - value_cache: torch.Tensor, # [num_blocks, num_kv_heads, block_size, head_size] - slot_mapping: torch.Tensor, # [num_tokens] — maps token → (block, offset) - kv_cache_dtype: str = "auto", - k_scale: float = 1.0, - v_scale: float = 1.0, -) -> None: - """Write new K,V into their assigned cache slots.""" - num_tokens = key.shape[0] - block_size = key_cache.shape[2] - - for i in range(num_tokens): - slot = slot_mapping[i].item() - if slot < 0: - continue - block_idx = slot // block_size - block_offset = slot % block_size - key_cache[block_idx, :, block_offset, :] = key[i] * k_scale - value_cache[block_idx, :, block_offset, :] = value[i] * v_scale - - -def copy_blocks_pytorch( - key_caches: List[torch.Tensor], - value_caches: List[torch.Tensor], - block_mapping: torch.Tensor, # [num_pairs, 2] src→dst -) -> None: - """Copy cache blocks (used for fork/copy-on-write).""" - num_pairs = block_mapping.shape[0] - num_layers = len(key_caches) - - for i in range(num_pairs): - src = block_mapping[i, 0].item() - dst = block_mapping[i, 1].item() - for layer in range(num_layers): - key_caches[layer][dst].copy_(key_caches[layer][src]) - value_caches[layer][dst].copy_(value_caches[layer][src]) - - -def swap_blocks_pytorch( - src: torch.Tensor, - dst: torch.Tensor, - block_mapping: torch.Tensor, -) -> None: - """Swap cache blocks between GPU and CPU.""" - for i in range(block_mapping.shape[0]): - src_idx = block_mapping[i, 0].item() - dst_idx = block_mapping[i, 1].item() - dst[dst_idx].copy_(src[src_idx]) diff --git a/enginex/ops/gdn.py b/enginex/ops/gdn.py deleted file mode 100644 index 9decf522..00000000 --- a/enginex/ops/gdn.py +++ /dev/null @@ -1,167 +0,0 @@ -""" -EngineX GDN (GatedDeltaNet) operators. - -From docker log: - Sub168 (working): corex_gdn.py:56 Loaded fused CoreX GDN decode from libcorex_gdn.so - Our run (broken): qwen3_5.py:445 NaN in prefill GatedDeltaNet layer 0 (frac=0.9998) - -The GDN is a linear attention variant with gated delta rule updates. -4 of 36 attention layers use GDN instead of full attention. - -Two paths: - - Prefill: chunked computation (L tokens split into chunks of C) - - Decode: single-step recurrent update (state @ query) - -CCCL parallel: maps to dispatch_scan pattern (state accumulation = prefix scan). -""" - -import ctypes -import logging -import math -from typing import Optional, Tuple - -import torch -import torch.nn.functional as F - -logger = logging.getLogger("enginex.ops.gdn") - - -# --------------------------------------------------------------------------- -# Tier 1: Native .so wrappers (dlopen libcorex_gdn.so) -# --------------------------------------------------------------------------- -def make_native_gdn_decode(handle: ctypes.CDLL): - """Wrap the native CoreX GDN decode operator loaded from .so.""" - # The actual C function signature would be discovered at integration time. - # For now, this is a placeholder that logs the call. - def native_gdn_decode(q, k, v, gate, beta, conv_state, temporal_state): - logger.debug("native_gdn_decode called via libcorex_gdn.so") - # Would call handle.corex_gdn_decode_forward(...) - raise NotImplementedError("Native .so integration pending on-device testing") - return native_gdn_decode - - -def make_native_gdn_prefill(handle: ctypes.CDLL): - """Wrap the native CoreX GDN prefill operator.""" - def native_gdn_prefill(q, k, v, gate, beta, state, chunk_size=64): - logger.debug("native_gdn_prefill called via libcorex_gdn.so") - raise NotImplementedError("Native .so integration pending on-device testing") - return native_gdn_prefill - - -def make_flashqla_gdn_prefill(so_path: str): - """Wrap our compiled FlashQLA SM70 kernel (gdn_forward.cu → .so).""" - def flashqla_prefill(q, k, v, gate, beta, state, chunk_size=64): - # This calls the JIT-compiled .so from flash_qla_sm70/ - try: - from qwen3_6_scripts.flash_qla_sm70 import chunk_gated_delta_rule_fwd_sm70 - return chunk_gated_delta_rule_fwd_sm70(q, k, v, gate, beta, state) - except ImportError: - logger.warning("FlashQLA SM70 not importable, falling back to PyTorch") - return gdn_prefill_pytorch(q, k, v, gate, beta, state, chunk_size) - return flashqla_prefill - - -# --------------------------------------------------------------------------- -# Tier 3: PyTorch fallback with numerical stability fixes -# --------------------------------------------------------------------------- -def gdn_decode_pytorch( - q: torch.Tensor, # [B, H, D] - k: torch.Tensor, # [B, H, D] - v: torch.Tensor, # [B, H, D] - gate: torch.Tensor, # [B, H] — gate (sigmoid applied externally) - beta: torch.Tensor, # [B, H] — delta rule learning rate - conv_state: torch.Tensor, # [B, H, conv_width, D] — causal conv1d state - temporal_state: torch.Tensor, # [B, H, D, D] — recurrent state -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Single-step recurrent GDN decode. - - The delta rule update: S' = gate * S + beta * (k^T @ v) - Output: o = S' @ q - - CCCL parallel: single-element "scan" — just the recurrent update. - """ - B, H, D = q.shape - - # Delta rule: state decay + write - # gate controls how much old state to retain - # beta controls how much new (k,v) pair to inject - kv_outer = torch.einsum('bhd,bhe->bhde', k, v) # [B, H, D, D] - - # Clamp to prevent NaN propagation (the fix for 99.98% NaN) - gate_expanded = gate.unsqueeze(-1).unsqueeze(-1).clamp(-5.0, 5.0) - beta_expanded = beta.unsqueeze(-1).unsqueeze(-1).clamp(-5.0, 5.0) - - # State update - new_state = gate_expanded * temporal_state + beta_expanded * kv_outer - - # Clamp state to prevent NaN accumulation across layers - new_state = new_state.clamp(-1e4, 1e4) - - # Output = state @ query - output = torch.einsum('bhde,bhd->bhe', new_state, q) # [B, H, D] - - return output, new_state - - -def gdn_prefill_pytorch( - q: torch.Tensor, # [1, L, H, D] - k: torch.Tensor, # [1, L, H, D] - v: torch.Tensor, # [1, L, H, D] - gate: torch.Tensor, # [1, L, H] - beta: torch.Tensor, # [1, L, H] - state: torch.Tensor, # [B, H, D, D] initial state - chunk_size: int = 16, # Reduced from 64→16 per CCCL overflow fix -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - Chunked GDN prefill — processes L tokens in chunks of chunk_size. - - This is the numerically-stable version that prevents the 99.98% NaN issue. - Key fixes applied: - 1. chunk_size 64→16 (fewer cumsum steps = less overflow) - 2. Clamp gate/beta before exp/cumsum - 3. Clamp state after each chunk - - CCCL parallel: maps to dispatch_scan two-phase pattern: - Phase 1: per-chunk local scan (intra-chunk attention) - Phase 2: cross-chunk state propagation (lookback) - """ - B, L, H, D = q.shape - - outputs = [] - current_state = state.clone() - - for start in range(0, L, chunk_size): - end = min(start + chunk_size, L) - C = end - start - - q_chunk = q[:, start:end] # [B, C, H, D] - k_chunk = k[:, start:end] - v_chunk = v[:, start:end] - g_chunk = gate[:, start:end].clamp(-5.0, 5.0) # [B, C, H] - b_chunk = beta[:, start:end].clamp(-5.0, 5.0) - - chunk_out = torch.zeros_like(q_chunk) - - # Intra-chunk: causal attention with delta rule - for t in range(C): - qt = q_chunk[:, t] # [B, H, D] - kt = k_chunk[:, t] - vt = v_chunk[:, t] - gt = g_chunk[:, t].unsqueeze(-1).unsqueeze(-1) # [B, H, 1, 1] - bt = b_chunk[:, t].unsqueeze(-1).unsqueeze(-1) - - kv_outer = torch.einsum('bhd,bhe->bhde', kt, vt) - - # Delta rule state update - current_state = gt * current_state + bt * kv_outer - current_state = current_state.clamp(-1e4, 1e4) - - # Query against state - ot = torch.einsum('bhde,bhd->bhe', current_state, qt) - chunk_out[:, t] = ot - - outputs.append(chunk_out) - - output = torch.cat(outputs, dim=1) # [B, L, H, D] - return output, current_state diff --git a/enginex/ops/moe.py b/enginex/ops/moe.py deleted file mode 100644 index d477de06..00000000 --- a/enginex/ops/moe.py +++ /dev/null @@ -1,201 +0,0 @@ -""" -EngineX MoE operators — replacements for missing ixformer MoE functions. - -From docker log (comp 168): - ERROR _custom_ops.py:58] module 'ixformer.functions' has no attribute 'vllm_moe_topk_softmax' - WARNING qwen3_5.py:913] FusedMoE native kernel failed, falling back to pure PyTorch - -This fires on EVERY MoE layer (36 per token), 4 workers = 144 error lines per forward pass. - -Three operators needed: - 1. moe_topk_softmax — gate logits → softmax → topk expert selection - 2. moe_fused_kernel — the actual expert GEMM dispatch - 3. moe_align_block_size — pad expert assignments to block boundaries -""" - -import torch -import torch.nn.functional as F - - -def moe_topk_softmax_pytorch( - topk_weights: torch.Tensor, # [num_tokens, topk] output - topk_ids: torch.Tensor, # [num_tokens, topk] output - token_expert_indices: torch.Tensor, # [num_tokens, topk] output - gating_output: torch.Tensor, # [num_tokens, num_experts] input -) -> None: - """ - Replacement for ixf_F.vllm_moe_topk_softmax. - - Computes softmax over expert gating logits, selects top-k experts per token. - This is the router in Qwen3.5's MoE layer (256 experts, topk=8). - - CCCL parallel: maps to tuning_batched_topk.cuh worker_policy pattern — - each token is a "segment", we find top-k within each segment. - """ - num_tokens = gating_output.shape[0] - topk = topk_weights.shape[1] - - # Softmax over experts (dim=-1) - probs = F.softmax(gating_output, dim=-1) - - # Top-k selection per token - weights, ids = torch.topk(probs, k=topk, dim=-1) - - # Renormalize weights to sum to 1 - weights = weights / (weights.sum(dim=-1, keepdim=True) + 1e-8) - - # Write outputs in-place (matches vllm calling convention) - topk_weights.copy_(weights) - topk_ids.copy_(ids) - - # token_expert_indices: flatten assignment for scatter - # Shape: [num_tokens, topk], value = token_idx * topk + local_expert_rank - if token_expert_indices.numel() > 0: - arange = torch.arange(num_tokens, device=gating_output.device) - token_expert_indices.copy_( - arange.unsqueeze(1) * topk + - torch.arange(topk, device=gating_output.device).unsqueeze(0) - ) - - -def moe_fused_kernel_pytorch( - hidden_states: torch.Tensor, # [num_tokens, hidden_dim] - w1: torch.Tensor, # [num_experts, hidden_dim, intermediate_dim] - w2: torch.Tensor, # [num_experts, intermediate_dim, hidden_dim] - topk_weights: torch.Tensor, # [num_tokens, topk] - topk_ids: torch.Tensor, # [num_tokens, topk] - inplace: bool = True, - override_config: dict = None, - use_fp8_w8a8: bool = False, - use_int8_w8a16: bool = False, - w1_scale: torch.Tensor = None, - w2_scale: torch.Tensor = None, - a1_scale: torch.Tensor = None, - a2_scale: torch.Tensor = None, -) -> torch.Tensor: - """ - Replacement for vllm_invoke_fused_moe_kernel. - - Dispatches tokens to their assigned experts, runs GEMM, combines results. - This is the hot inner loop — called 36 times per forward pass. - - Sub168 log shows kernel=expert-grouped-wmma, meaning the native kernel - groups tokens by expert and runs WMMA (tensor core) GEMMs. - Our fallback loops over experts — correct but slow. - - CCCL parallel: maps to dispatch_segmented_sort + dispatch_reduce pattern. - """ - num_tokens, hidden_dim = hidden_states.shape - topk = topk_ids.shape[1] - - # Group tokens by expert - # For each expert, collect which tokens use it and their weights - output = torch.zeros_like(hidden_states) - num_experts = w1.shape[0] - - for expert_idx in range(num_experts): - # Find tokens assigned to this expert - mask = (topk_ids == expert_idx) # [num_tokens, topk] - if not mask.any(): - continue - - # Get token indices and their weights for this expert - token_indices, topk_positions = mask.nonzero(as_tuple=True) - - if token_indices.numel() == 0: - continue - - weights = topk_weights[token_indices, topk_positions] # [n_assigned] - expert_input = hidden_states[token_indices] # [n_assigned, hidden_dim] - - # Expert forward: gate_up → silu → down - # w1 is [hidden_dim, intermediate_dim*2] (gate + up fused) - expert_w1 = w1[expert_idx] # [hidden_dim, intermediate_dim*2] - expert_w2 = w2[expert_idx] # [intermediate_dim, hidden_dim] - - # gate_up = input @ w1 → [n_assigned, intermediate_dim*2] - gate_up = expert_input @ expert_w1 - intermediate_dim = gate_up.shape[-1] // 2 - gate = gate_up[..., :intermediate_dim] - up = gate_up[..., intermediate_dim:] - - # SiLU(gate) * up - activated = F.silu(gate) * up - - # down = activated @ w2 - expert_output = activated @ expert_w2 # [n_assigned, hidden_dim] - - # Weighted accumulate - output.index_add_( - 0, token_indices, - expert_output * weights.unsqueeze(-1) - ) - - return output - - -def moe_align_block_size_pytorch( - topk_ids: torch.Tensor, # [num_tokens, topk] - num_experts: int, - block_size: int, - sorted_ids: torch.Tensor, # output - expert_ids: torch.Tensor, # output - num_tokens_post_pad: torch.Tensor, # output -) -> None: - """ - Replacement for ixf_F.vllm_moe_align_block_size. - - Pads expert assignments so each expert's token count is a multiple of - block_size (for efficient GEMM tiling). This is the MoE equivalent of - CCCL's dispatch_batch_memcpy tile alignment. - """ - num_tokens = topk_ids.shape[0] - topk = topk_ids.shape[1] - - # Flatten expert assignments - flat_ids = topk_ids.flatten() # [num_tokens * topk] - - # Count tokens per expert - counts = torch.zeros(num_experts, dtype=torch.int32, - device=topk_ids.device) - for e in range(num_experts): - counts[e] = (flat_ids == e).sum() - - # Pad counts to block_size multiples - padded_counts = ((counts + block_size - 1) // block_size) * block_size - total_padded = padded_counts.sum().item() - - # Sort tokens by expert, pad with dummy tokens - offsets = torch.zeros(num_experts + 1, dtype=torch.int32, - device=topk_ids.device) - offsets[1:] = torch.cumsum(padded_counts, dim=0) - - # Fill sorted_ids: real tokens first, then padding - write_pos = torch.zeros(num_experts, dtype=torch.int32, - device=topk_ids.device) - - for i in range(num_tokens * topk): - token_idx = i // topk - expert = flat_ids[i].item() - pos = offsets[expert].item() + write_pos[expert].item() - if pos < sorted_ids.numel(): - sorted_ids[pos] = token_idx - write_pos[expert] += 1 - - # Fill padding positions with 0 (dummy token) - for e in range(num_experts): - start = offsets[e].item() + counts[e].item() - end = offsets[e].item() + padded_counts[e].item() - if start < sorted_ids.numel() and end <= sorted_ids.numel(): - sorted_ids[start:end] = 0 - - # Expert ids: one per block - idx = 0 - for e in range(num_experts): - n_blocks = padded_counts[e].item() // block_size - for b in range(n_blocks): - if idx < expert_ids.numel(): - expert_ids[idx] = e - idx += 1 - - num_tokens_post_pad.fill_(total_padded) diff --git a/enginex/ops/norm.py b/enginex/ops/norm.py deleted file mode 100644 index b8857a45..00000000 --- a/enginex/ops/norm.py +++ /dev/null @@ -1,36 +0,0 @@ -""" -EngineX norm operators. - -RMSNorm is called 128 times per forward pass (pre-attn + post-attn × 64 layers). -fused_add_rms_norm fuses residual addition with normalization. - -ixformer provides both natively. Fallbacks for environments without ixformer. -""" - -import torch - - -def rms_norm_pytorch( - input: torch.Tensor, - weight: torch.Tensor, - output: torch.Tensor, - epsilon: float = 1e-6, -) -> None: - """RMSNorm: output = (input / rms(input)) * weight""" - variance = input.to(torch.float32).pow(2).mean(-1, keepdim=True) - normed = input * torch.rsqrt(variance + epsilon) - output.copy_(normed * weight) - - -def fused_add_rms_norm_pytorch( - input: torch.Tensor, - residual: torch.Tensor, - weight: torch.Tensor, - epsilon: float = 1e-6, -) -> None: - """Fused: input = RMSNorm(input + residual); residual = input + residual""" - # In-place: residual += input, then normalize - residual.add_(input) - variance = residual.to(torch.float32).pow(2).mean(-1, keepdim=True) - normed = residual * torch.rsqrt(variance + epsilon) - input.copy_(normed * weight) diff --git a/enginex/ops/sampling.py b/enginex/ops/sampling.py deleted file mode 100644 index 096293fe..00000000 --- a/enginex/ops/sampling.py +++ /dev/null @@ -1,53 +0,0 @@ -""" -EngineX sampling operators. - -rotary_embedding: applies RoPE (Rotary Position Embedding) to Q and K. -Called once per attention layer per forward pass. - -ixformer provides vllm_rotary_embedding_neox natively. -""" - -import torch - - -def rotary_embedding_pytorch( - positions: torch.Tensor, # [num_tokens] - query: torch.Tensor, # [num_tokens, num_heads * head_size] - key: torch.Tensor, # [num_tokens, num_kv_heads * head_size] - head_size: int, - cos_sin_cache: torch.Tensor, # [max_position, rotary_dim] - is_neox: bool = True, -) -> None: - """Apply rotary position embedding in-place on query and key.""" - rotary_dim = cos_sin_cache.shape[1] - half_rot = rotary_dim // 2 - - # Gather cos/sin for each token's position - cos = cos_sin_cache[positions, :half_rot] # [num_tokens, half_rot] - sin = cos_sin_cache[positions, half_rot:] # [num_tokens, half_rot] - - def _apply_rotary(x, cos, sin, head_size, rotary_dim): - """Apply rotary embedding to a reshaped tensor.""" - num_tokens = x.shape[0] - num_heads = x.shape[1] // head_size - x_view = x.view(num_tokens, num_heads, head_size) - - rot = x_view[..., :rotary_dim] - pass_through = x_view[..., rotary_dim:] - - x1 = rot[..., :half_rot] - x2 = rot[..., half_rot:] - - cos_exp = cos.unsqueeze(1) # [num_tokens, 1, half_rot] - sin_exp = sin.unsqueeze(1) - - rot_out = torch.cat([ - x1 * cos_exp - x2 * sin_exp, - x2 * cos_exp + x1 * sin_exp, - ], dim=-1) - - x_view[..., :rotary_dim] = rot_out - x.copy_(x_view.reshape(num_tokens, -1)) - - _apply_rotary(query, cos, sin, head_size, rotary_dim) - _apply_rotary(key, cos, sin, head_size, rotary_dim) diff --git a/enginex/tests/__init__.py b/enginex/tests/__init__.py deleted file mode 100644 index e69de29b..00000000 diff --git a/enginex/tests/test_ops.py b/enginex/tests/test_ops.py deleted file mode 100644 index 3fccbd3d..00000000 --- a/enginex/tests/test_ops.py +++ /dev/null @@ -1,207 +0,0 @@ -""" -EngineX operator tests. - -Tests each operator implementation against known-correct behavior. -Run: python -m pytest enginex/tests/test_ops.py -v -""" - -import sys -import os -sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.dirname(__file__)))) - -import torch -import torch.nn.functional as F -import pytest - - -class TestMoETopkSoftmax: - """Test the critical missing operator.""" - - def test_basic_topk_selection(self): - from enginex.ops.moe import moe_topk_softmax_pytorch - - num_tokens = 4 - num_experts = 8 - topk = 2 - - gating = torch.randn(num_tokens, num_experts) - topk_weights = torch.empty(num_tokens, topk) - topk_ids = torch.empty(num_tokens, topk, dtype=torch.long) - token_expert_indices = torch.empty(num_tokens, topk, dtype=torch.long) - - moe_topk_softmax_pytorch(topk_weights, topk_ids, token_expert_indices, gating) - - # Weights should sum to ~1 per token - sums = topk_weights.sum(dim=-1) - assert torch.allclose(sums, torch.ones(num_tokens), atol=0.01) - - # IDs should be valid expert indices - assert (topk_ids >= 0).all() and (topk_ids < num_experts).all() - - # Should pick the actual top-k from softmax - probs = F.softmax(gating, dim=-1) - for i in range(num_tokens): - expected_ids = torch.topk(probs[i], k=topk).indices - # Same experts selected (order may differ) - assert set(topk_ids[i].tolist()) == set(expected_ids.tolist()) - - def test_large_expert_count(self): - """Qwen3.5 has 128 experts with topk=8.""" - from enginex.ops.moe import moe_topk_softmax_pytorch - - num_tokens = 16 - num_experts = 128 - topk = 8 - - gating = torch.randn(num_tokens, num_experts) - topk_weights = torch.empty(num_tokens, topk) - topk_ids = torch.empty(num_tokens, topk, dtype=torch.long) - token_expert_indices = torch.empty(num_tokens, topk, dtype=torch.long) - - moe_topk_softmax_pytorch(topk_weights, topk_ids, token_expert_indices, gating) - - assert topk_weights.shape == (num_tokens, topk) - assert (topk_ids >= 0).all() and (topk_ids < num_experts).all() - assert not torch.isnan(topk_weights).any() - - -class TestGDN: - """Test GatedDeltaNet implementations.""" - - def test_decode_no_nan(self): - """The critical test — decode must not produce NaN.""" - from enginex.ops.gdn import gdn_decode_pytorch - - B, H, D = 1, 4, 128 - q = torch.randn(B, H, D) - k = torch.randn(B, H, D) - v = torch.randn(B, H, D) - gate = torch.sigmoid(torch.randn(B, H)) - beta = torch.sigmoid(torch.randn(B, H)) * 0.1 - conv_state = torch.randn(B, H, 4, D) - temporal_state = torch.randn(B, H, D, D) * 0.01 - - output, new_state = gdn_decode_pytorch( - q, k, v, gate, beta, conv_state, temporal_state) - - assert not torch.isnan(output).any(), "GDN decode produced NaN!" - assert not torch.isnan(new_state).any(), "GDN state has NaN!" - - def test_prefill_no_nan(self): - """Prefill with chunk_size=16 must not NaN (was 99.98% NaN with 64).""" - from enginex.ops.gdn import gdn_prefill_pytorch - - B, L, H, D = 1, 64, 4, 128 - q = torch.randn(B, L, H, D) * 0.1 - k = torch.randn(B, L, H, D) * 0.1 - v = torch.randn(B, L, H, D) * 0.1 - gate = torch.sigmoid(torch.randn(B, L, H)) - beta = torch.sigmoid(torch.randn(B, L, H)) * 0.1 - state = torch.zeros(B, H, D, D) - - output, final_state = gdn_prefill_pytorch( - q, k, v, gate, beta, state, chunk_size=16) - - nan_frac = torch.isnan(output).float().mean().item() - assert nan_frac < 0.01, f"GDN prefill NaN fraction: {nan_frac:.4f}" - - def test_prefill_state_updates(self): - """State should be different after processing tokens.""" - from enginex.ops.gdn import gdn_prefill_pytorch - - B, L, H, D = 1, 32, 2, 64 - q = torch.randn(B, L, H, D) - k = torch.randn(B, L, H, D) - v = torch.randn(B, L, H, D) - gate = torch.sigmoid(torch.randn(B, L, H)) - beta = torch.sigmoid(torch.randn(B, L, H)) * 0.1 - state = torch.zeros(B, H, D, D) - - _, final_state = gdn_prefill_pytorch( - q, k, v, gate, beta, state, chunk_size=16) - - assert not torch.allclose(final_state, state), "State unchanged after prefill!" - - -class TestActivations: - def test_silu_and_mul(self): - from enginex.ops.activations import silu_and_mul_pytorch - - d = 128 - x = torch.randn(4, d * 2) - out = torch.empty(4, d) - silu_and_mul_pytorch(x, out) - - expected = F.silu(x[..., :d]) * x[..., d:] - assert torch.allclose(out, expected, atol=1e-5) - - def test_gelu_and_mul(self): - from enginex.ops.activations import gelu_and_mul_pytorch - - d = 128 - x = torch.randn(4, d * 2) - out = torch.empty(4, d) - gelu_and_mul_pytorch(x, out) - - expected = F.gelu(x[..., :d]) * x[..., d:] - assert torch.allclose(out, expected, atol=1e-5) - - -class TestNorm: - def test_rms_norm(self): - from enginex.ops.norm import rms_norm_pytorch - - hidden_size = 256 - x = torch.randn(4, hidden_size) - w = torch.ones(hidden_size) - out = torch.empty_like(x) - - rms_norm_pytorch(x, w, out, epsilon=1e-6) - - # Manual check - variance = x.float().pow(2).mean(-1, keepdim=True) - expected = (x * torch.rsqrt(variance + 1e-6)) * w - assert torch.allclose(out, expected, atol=1e-4) - - -class TestRegistry: - def test_registry_creates(self): - from enginex.dispatch.registry import get_registry - reg = get_registry() - assert reg is not None - - def test_probe_runs(self): - from enginex.dispatch.registry import OperatorRegistry - reg = OperatorRegistry() - reg.probe() - # Should have registered operators - assert len(reg.ops) > 0 - - def test_pytorch_fallbacks_always_available(self): - from enginex.dispatch.registry import OperatorRegistry, Backend - reg = OperatorRegistry() - reg.probe() - - # These must ALWAYS have a fallback - critical_ops = [ - "moe_topk_softmax", - "gdn_decode", - "gdn_prefill", - ] - for op_name in critical_ops: - entry = reg.ops.get(op_name) - assert entry is not None, f"{op_name} not registered" - assert entry.active is not None, f"{op_name} has no active impl" - assert entry.active.available, f"{op_name} impl not available" - - def test_summary(self): - from enginex.dispatch.registry import OperatorRegistry - reg = OperatorRegistry() - reg.probe() - summary = reg.summary() - assert "EngineX Operator Registry Summary" in summary - assert "moe_topk_softmax" in summary - - -if __name__ == "__main__": - pytest.main([__file__, "-v"])