Repository navigation
feat(cpu_jit): add CPUID-aware vectorized JIT and cache topology tiling engine - #2467
Open
presentyourlove wants to merge 1 commit into
Open
presentyourlove wants to merge 1 commit into
presentyourlove wants to merge 1 commit into
Conversation
…ng engine - Implement CPUIDInstructionProber with cross-platform SIMD vector flag detection and graceful degradation to prevent Illegal Instruction crashes on legacy hardware - Implement CacheTopologyProber and CacheTilingPlanner for L1/L2 cache capacity matching - Implement VectorizedJITDispatcher supporting AVX512, AVX2, SSE4.1/4.2 unrolled, NEON, and scalar safe kernels with ternary weight unpacking - Implement CacheTiledGEMMEngine with 3-level loop tiling and CoreAffinityManager - Add comprehensive unit test suite in test_cpu_jit.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This PR implements Frontier Dimension 5.1: CPUID-Aware Vectorized JIT & Cache Topology Tiling Engine, resolving compatibility crashes and performance bottlenecks on legacy edge CPUs lacking advanced vector instruction sets.
Background & Motivation
Many edge devices and older PCs (e.g. Intel Core 2 Duo, 1st–4th Gen Core, legacy AMD FX/Phenom, or older ARM SBCs) lack modern vector extensions such as AVX-512, AVX2, or even AVX. Running pre-compiled matrix math binaries directly triggers
Illegal Instruction(SIGILL) crashes. Falling back to non-vectorized scalar operations results in severe performance collapse. Furthermore, small L1/L2 caches on legacy hardware frequently suffer from cache thrashing during matrix multiplication.This PR introduces an architectural CPU JIT subsystem that dynamically detects hardware vector capabilities, gracefully degrades to specialized 128-bit unrolled SSE4 or safe scalar kernels, and calculates optimal 3-level loop tiling dimensions based on hardware L1/L2 cache capacity.
Key Changes
CPU Capability & Cache Geometry Types (
src/exo/shared/types/cpu_jit.py):CPUCapabilityInfo: Captures active instruction set levels (AVX512,AVX2,AVX,SSE4_2,SSE4_1,ARM_NEON,SCALAR), hardware flags, vendor, and core counts.CPUCacheGeometry: Models L1D, L2, L3 cache capacities and cache line sizes with validation.CacheTilingPlan&JITExecutionReport: Captures dynamic M/N/K tile blocks, working set footprints, estimated cache hit rates, and kernel GFLOPS.Cross-Platform Hardware Probers (
src/exo/worker/engines/cpu_jit/cpu_prober.py):CPUIDInstructionProber: Cross-platform probing (Linux/proc/cpuinfo, macOSsysctl, Windows processor identifiers) with support for simulated feature overrides and infallible fallback toSCALAR.CacheTopologyProber: Queries OS sysfs cache topologies with empirical fallback models.Micro-Architecture Vectorized Dispatcher (
src/exo/worker/engines/cpu_jit/vector_jit.py):VectorizedJITDispatcher: Dynamically binds optimal compute kernels:unpack_ternary_weights: Real-time on-CPU unpacking for 1.58-bit ternary {-1, 0, +1} quantized weights.Cache Tiling & Execution Engine (
src/exo/worker/engines/cpu_jit/cache_tiling_engine.py):CacheTilingPlanner: Solves optimal tile dimensions (CacheTiledGEMMEngine: Executes 3-level cache-partitioned matrix multiplication with micro-kernel dispatch and telemetry reporting.CoreAffinityManager: Safe, non-throwing CPU thread and core affinity pinning.Unit Tests (
src/exo/shared/tests/test_cpu_jit.py):Verification
.venv/Scripts/pytest.exe src/exo/shared/tests/test_cpu_jit.py -v # 8 passed in 0.12s.venv/Scripts/ruff.exe check src/exo/shared/types/cpu_jit.py src/exo/worker/engines/cpu_jit src/exo/shared/tests/test_cpu_jit.py # All checks passed!Related Issues
人人都能跑llm.md(Legacy CPU Microarchitecture Vectorized JIT & Cache Tiling)