cpu-device: OR-switch device resolver + RoPE multi-head fallback fix - #160
Open
spaz122772 wants to merge 4 commits into
Open
cpu-device: OR-switch device resolver + RoPE multi-head fallback fix#160spaz122772 wants to merge 4 commits into
spaz122772 wants to merge 4 commits into
Conversation
- NEW engine/device_switch.py: resolve_device()/make_stream()/guard_cuda() OR-switch
- engine.py init now routes through device_switch instead of hard-coded cuda:{rank}
- CI smoke workflow on free ubuntu-latest (reproducible, no paid cloud)
Verified: resolve_device(0)==cpu on M4 (no NVIDIA GPU); switch logic asserted.
Free/open-source CPU-fallback path; porting fork, not upstream.
- setup.py: FREETOKEN_CPU_ONLY=1 builds C++ exts with stub cuda runtime (no cudart/nvcc) - NEW stub_cuda_runtime.h: faithful no-op CUDA symbols + CUDART_CB + enum flags - pinned_tensor.cpp / cpu_moe_ext.cpp: #ifdef FREETOKEN_CPU_ONLY -> stub header - engine.py: OR-switch ALL torch.cuda.* calls behind has_cuda() (init, sync, events, allocator, graph teardown, forward assert, ForwardOutput annotation) - offload_cache.py / offload_kernels.py: flashlib/triton imports guarded so the package imports on CPU-only hardware (GPU path untouched when flashlib present) - CI: free ubuntu-latest verifies OR-switch + CPU-only build + CPU smoke-serve VERIFIED LOCALLY on M4 (arm64, no CUDA): - package imports without triton/flashlib - FREETOKEN_CPU_ONLY=1 build -> _pinned_tensor + _cpu_moe compile, link, import - AVX-512 kernels use runtime dispatch w/ scalar fallback -> portable build - OR-switch resolve_device(0) -> cpu Remaining (needs x86-Linux or free CI runner): real model generation smoke test.
Root cause: apply_rope_with_cos_sin_cache_inplace mishandled the engine's multi-head concatenated query (seq, NQ*head_dim)=(seq,2048). It treated the whole strip as one head (2048-wide) and rotated only the first rotary_dim dims, leaving 15 of 16 heads un-rotated -> attention divergence (engine vs HF post-RoPE cosine -0.02). Fix: reshape (seq, NQ*head_dim) -> (seq*NQ, head_dim), rotate each head with 2-D cos/sin (no phantom unsqueeze dim), reshape back. Isolated unit test tools/rope_fallback_test.py: cos=1.000000 (single- and multi-head). Verified on M4 (arm64, NO CUDA, NO triton): - qwen3-0.6b serves via --moe-backend cpu, answers 'The capital of France is' -> 'Paris.' - sustained 120-tok decode coherent; temp>0 sampling diverges across runs. - cold-start reproducibility confirmed from clean process/port state. Also: OR-switch device resolver (CUDA-preferred, CPU-fallback), CPU-only C++ build via no-op CUDA stub, flag-guarded probe hooks (zero cost when off). Consolidated record: docs/freetoken-deployment/phase5-cpu-device/CONSOLIDATED.md. Local model weights (_test_models/) gitignored — not committed.
…tive)
Root cause: setup.py passed CppExtension sources as absolute paths
(str(ROOT / '...')) and pyproject listed csrc/**/* as package-data.
Modern setuptools (>=77, what the free CI runner pulls) runs
build_py.assert_relative() and rejects absolute paths
(DistutilsSetupError: setup script specifies an absolute path),
breaking the CPU-only editable build on GitHub Actions.
Fix:
- setup.py: pass extension sources as paths relative to setup.py
('python/freetoken/kernel/csrc/...'), the portable form setuptools
accepts. Applies to both the CPU-only and CUDA build branches.
- pyproject.toml: drop csrc/**/* from package-data; C++ sources are
build inputs, not runtime data, and must not be shipped as such.
Verified locally with setuptools 84.0.0 (editable CPU-only install):
FREETOKEN_CPU_ONLY=1 pip install -e . --no-build-isolation -> OK
from freetoken.kernel import _cpu_moe, _pinned_tensor -> OK
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
1 participant
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.
Summary
engine/device_switch.py): prefer CUDA when available, fall back to CPU. Same branch runs on GPU or CPU-only boxes unchanged.FREETOKEN_CPU_ONLY=1build.kernel/torch_fallback.py: the engine passes a multi-head concatenated query(seq, NQ*head_dim)=(seq, 2048)(16 heads x 128); the old code rotated only the first head's dims, leaving 15/16 heads un-rotated (engine-vs-reference post-RoPE cosine was -0.02). Fix reshapes per-head, keeps cos/sin 2-D (no phantom broadcast dim), rotates each head, reassembles. Unit test cos = 1.000000 (single- and multi-head).Verified (real execution, arm64 CPU-only, no CUDA/triton)
qwen3-0.6bserves via--moe-backend cpu; "The capital of France is" -> "Paris."tools/rope_fallback_test.py.Notes / limitations
qwen3-0.6bexercised locally; fix lives in the shared fallback kernel (arch-general for the CPU/MoE path).cpu-device-smoke.yml) added for free-runner verification.Branch: https://github.com/spaz122772/FreeToken/tree/cpu-device